| 55 | |
| 56 | @app.post("/search", response_model=Union[Tuple[List[Document], List[float]], List[Document]]) |
| 57 | async def search(request: QueryRequest): |
| 58 | query = request.query |
| 59 | top_n = request.top_n |
| 60 | return_score = request.return_score |
| 61 | |
| 62 | if not query or not query.strip(): |
| 63 | print(f"Query content cannot be empty: {query}") |
| 64 | raise HTTPException( |
| 65 | status_code=400, |
| 66 | detail="Query content cannot be empty" |
| 67 | ) |
| 68 | |
| 69 | async with retriever_semaphore: |
| 70 | async with retriever_lock: |
| 71 | retriever_idx = available_retrievers.popleft() |
| 72 | try: |
| 73 | loop = asyncio.get_event_loop() |
| 74 | with ThreadPoolExecutor() as executor: |
| 75 | if return_score: |
| 76 | results, scores = await loop.run_in_executor( |
| 77 | executor, |
| 78 | retriever_list[retriever_idx].search, |
| 79 | query, top_n, return_score |
| 80 | ) |
| 81 | return [Document(id=result['id'], contents=result['contents']) for result in results], scores |
| 82 | else: |
| 83 | results = await loop.run_in_executor( |
| 84 | executor, |
| 85 | retriever_list[retriever_idx].search, |
| 86 | query, top_n, return_score |
| 87 | ) |
| 88 | return [Document(id=result['id'], contents=result['contents']) for result in results] |
| 89 | finally: |
| 90 | async with retriever_lock: |
| 91 | available_retrievers.append(retriever_idx) |
| 92 | |
| 93 | @app.post("/batch_search", response_model=Union[List[List[Document]], Tuple[List[List[Document]], List[List[float]]]]) |
| 94 | async def batch_search(request: BatchQueryRequest): |