Update scripts/router_chain.py

scripts/router_chain.py

@@ -44,23 +44,52 @@ User request: {input}
             #     chain = prompt | llm | StrOutputParser()
             #     return {"result": chain.invoke({"input": input_dict["input"]})}
 
+            #elif category == "summarize":
+            #    # 1. Use RAG to retrieve relevant docs
+            #    rag_result = general_qa({"query": input_dict["input"]})
+
+            #    # 2. Extract docs and prepare text
+            #    source_docs = rag_result.get("source_documents", [])
+            #    combined_text = "\n\n".join([doc.page_content for doc in source_docs])
+
+            #    # 3. Run the summarizer chain on the retrieved text
+            #    from scripts.summarizer import get_summarizer
+            #    summarizer_chain = get_summarizer()
+            #    summary = summarizer_chain.run(combined_text)
+
+            #    # 4. Add sources if any
+            #    sources = list({str(doc.metadata.get("source", "unknown")) for doc in source_docs})
+            #    if sources:
+            #        summary += f"\n\n📚 Sources: {', '.join(sources)}"
+
+            #    return {"result": summary}
+            
+
             elif category == "summarize":
-                # 1. Use RAG to retrieve relevant docs
+                # 1) Retrieve relevant documents via your existing RAG chain
                 rag_result = general_qa({"query": input_dict["input"]})
 
-                # 2. Extract docs and prepare text
-                source_docs = rag_result.get("source_documents", [])
-                combined_text = "\n\n".join([doc.page_content for doc in source_docs])
+                # 2) Get the retrieved docs (already LangChain Document objects)
+                source_docs = rag_result.get("source_documents", []) or []
 
-                # 3. Run the summarizer chain on the retrieved text
+                # 3) Build the summarizer and prepare the docs list
+                from langchain.docstore.document import Document
                 from scripts.summarizer import get_summarizer
+
                 summarizer_chain = get_summarizer()
-                summary = summarizer_chain.run(combined_text)
 
-                # 4. Add sources if any
-                sources = list({str(doc.metadata.get("source", "unknown")) for doc in source_docs})
-                if sources:
-                    summary += f"\n\n📚 Sources: {', '.join(sources)}"
+                # If retrieval returned nothing, fall back to summarizing the user’s text
+                docs = source_docs if source_docs else [Document(page_content=input_dict["input"])]
+
+                # 4) Summarize — load_summarize_chain returns {"output_text": "..."}
+                out = summarizer_chain.invoke(docs)
+                summary = out["output_text"] if isinstance(out, dict) and "output_text" in out else str(out)
+
+                # 5) Append sources (only if we actually had retrieved docs)
+                if source_docs:
+                    sources = sorted({str(d.metadata.get("source", "unknown")) for d in source_docs})
+                    if sources:
+                        summary += f"\n\n📚 Sources: {', '.join(sources)}"
 
                 return {"result": summary}