feat: improve citation instruction

main.py

@@ -12,15 +12,14 @@ class Article(BaseModel):
     title: str
     summary: str | None
     abstract: str | None
-
-
-class Citation(BaseModel):
-    source_id: str
+    venue: str
+    year: int
+    citations: int | None
 
 
 class Statement(BaseModel):
     text: str
-    citation: Citation | None
+    source_ids: list[str] | None
 
 
 def improve_prompt(client: genai.Client, prompt: str) -> str:
@@ -44,12 +43,33 @@ def format_sources(articles: list[Article]) -> str:
     return "\n".join(sources)
 
 
+ANSWER_INSTRUCTION = """\
+You are a medical research expert.
+
+Please answer the user's query clearly and concisely, using no more than 100 words.
+
+Base every claim or statement strictly on the provided sources. For each claim, include a citation referencing the source's ID (do not include the citation in the `text` field). A claim may be supported by one or multiple sources, but only cite sources that directly support the claim. Do not add unnecessary citations.
+
+If none of the sources contain relevant information to answer the query, politely inform the user that an answer cannot be provided, and do not use any citations.
+
+If the query is not related to medicine, politely decline to answer.
+
+<query>{query}</query>
+
+<sources>
+{sources}
+</sources>
+"""
+
+
 def generate_answer(
     client: genai.Client, query: str, articles: list[Article]
 ) -> list[Statement]:
     response = client.models.generate_content(
         model="gemini-2.5-flash-preview-04-17",
-        contents=f"Answer the query based solely on the provided sources. The answer should be less than 100 words. Justify the answer by citing from the sources. Refuse to answer non-medical related query.\n\n<query>{query}</query>\n\n<sources>{format_sources(articles)}</sources>",
+        contents=ANSWER_INSTRUCTION.format(
+            query=query, sources=format_sources(articles)
+        ),
         config={
             "response_mime_type": "application/json",
             "response_schema": list[Statement],
@@ -93,6 +113,9 @@ def semantic_scholar(
             title=article["title"],
             summary=article["tldr"]["text"] if article["tldr"] else "",
             abstract=article["abstract"],
+            venue=article["venue"],
+            year=article["year"],
+            citations=article["citationCount"],
         )
         articles.append(article)
     return articles
@@ -122,7 +145,7 @@ def pubmed(query: str, top_k: int = 10, db: str = "pubmed"):
 
 
 def main():
-    semantic_scholar_client = httpx.Client(transport=httpx.HTTPTransport(retries=1))
+    semantic_scholar_client = httpx.Client(transport=httpx.HTTPTransport(retries=3))
     gemini_client = genai.Client(api_key=os.environ["GOOGLE_API_KEY"])
 
     st.title("Ask ~~Jeeves~~ Elna")
@@ -142,22 +165,18 @@ def main():
                     citations = {}
                     statements = generate_answer(gemini_client, query, papers)
                     for statement in statements:
-                        if statement.citation:
-                            if not (
-                                citation_id := citations.get(
-                                    statement.citation.source_id
-                                )
-                            ):
+                        sentence = statement.text
+                        for source_id in statement.source_ids or []:
+                            if not (citation_id := citations.get(source_id)):
                                 citation_id = len(citations) + 1
-                                citations[statement.citation.source_id] = citation_id
-                            sentences.append(f"{statement.text}[^{citation_id}]")
-                        else:
-                            sentences.append(statement.text)
+                                citations[source_id] = citation_id
+                            sentence += f"[^{citation_id}] "
+                        sentences.append(sentence.strip())
                     answer = " ".join(sentences)
                     footnotes = ""
                     if citations:
                         footnotes = "\n".join(
-                            f"[^{v}]: [{paper_map[k].title}](https://doi.org/{paper_map[k].id})"
+                            f"[^{v}]: :grey-badge[:material/attribution: {paper_map[k].citations}] [{paper_map[k].title}](https://doi.org/{paper_map[k].id}). _{paper_map[k].venue}_ {paper_map[k].year}."
                             for k, v in citations.items()
                         )