Feature/source citations in query result

.gitignore

@@ -35,12 +35,14 @@ dist/
 *.log
 
 # Data stores (explicit)
-chroma_db
 notion_data.json
 
 # SQLite database
 workmate.db
 
+# ChromaDB
+workmate_db/
+
 # Antigravity SKILLS
 .agent/
 CLAUDE.md

src/backend/config.py

@@ -13,7 +13,7 @@ class Settings(BaseSettings):
     GEMINI_API_KEY: str = ""
     NOTION_TOKEN: str = ""
 
-    model_config = {"env_file": ".env"}
+    model_config = {"env_file": ".env", "extra": "ignore"}
 
 
 settings = Settings()

src/backend/llm/gemini_client.py

@@ -20,15 +20,17 @@ class GeminiClient:
     """
 
     def __init__(self, model_id: Optional[str] = None):
-        api_key = get_required_env("GEMINI_API_KEY")
+        api_key = get_required_env("GEMINI_KEY")
         self.client = genai.Client(api_key=api_key)
         self.model_id = model_id or DEFAULT_GEMINI_MODEL_ID
 
-    def ask_workmate(self, chunks: List[Dict[str, Any]], user_question: str) -> str:
+    def ask_workmate(
+        self, chunks: List[Dict[str, Any]], user_question: str, debug: bool = False
+    ) -> str:
         """
         Generate an answer using ONLY the provided top-k context chunks.
         """
-        final_prompt = get_rag_prompt(chunks, user_question)
+        final_prompt = get_rag_prompt(chunks, user_question, debug)
 
         # Keep outputs grounded + stable
         cfg = types.GenerateContentConfig(
@@ -56,4 +58,4 @@ def ask_workmate(self, chunks: List[Dict[str, Any]], user_question: str) -> str:
                 )
                 return f"I'm currently unable to respond due to API rate limits. Please try again in about {retry_secs} seconds."
             logger.error(f"❌ Gemini error: {e}")
-            return "Sorry, something went wrong while generating a response. Please try again."
+            return "Sorry, something went wrong while generating a response. Please try again."

src/backend/llm/prompts.py

@@ -13,35 +13,47 @@
 - Always include citations in the required format.
 
 OUTPUT FORMAT (exact):
-Answer: <a thorough, multi-sentence answer that fully addresses the question>
-Sources:
-- <chunk_id> | <page_title>
+Answer: <1-3 sentences>
 Confidence: <High/Medium/Low>
 """.strip()
 
 
 def _format_chunks(chunks: List[Dict[str, Any]]) -> str:
     """
-    Build a clean, deterministic context block with explicit chunk IDs + page titles.
-    Each chunk must contain: chunk_id, page_title, text
+    Build a clean, deterministic context block with explicit chunk IDs + page titles + section + paragraph.
+    Each chunk must contain: chunk_id, page_title, text, (optionally section, paragraph)
     """
-    lines = ["CONTEXT CHUNKS (Top-5):"]
+    lines = ["CONTEXT CHUNKS:"]
     for i, ch in enumerate(chunks, start=1):
         chunk_id = ch.get("chunk_id", f"chunk_{i}")
         page_title = ch.get("page_title", "Unknown Page")
+        section = ch.get("section", "")
+        paragraph = ch.get("paragraph", "")
         text = (ch.get("text") or "").strip()
 
-        lines.append(f"\n[CHUNK {i}] id={chunk_id} | page={page_title}\n{text}")
+        meta_info = f"id={chunk_id} | page={page_title}"
+        if section:
+            meta_info += f" | section={section}"
+        if paragraph:
+            meta_info += f" | paragraph={paragraph}"
+
+        lines.append(f"\n[CHUNK {i}] {meta_info}\n{text}")
 
     return "\n".join(lines).strip()
 
 
-def get_rag_prompt(chunks: List[Dict[str, Any]], question: str) -> str:
+def get_rag_prompt(
+    chunks: List[Dict[str, Any]], question: str, debug: bool = False
+) -> str:
     """
     Combines top-k chunks with the user question into a structured prompt.
     """
     context_block = _format_chunks(chunks)
 
+    debug_instruction = ""
+    if debug:
+        debug_instruction = "\n    5) DEBUG MODE ACTIVE: In addition to the standard citation, also append the chunk ID, e.g., [Document name, section, paragraph, Chunk ID: <id>]."
+
     return f"""
     {context_block}
 
@@ -52,7 +64,7 @@ def get_rag_prompt(chunks: List[Dict[str, Any]], question: str) -> str:
     1) Answer ONLY using the context chunks above.
     2) If multiple chunks mention the topic, prefer the most specific one.
     3) If no chunk clearly answers, reply exactly: "I cannot find this in your Notion docs."
-    4) Output MUST match the required format and include citations (chunk_id + page_title).
+    4) Output MUST include inline citations for the source of its findings in the exact format: [Document name, section, paragraph] (section and paragraph are optional if not provided).{debug_instruction}
 
     Remember: Do not use outside knowledge.
-    """.strip()
+    """.strip()

src/backend/load/chroma_manager.py

@@ -11,9 +11,13 @@
 
 logger = logging.getLogger(__name__)
 
+# Resolve the project root (3 levels up from src/backend/load/chroma_manager.py)
+PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../"))
+DEFAULT_DB_PATH = os.path.join(PROJECT_ROOT, "workmate_db")
+
 
 class ChromaManager:
-    def __init__(self, db_path="chroma_db", collection_name="notion_docs"):
+    def __init__(self, db_path=DEFAULT_DB_PATH, collection_name="notion_docs"):
         """
         Initialize the ChromaDB client and collection.
         :param db_path: Path to the persistent database directory.

src/backend/load/google_embedder.py

@@ -18,11 +18,9 @@ class GoogleEmbedder(EmbeddingFunction):
     def __init__(self, model_name="gemini-embedding-001"):
         self.model_name = model_name
 
-        api_key = os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")
+        api_key = os.getenv("GEMINI_KEY")
         if not api_key:
-            print(
-                "⚠️ WARNING: GEMINI_API_KEY or GOOGLE_API_KEY not found in environment."
-            )
+            print("⚠️ WARNING: GEMINI_KEY not found in environment.")
 
         self.client = genai.Client(api_key=api_key)
 

src/backend/routers/chat.py

@@ -18,9 +18,7 @@ def get_chroma_manager() -> ChromaManager:
     global _chroma_manager
     if _chroma_manager is None:
         try:
-            _chroma_manager = ChromaManager(
-                db_path="chroma_db", collection_name="notion_docs"
-            )
+            _chroma_manager = ChromaManager()
         except Exception as e:
             logger.error(f"Failed to initialize ChromaManager: {e}")
             raise HTTPException(
@@ -56,7 +54,7 @@ async def ask_question(
         # Step 1: Retrieval
         results = chroma.query(request.question, n_results=3)
 
-        # Step 2: Augmentation – build the context string
+        # Step 2: Augmentation – extract chunks
         chunks = []
         if results and results.get("documents") and results["documents"][0]:
             docs = results["documents"][0]
@@ -81,6 +79,7 @@ async def ask_question(
         answer = gemini.ask_workmate(
             chunks=chunks,
             user_question=request.question,
+            debug=request.debug,
         )
 
         return ChatResponse(answer=answer)

src/backend/schemas/chat.py

@@ -2,8 +2,15 @@
 
 
 class ChatRequest(BaseModel):
-    question: str = Field(..., description="The user's question to ask the RAG pipeline.")
+    question: str = Field(
+        ..., description="The user's question to ask the RAG pipeline."
+    )
+    debug: bool = Field(
+        False, description="Enable debug mode to include chunk IDs in the response."
+    )
 
 
 class ChatResponse(BaseModel):
-    answer: str = Field(..., description="The generated context-aware answer from the Gemini LLM.")
+    answer: str = Field(
+        ..., description="The generated context-aware answer from the Gemini LLM."
+    )

src/backend/transform/notion_ingestory.py

@@ -1,5 +1,4 @@
 import json
-import os
 from pathlib import Path
 from langchain_text_splitters import (
     MarkdownHeaderTextSplitter,
@@ -9,7 +8,7 @@
 
 # Default path resolving to src/data/notion_data.json
 DEFAULT_DATA_PATH = str(
-    Path(__file__).resolve().parent.parent.parent / "data" / "notion_data.json"
+    Path(__file__).resolve().parent.parent.parent / "data/notion_data.json"
 )