[KM 436-439] adjust endpoint for new features

src/agents/orchestration.py

@@ -35,6 +35,11 @@ Intent Routing:
 - greeting -> needs_search=False, direct_response="Hello! How can I assist you today?"
 - goodbye -> needs_search=False, direct_response="Goodbye! Have a great day!"
 - other -> needs_search=True, search_query=<standalone rewritten query>
+
+Source Routing (set source_hint):
+- Columns, tables, sheets, data types, schema, row counts, statistics -> source_hint=schema
+- Document content, paragraphs, reports, articles, text -> source_hint=document
+- Unclear or spans both -> source_hint=both
 """),
             MessagesPlaceholder(variable_name="history"),
             ("user", "{message}")

src/api/v1/chat.py

@@ -9,6 +9,9 @@ from src.db.postgres.models import ChatMessage, MessageSource
 from src.agents.orchestration import orchestrator
 from src.agents.chatbot import chatbot
 from src.rag.retriever import retriever
+from src.rag.base import RetrievalResult
+from src.query.query_executor import query_executor
+from src.query.base import QueryResult
 from src.db.redis.connection import get_redis
 from src.config.settings import settings
 from src.middlewares.logging import get_logger, log_execution
@@ -61,7 +64,7 @@ def _extract_sources(results: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
     seen = set()
     sources = []
     for result in results:
-        if "document_id" in result["metadata"]["data"]:
+        if "document_id" in result["metadata"].get("data", {}):
             meta = result["metadata"]
             key = (meta.get("data", {}).get("document_id"), meta.get("data", {}).get("page_label"))
             if key not in seen:
@@ -88,6 +91,22 @@ def _extract_sources(results: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
     return sources
 
 
+def _format_query_results(results: list[QueryResult]) -> str:
+    if not results:
+        return ""
+    lines = []
+    for r in results:
+        name = r.metadata.get("client_name", r.source_id)
+        lines.append(f"[Query result — {name}, tables: {r.table_or_file}]")
+        lines.append(f"SQL: {r.metadata.get('sql', '')}")
+        if r.columns and r.rows:
+            lines.append(" | ".join(r.columns))
+            for row in r.rows[:20]:
+                lines.append(" | ".join(str(row.get(c, "")) for c in r.columns))
+        lines.append(f"({r.row_count} rows total)\n")
+    return "\n".join(lines)
+
+
 async def get_cached_response(redis, cache_key: str) -> Optional[str]:
     cached = await redis.get(cache_key)
     if cached:
@@ -182,16 +201,25 @@ async def chat_stream(request: ChatRequest, db: AsyncSession = Depends(get_db)):
 
             if not intent_result.get("needs_search"):
                 retrieval_task.cancel()
+                try:
+                    await retrieval_task
+                except asyncio.CancelledError:
+                    pass
                 raw_results = []
             else:
                 search_query = intent_result.get("search_query", request.message)
                 logger.info(f"Searching for: {search_query}")
                 if search_query != request.message:
                     retrieval_task.cancel()
+                    try:
+                        await retrieval_task
+                    except asyncio.CancelledError:
+                        pass
                     raw_results = await retriever.retrieve(
                         query=search_query,
                         user_id=request.user_id,
                         db=db,
+                        source_hint=intent_result.get("source_hint", "both"),
                     )
                 else:
                     raw_results = await retrieval_task
@@ -199,6 +227,27 @@ async def chat_stream(request: ChatRequest, db: AsyncSession = Depends(get_db)):
             context = _format_context(raw_results)
             sources = _extract_sources(raw_results)
 
+            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,
+                    user_id=request.user_id,
+                    db=db,
+                    question=intent_result.get("search_query") or request.message,
+                )
+                query_context = _format_query_results(query_results)
+                if query_context:
+                    context = query_context + "\n\n" + context
+
         # Step 3: Direct response for greetings / non-document intents
         if intent_result.get("direct_response"):
             response = intent_result["direct_response"]

src/api/v1/db_client.py

@@ -458,7 +458,7 @@ async def ingest_database_client(
             db_type=client.db_type,
             credentials=creds,
         ) as engine:
-            total = await db_pipeline_service.run(user_id=user_id, engine=engine)
+            total = await db_pipeline_service.run(user_id=user_id, client_id=client_id, engine=engine)
     except NotImplementedError as e:
         raise HTTPException(status_code=status.HTTP_501_NOT_IMPLEMENTED, detail=str(e))
     except Exception as e: