modified graph.py

graph.py

@@ -7,26 +7,38 @@ from transformers import pipeline
 # 🧠 Retrieval Node
 def retrieve(state: Dict[str, Any]) -> Dict[str, Any]:
     query = state["query"]
-    retriever = state["retriever"]
+    retriever = state.get("retriever")        # optional
     top_k = state.get("top_k", 3)
     summarizer = state.get("summarizer")
-    raw_text = state.get("raw_text")  # <-- add this
-
-    # Retrieve top-k documents
-    retrieved_docs = retriever.invoke(query, k=top_k)
-
-    # --- NEW: Use uploaded file if no docs retrieved ---
-    if not retrieved_docs and raw_text:
-        context = raw_text  # use uploaded file directly
-    elif not retrieved_docs:
-        # fallback to sample.txt only if no uploaded file
-        with open("data/sample.txt", "r", encoding="utf-8") as f:
-            context = f.read().strip()
-    else:
-        # Limit each chunk to 500 chars to reduce prompt size
+    raw_text = state.get("raw_text")          # uploaded file
+
+    retrieved_docs = []
+    context = ""
+
+    # --- Use FAISS retriever if available ---
+    if retriever:
+        try:
+            retrieved_docs = retriever.invoke(query, k=top_k)
+        except Exception:
+            retrieved_docs = []
+
+    # --- Determine context ---
+    if retrieved_docs:
+        # Limit each chunk to 500 chars
         context = "\n\n".join([doc.page_content[:500] for doc in retrieved_docs])
+    elif raw_text:
+        # Use uploaded file directly if no retrieved docs
+        context = raw_text
+    else:
+        # Fallback to sample.txt
+        sample_path = "data/sample.txt"
+        if os.path.exists(sample_path):
+            with open(sample_path, "r", encoding="utf-8") as f:
+                context = f.read().strip()
+        else:
+            context = "No content available."
 
-    # Summarize if context too large
+    # --- Summarize if context too long ---
     if len(context) > 1500 and summarizer:
         context = summarizer(
             context,
@@ -42,7 +54,7 @@ def retrieve(state: Dict[str, Any]) -> Dict[str, Any]:
         "context": context,
         "llm": state["llm"],
         "temperature": state.get("temperature", 0.7),
-        "raw_text": raw_text  # keep it in state for next steps
+        "raw_text": raw_text
     }
 
 # 🧠 Generation Node