Improve chat agent prompts, citation linking, and tool-call handling

src/app/api/api_v1/endpoints/chat.py

@@ -20,6 +20,7 @@
 from src.app.models import chat as models
 from src.app.search.services.search import SearchService, get_search_service
 from src.app.services.data_collection import get_data_collection_service
+from src.app.services.helpers import linkify_missing_citations
 from src.app.shared.domain.constants import subjects as subjectsDict
 from src.app.shared.domain.exceptions import (
     EmptyQueryError,
@@ -71,6 +72,7 @@ def get_params(body: models.Context) -> models.ContextOut:
         query=body.query,
         subject=body.subject,
         conversation_id=None,
+        lang=body.lang,
     )
 
 
@@ -154,7 +156,7 @@ async def q_and_a_new_questions(
 ):
     try:
         new_questions = await chatfactory.get_new_questions(
-            query=body.query, history=body.history
+            query=body.query, history=body.history, lang=body.lang
         )
 
         return new_questions
@@ -487,13 +489,17 @@ async def agent_response(
                 sp=sp,
             )
 
-        if isinstance(res["messages"][-2], ToolMessage):
-            docs = res["messages"][-2].artifact
-        else:
-            docs = None
+        all_docs = []
+        for msg in res["messages"]:
+            if isinstance(msg, ToolMessage) and getattr(msg, "artifact", None):
+                all_docs.extend(msg.artifact)
+        docs = all_docs if all_docs else None
+        content = linkify_missing_citations(
+            cast(str, res["messages"][-1].content), docs or []
+        )
 
         agent_ans = {
-            "content": cast(str, res["messages"][-1].content),
+            "content": content,
             "docs": docs,
             "thread_id": thread_id,
         }
@@ -503,7 +509,7 @@ async def agent_response(
                 session_id=session_id,
                 user_query=body.query,
                 conversation_id=thread_id,
-                answer_content=res["messages"][-1].content,
+                answer_content=content,
                 sources=docs,
             )
 

src/app/api/api_v1/endpoints/chat_utils.py

@@ -11,6 +11,7 @@
 
 from src.app.models import chat as models
 from src.app.search.services.search import SearchService
+from src.app.services.helpers import linkify_missing_citations
 from src.app.utils.logger import logger as utils_logger
 
 logger = utils_logger(__name__)
@@ -154,7 +155,7 @@ async def _stream_agent_response(
     thread_id: UUID,
 ) -> AsyncGenerator[str, None]:
     final_content = ""
-    docs = []
+    docs = None
     has_streamed_content = False
 
     stream = _stream_agent_with_memory(
@@ -178,6 +179,8 @@ async def _stream_agent_response(
         except Exception as e:
             logger.error("Error while yielding chunk: %s", e)
 
+    final_content = linkify_missing_citations(final_content, docs)
+
     final_payload = _build_final_stream_payload(
         final_content=final_content,
         docs=docs,

src/app/models/chat.py

@@ -15,6 +15,7 @@ class Context(BaseModel):
     history: list[dict] | None = []
     query: str | None = None
     subject: str | None = Field(None)
+    lang: str | None = None
 
 
 class ContextOut(BaseModel):
@@ -23,6 +24,7 @@ class ContextOut(BaseModel):
     query: str
     subject: str | None = Field(None)
     conversation_id: uuid.UUID | None = Field(None)
+    lang: str | None = None
 
 
 class Role(Enum):

src/app/search/api/router.py

@@ -148,7 +148,6 @@ async def search_all_slices_by_lang(
     sp: SearchService = Depends(get_search_service),
 ):
     try:
-
         res = await sp.search_handler(
             qp=qp, method=SearchMethods.BY_SLICES, background_tasks=background_tasks
         )

src/app/services/agent.py

@@ -20,6 +20,14 @@ async def _get_resources_about_sustainability(
 ) -> Tuple[str, List[Document]]:
     """Core logic for getting relevant resources about sustainability from WeLearn database."""
 
+    tool_called: list = config["configurable"].get("tool_called", [False])
+    if tool_called[0]:
+        return (
+            "Search has already been performed. Use the documents already retrieved to answer the question.",
+            [],
+        )
+    tool_called[0] = True
+
     qp = EnhancedSearchQuery(
         query=rag_query,
         sdg_filter=config["configurable"].get("sdg_filter"),

src/app/services/helpers.py

@@ -1,5 +1,6 @@
+import re
 from functools import cache
-from typing import Any, List, cast
+from typing import Any, List, Optional, cast
 
 import json_repair
 import numpy
@@ -139,6 +140,43 @@ def stringify_docs_content(docs: List[Any]) -> str:
     return documents.strip()
 
 
+_CITATION_RE = re.compile(r'(?<!target="_blank">)\[Doc\s*(\d+)\]')
+
+
+def linkify_missing_citations(text: str, docs: List[Any]) -> str:
+    """
+    Wraps any bare `[Doc N]` marker in `text` with the `<a href=... target="_blank">`
+    tag for document N, using the URL from `docs[N-1]`. Markers already wrapped in an
+    <a> tag are left untouched. Safety net for when the LLM forgets to format a
+    citation as a link itself.
+
+    Args:
+        text: The assembled answer text.
+        docs: The full list of retrieved documents, 1-indexed by citation number.
+
+    Returns:
+        str: `text` with any missed citations auto-linked. Unresolvable markers
+        (out-of-range N, or no docs) are left as-is.
+    """
+    if not text or not docs:
+        return text
+
+    def _url_for(n: int) -> Optional[str]:
+        if not (1 <= n <= len(docs)):
+            return None
+        payload = normalize_payload(getattr(docs[n - 1], "payload", docs[n - 1]))
+        return str(payload.get("document_url", "")).strip() or None
+
+    def _replace(match: "re.Match[str]") -> str:
+        n = int(match.group(1))
+        url = _url_for(n)
+        if not url:
+            return match.group(0)
+        return f'<a href="{url}" target="_blank">[Doc {n}]</a>'
+
+    return _CITATION_RE.sub(_replace, text)
+
+
 def extract_json_from_response(
     response: str,
 ) -> JSONReturnType | tuple[JSONReturnType, list[dict[str, str]]] | str: