Create generator.py

rag_pipeline/generator.py

@@ -0,0 +1,73 @@
+"""
+Final answer generation using Anthropic API.
+Returns answer + the exact prompt sent to the model for full transparency.
+"""
+
+import os
+from typing import List, Tuple
+import anthropic
+
+from .reranker import RerankResult
+
+
+SYSTEM_PROMPT = """You are a precise scientific assistant. Answer questions using ONLY
+the provided context excerpts from the paper. Follow these rules:
+
+1. Cite the specific context excerpt (by number) that supports each claim.
+2. If the context does not contain enough information, say so explicitly.
+3. Do not add knowledge from outside the provided context.
+4. Be concise and precise — this is a scientific document.
+5. Format your answer clearly with citations like [Excerpt 1], [Excerpt 2].
+"""
+
+
+def build_prompt(query: str, kept_results: List[RerankResult]) -> str:
+    context_blocks = []
+    for i, r in enumerate(kept_results, 1):
+        context_blocks.append(f"[Excerpt {i}]\n{r.text}")
+
+    context_str = "\n\n".join(context_blocks)
+
+    return f"""Here are the relevant excerpts retrieved from the scientific paper:
+
+{context_str}
+
+---
+
+Question: {query}
+
+Answer based strictly on the excerpts above. Cite which excerpt(s) support each point."""
+
+
+def generate_answer(
+    query: str,
+    kept_results: List[RerankResult],
+    model: str = "claude-haiku-4-5-20251001",
+    api_key: str = None,
+) -> Tuple[str, str]:
+    """
+    Returns (answer, prompt_sent_to_model).
+    Shows the prompt for transparency.
+    """
+    if api_key:
+        os.environ["ANTHROPIC_API_KEY"] = api_key
+
+    key = os.environ.get("ANTHROPIC_API_KEY")
+    if not key:
+        return "❌ No API key provided.", ""
+
+    prompt = build_prompt(query, kept_results)
+
+    try:
+        client = anthropic.Anthropic(api_key=key)
+        message = client.messages.create(
+            model=model,
+            max_tokens=1024,
+            system=SYSTEM_PROMPT,
+            messages=[{"role": "user", "content": prompt}],
+        )
+        answer = message.content[0].text
+    except Exception as e:
+        answer = f"❌ API error: {e}"
+
+    return answer, prompt