Spaces:
Sleeping
Sleeping
fix(search): add missing return statement in search service
Browse files- api/schemas/search.py +2 -1
- api/v1/search.py +4 -2
- services/search_service.py +13 -6
api/schemas/search.py
CHANGED
|
@@ -13,7 +13,8 @@ class SearchRequest(BaseModel):
|
|
| 13 |
top_k: int = Field(default=5, ge=1, le=50)
|
| 14 |
|
| 15 |
# optional
|
| 16 |
-
|
|
|
|
| 17 |
|
| 18 |
# ---------------------------
|
| 19 |
# Document metadata
|
|
|
|
| 13 |
top_k: int = Field(default=5, ge=1, le=50)
|
| 14 |
|
| 15 |
# optional
|
| 16 |
+
use_reranking: Optional[bool] = True
|
| 17 |
+
include_llm_context: Optional[bool] = True
|
| 18 |
|
| 19 |
# ---------------------------
|
| 20 |
# Document metadata
|
api/v1/search.py
CHANGED
|
@@ -28,7 +28,9 @@ async def execute_search(
|
|
| 28 |
try:
|
| 29 |
search_output = search_service.search(
|
| 30 |
query=request_data.query,
|
| 31 |
-
top_k=request_data.top_k
|
|
|
|
|
|
|
| 32 |
)
|
| 33 |
return SearchResponse(
|
| 34 |
query=search_output["query"],
|
|
@@ -66,7 +68,7 @@ async def demo_page_post(
|
|
| 66 |
Handle form submission from the demo page, execute the search, and render results in the same template.
|
| 67 |
"""
|
| 68 |
try:
|
| 69 |
-
search_output = search_service.search(query=query, top_k=5)
|
| 70 |
|
| 71 |
return templates.TemplateResponse(
|
| 72 |
"index.html",
|
|
|
|
| 28 |
try:
|
| 29 |
search_output = search_service.search(
|
| 30 |
query=request_data.query,
|
| 31 |
+
top_k=request_data.top_k,
|
| 32 |
+
use_reranking=request_data.use_reranking,
|
| 33 |
+
include_llm_context=request_data.include_llm_context
|
| 34 |
)
|
| 35 |
return SearchResponse(
|
| 36 |
query=search_output["query"],
|
|
|
|
| 68 |
Handle form submission from the demo page, execute the search, and render results in the same template.
|
| 69 |
"""
|
| 70 |
try:
|
| 71 |
+
search_output = search_service.search(query=query, top_k=5, use_reranking=False, include_llm_context=False)
|
| 72 |
|
| 73 |
return templates.TemplateResponse(
|
| 74 |
"index.html",
|
services/search_service.py
CHANGED
|
@@ -24,13 +24,14 @@ class HybridSearchService:
|
|
| 24 |
self.embedder = embedder
|
| 25 |
self.reranker = reranker
|
| 26 |
|
| 27 |
-
def search(self, query: str, top_k: int = 5, limit: int = 50, include_llm_context: bool = True) -> Dict[str, Any]:
|
| 28 |
"""
|
| 29 |
Receives user queries and performs hybrid search and reranking.
|
| 30 |
|
| 31 |
:param query: User search query
|
| 32 |
:param top_k: Number of documents to return (after reranking)
|
| 33 |
:param limit: Number of candidate documents to fetch from Qdrant (after RRF fusion, before reranking)
|
|
|
|
| 34 |
:param include_llm_context: Whether to include LLM context in the response (formatted text for LLM consumption)
|
| 35 |
:return: A dictionary containing the original query, a list of search results, and latency information. Each search result includes chunk_id, text, relevance score, and metadata.
|
| 36 |
"""
|
|
@@ -80,11 +81,15 @@ class HybridSearchService:
|
|
| 80 |
|
| 81 |
# 5. Perform Cross-Encoder Reranking
|
| 82 |
# Return a list sorted in descending order after recalculating context-based precise scores
|
| 83 |
-
|
| 84 |
-
|
| 85 |
-
|
| 86 |
-
|
| 87 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 88 |
|
| 89 |
# 6. Top-K Truncation and Mapping to Pydantic Schema (SearchResultItem) Specification
|
| 90 |
final_results = []
|
|
@@ -108,6 +113,8 @@ class HybridSearchService:
|
|
| 108 |
if include_llm_context:
|
| 109 |
# 7. Optional: Format results into LLM-friendly context (Markdown/XML mixed format)
|
| 110 |
response["llm_context"] = self.format_for_llm(final_results)
|
|
|
|
|
|
|
| 111 |
|
| 112 |
except Exception as e:
|
| 113 |
# Wrap unexpected errors in custom errors and throw them to the router
|
|
|
|
| 24 |
self.embedder = embedder
|
| 25 |
self.reranker = reranker
|
| 26 |
|
| 27 |
+
def search(self, query: str, top_k: int = 5, limit: int = 50, use_reranking: bool = True, include_llm_context: bool = True) -> Dict[str, Any]:
|
| 28 |
"""
|
| 29 |
Receives user queries and performs hybrid search and reranking.
|
| 30 |
|
| 31 |
:param query: User search query
|
| 32 |
:param top_k: Number of documents to return (after reranking)
|
| 33 |
:param limit: Number of candidate documents to fetch from Qdrant (after RRF fusion, before reranking)
|
| 34 |
+
:param use_reranking: Whether to use reranking (if False, it will skip the reranking step and return Qdrant results directly, still mapped with SQLite data)
|
| 35 |
:param include_llm_context: Whether to include LLM context in the response (formatted text for LLM consumption)
|
| 36 |
:return: A dictionary containing the original query, a list of search results, and latency information. Each search result includes chunk_id, text, relevance score, and metadata.
|
| 37 |
"""
|
|
|
|
| 81 |
|
| 82 |
# 5. Perform Cross-Encoder Reranking
|
| 83 |
# Return a list sorted in descending order after recalculating context-based precise scores
|
| 84 |
+
if use_reranking:
|
| 85 |
+
reranked_docs = self.reranker.rerank(
|
| 86 |
+
query=query,
|
| 87 |
+
documents=chunks_for_reranking,
|
| 88 |
+
text_key="text"
|
| 89 |
+
)
|
| 90 |
+
else:
|
| 91 |
+
# If reranking is disabled, use the Qdrant results directly
|
| 92 |
+
reranked_docs = chunks_for_reranking
|
| 93 |
|
| 94 |
# 6. Top-K Truncation and Mapping to Pydantic Schema (SearchResultItem) Specification
|
| 95 |
final_results = []
|
|
|
|
| 113 |
if include_llm_context:
|
| 114 |
# 7. Optional: Format results into LLM-friendly context (Markdown/XML mixed format)
|
| 115 |
response["llm_context"] = self.format_for_llm(final_results)
|
| 116 |
+
|
| 117 |
+
return response
|
| 118 |
|
| 119 |
except Exception as e:
|
| 120 |
# Wrap unexpected errors in custom errors and throw them to the router
|