[NOTICKET] minor refactoring

src/api/v1/chat.py

@@ -48,43 +48,43 @@ class ChatRequest(BaseModel):
     message: str
 
 
-def _format_context(results: List[Dict[str, Any]]) -> str:
+def _format_context(results: List[RetrievalResult]) -> str:
     """Format retrieval results as context string for the LLM."""
     lines = []
     for result in results:
-        filename = result["metadata"].get("filename", "Unknown")
-        page = result["metadata"].get("page_label")
+        filename = result.metadata.get("filename", "Unknown")
+        page = result.metadata.get("page_label")
         source_label = f"{filename}, p.{page}" if page else filename
-        lines.append(f"[Source: {source_label}]\n{result['content']}\n")
+        lines.append(f"[Source: {source_label}]\n{result.content}\n")
     return "\n".join(lines)
 
 
-def _extract_sources(results: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+def _extract_sources(results: List[RetrievalResult]) -> List[Dict[str, Any]]:
     """Extract deduplicated source references from retrieval results."""
     seen = set()
     sources = []
     for result in results:
-        if "document_id" in result["metadata"].get("data", {}):
-            meta = result["metadata"]
-            key = (meta.get("data", {}).get("document_id"), meta.get("data", {}).get("page_label"))
+        meta = result.metadata
+        data = meta.get("data", {})
+        if "document_id" in data:
+            key = (data.get("document_id"), data.get("page_label"))
             if key not in seen:
                 seen.add(key)
                 sources.append({
-                    "document_id": meta.get("data", {}).get("document_id"),
-                    "filename": meta.get("data", {}).get("filename", "Unknown"),
-                    "page_label": meta.get("data", {}).get("page_label", "Unknown"),
+                    "document_id": data.get("document_id"),
+                    "filename": data.get("filename", "Unknown"),
+                    "page_label": data.get("page_label", "Unknown"),
                 })
         else:
-            meta = result["metadata"]
-            key = (meta.get("data", {}).get("table_name"), meta.get("data", {}).get("column_name"))
+            key = (data.get("table_name"), data.get("column_name"))
             if key not in seen:
                 seen.add(key)
-                table_name = meta.get("data", {}).get("table_name")
+                table_name = data.get("table_name")
                 user_id = meta.get("user_id")
                 sources.append({
                     "document_id": f"{user_id}_{table_name}",
-                    "filename": meta.get("data", {}).get("table_name", "Unknown"),
-                    "page_label": meta.get("data", {}).get("column_name", "Unknown"),
+                    "filename": data.get("table_name", "Unknown"),
+                    "page_label": data.get("column_name", "Unknown"),
                 })
 
     logger.debug(f"Extracted sources: {sources}")
@@ -229,17 +229,8 @@ async def chat_stream(request: ChatRequest, db: AsyncSession = Depends(get_db)):
 
             source_hint = intent_result.get("source_hint", "both")
             if source_hint in ("schema", "both"):
-                retrieval_objects = [
-                    RetrievalResult(
-                        content=r["content"],
-                        metadata=r["metadata"],
-                        score=0.0,
-                        source_type=r["metadata"].get("source_type", ""),
-                    )
-                    for r in raw_results
-                ]
                 query_results = await query_executor.execute(
-                    results=retrieval_objects,
+                    results=raw_results,
                     user_id=request.user_id,
                     db=db,
                     question=intent_result.get("search_query") or request.message,

src/rag/retriever.py

@@ -1,10 +1,9 @@
 """Public retrieval API — thin wrapper around RetrievalRouter."""
 
-from typing import Any
-
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from src.middlewares.logging import get_logger
+from src.rag.base import RetrievalResult
 from src.rag.retrievers.document import document_retriever
 from src.rag.retrievers.schema import schema_retriever
 from src.rag.router import RetrievalRouter, SourceHint
@@ -16,7 +15,11 @@ class RetrieverService:
     """Public retrieval service used by chat.py and search tools.
 
     Delegates to RetrievalRouter which dispatches based on source_hint.
-    Returns List[Dict] to preserve backward compatibility with chat.py.
+    Returns RetrievalResult objects directly so downstream consumers
+    (db_executor, tabular_executor) can be fed without lossy dict
+    conversion. The `db` parameter is accepted for call-site compatibility
+    but currently unused — retrieval reads PGVector via _pgvector_engine
+    inside each retriever.
     """
 
     def __init__(self):
@@ -32,10 +35,9 @@ class RetrieverService:
         db: AsyncSession,
         k: int = 5,
         source_hint: SourceHint = "both",
-    ) -> list[dict[str, Any]]:
+    ) -> list[RetrievalResult]:
         try:
-            results = await self._router.retrieve(query, user_id, source_hint, k)
-            return [{"content": r.content, "metadata": r.metadata} for r in results]
+            return await self._router.retrieve(query, user_id, source_hint, k)
         except Exception as e:
             logger.error("retrieval failed", error=str(e))
             return []

src/rag/retrievers/baseline.py

@@ -13,8 +13,14 @@ logger = get_logger("retriever")
 _RETRIEVAL_CACHE_TTL = 3600  # 1 hour
 
 
-class RetrieverService:
-    """Service for retrieving relevant documents."""
+class BaselineRetrieverService:
+    """Baseline (pre-Phase-1) retriever — preserved for benchmark comparison.
+
+    Renamed from RetrieverService so it doesn't shadow the production wrapper
+    at src/rag/retriever.py. Production code imports from src.rag.retriever;
+    benchmark scripts that want this baseline must import explicitly from
+    src.rag.retrievers.baseline.
+    """
 
     def __init__(self):
         self.vector_store = get_vector_store()
@@ -67,4 +73,4 @@ class RetrieverService:
             return []
 
 
-retriever = RetrieverService()
+baseline_retriever = BaselineRetrieverService()

src/tools/search.py

@@ -34,10 +34,10 @@ async def search_documents(
 
         formatted_results = []
         for result in results:
-            filename = result["metadata"].get("filename", "Unknown")
-            page = result["metadata"].get("page_label")
+            filename = result.metadata.get("filename", "Unknown")
+            page = result.metadata.get("page_label")
             source_label = f"{filename}, p.{page}" if page else filename
-            formatted_results.append(f"[Source: {source_label}]\n{result['content']}\n")
+            formatted_results.append(f"[Source: {source_label}]\n{result.content}\n")
 
         return "\n".join(formatted_results)