feat: expand source metadata

app/agent.py

@@ -1,10 +1,11 @@
-import json
 import re
 from pathlib import Path
 
 from google import genai
 from google.genai import types
+from pydantic import ValidationError
 
+from app import models
 from app.tools import dailymed, literature
 
 CONFIG = types.GenerateContentConfig(
@@ -16,32 +17,59 @@ CONFIG = types.GenerateContentConfig(
     system_instruction=(Path(__file__).parent / "system_instruction.txt").read_text(),
 )
 
+SOURCE_TOOL_NAMES = {
+    literature.search_medical_literature.__name__,
+    dailymed.find_drug_set_ids.__name__,
+}
 
-def respond(client: genai.Client, query: str) -> list[dict]:
-    config = types.GenerateContentConfig(
-        tools=[
-            dailymed.find_drug_set_ids,
-            dailymed.find_drug_instruction,
-            literature.search_medical_literature,
-        ],
-        system_instruction=SYSTEM_INSTRUCTION,
+
+def hydrate_sources(
+    statements: models.Statements, calling_history: list[types.Content]
+) -> models.Statements:
+    sources = {}
+    for call in calling_history:
+        for part in call.parts or []:
+            if (
+                (func := part.function_response)
+                and func.name in SOURCE_TOOL_NAMES
+                and func.response
+            ):
+                for source in func.response["result"]:
+                    sources[source["url"]] = source
+
+    for statement in statements.statements:
+        if statement.sources:
+            statement.sources = [
+                models.Source.model_validate(sources[source.url])
+                for source in statement.sources
+            ]
+
+    return statements
+
+
+def validate_response(response: types.GenerateContentResponse) -> models.Statements:
+    text = (response.text or "").strip()
+
+    # Extract content inside the first markdown code block (``` or ```json)
+    match = re.match(r"^```(?:json)?\s*([\s\S]*?)\s*```", text)
+    if match:
+        text = match.group(1).strip()
+
+    try:
+        statements = models.Statements.model_validate_json(f'{{"statements":{text}}}')
+    except ValidationError:
+        statements = models.Statements(statements=[models.Statement(text=text)])
+
+    statements = hydrate_sources(
+        statements, response.automatic_function_calling_history or []
     )
+    return statements
+
+
+def respond(client: genai.Client, query: str) -> models.Statements:
     resp = client.models.generate_content(
         model="gemini-2.5-flash-preview-04-17",
         contents=query,
         config=CONFIG,
     )
-
-    output = ((resp.text) or "").strip()
-
-    if output.startswith("```"):
-        # Extract content inside the first markdown code block (``` or ```json)
-        match = re.match(r"^```(?:json)?\s*([\s\S]*?)\s*```", output)
-        if match:
-            output = match.group(1).strip()
-
-    try:
-        return json.loads(output)
-    except json.decoder.JSONDecodeError as err:
-        print(err)
-        return [{"text": output}]
+    return validate_response(resp)

app/main.py

@@ -14,5 +14,4 @@ def health_check():
 
 @app.get("/ask", response_model=models.Statements)
 def ask(query: str):
-    output = agent.respond(gemini, query)
-    return {"statements": output}
+    return agent.respond(gemini, query)

app/tools/dailymed.py

@@ -18,8 +18,10 @@ def find_drug_set_ids(name: str) -> list[dict]:
     )
     return [
         {
-            "name": row["title"],
+            "title": row["title"],
             "set_id": row["setid"],
+            "venue": "DailyMed",
+            "year": row["published_date"][-4:],  # Original format: "May 05, 2025"
             "url": f"https://dailymed.nlm.nih.gov/dailymed/drugInfo.cfm?setid={row['setid']}",
         }
         for row in resp.json()["data"]

app/tools/literature.py

@@ -49,26 +49,20 @@ def get_pubmed_abstracts(pmids: list[int]) -> dict[str, dict]:
     return abstracts
 
 
-def format_publication(publication: dict) -> str:
-    title = publication["title"]
-    summary = (publication["tldr"] or {}).get("text", "")
-    abstract = publication["abstract"]
-    venue = publication["venue"]
-    year = publication["year"]
-    citations = publication["citationCount"]
-    influential_citations = publication["influentialCitationCount"]
-    doi = publication["externalIds"].get("DOI")
-    url = f"https://doi.org/{doi}" if doi else publication["url"]
-    return (
-        f"<publication title={title}>\n<url>{url}</url>\n"
-        f"<summary>{summary}</summary>\n<abstract>{abstract}</abstract>\n"
-        f"<venue>{venue}</venue>\n<year>{year}</year>\n"
-        f"<citationCount>{citations}</citationCount>\n<influentialCitationCount>{influential_citations}</influentialCitationCount>\n"
-        "</publication>"
-    )
-
-
-def search_medical_literature(query: str) -> str:
+def format_publication(publication: dict) -> dict:
+    tldr = publication.pop("tldr") or {}
+    external_ids = publication.pop("externalIds")
+    doi = external_ids.get("DOI")
+    publication["summary"] = tldr.get("text", "")
+    publication["citations"] = publication.pop("citationCount")
+    publication["influential_citations"] = publication.pop("influentialCitationCount")
+    publication["doi"] = doi
+    if doi:
+        publication["url"] = f"https://doi.org/{doi}"
+    return publication
+
+
+def search_medical_literature(query: str) -> list[dict]:
     """Get medical literature related to the query.
 
     Args:
@@ -98,8 +92,4 @@ def search_medical_literature(query: str) -> str:
 
         outputs.append(format_publication(publication))
 
-    return (
-        f"<publications>\n{'\n'.join(outputs)}\n</publications>"
-        if outputs
-        else "No literature found"
-    )
+    return outputs

main.py

@@ -1,32 +1,27 @@
 import streamlit as st
 from google import genai
 
-from app import agent, config
-
-
-def format_output(response: list[dict]) -> tuple[str, str]:
-    try:
-        answer = ""
-        citations = {}
-        for statement in response:
-            text = statement["text"].strip()
-            answer = (
-                answer + f"\n{text}"
-                if text.startswith("*") or text.startswith("-")
-                else answer + f" {text}"
-            )
+from app import agent, config, models
+
+
+def format_output(statements: models.Statements) -> tuple[str, str]:
+    answer = ""
+    citations = {}
+
+    for statement in statements.statements:
+        if statement.text.startswith(("*", "-")):
+            # Bullet points should be on a newline.
+            answer += "\n"
+        answer += statement.text
+
+        if statement.sources:
             citation_ids = []
-            for source in statement.get("sources", []):
-                source_str = f"[{source['title']}]({source['url']})"
-                if not (citation_id := citations.get(source_str)):
+            for source in statement.sources:
+                if not (citation_id := citations.get(source.citation)):
                     citation_id = len(citations) + 1
-                    citations[source_str] = citation_id
+                    citations[source.citation] = citation_id
                 citation_ids.append(citation_id)
-            if citation_ids:
-                answer += " ".join(f"[^{i}]" for i in sorted(citation_ids))
-    except KeyError as err:
-        print(err)
-        return str(response), ""
+            answer += " ".join(f"[^{i}]" for i in sorted(citation_ids))
 
     footnotes = "\n".join(f"[^{id}]: {citation}" for citation, id in citations.items())
     return answer, footnotes