Update scripts/router_chain.py

scripts/router_chain.py

@@ -37,12 +37,33 @@ User request: {input}
                 chain = prompt | llm | StrOutputParser()
                 return {"result": chain.invoke({"input": input_dict["input"]})}
 
+            # elif category == "summarize":
+            #     prompt = ChatPromptTemplate.from_template(
+            #         "Provide a concise summary about: {input}\nSummary:"
+            #     )
+            #     chain = prompt | llm | StrOutputParser()
+            #     return {"result": chain.invoke({"input": input_dict["input"]})}
+
             elif category == "summarize":
-                prompt = ChatPromptTemplate.from_template(
-                    "Provide a concise summary about: {input}\nSummary:"
-                )
-                chain = prompt | llm | StrOutputParser()
-                return {"result": chain.invoke({"input": input_dict["input"]})}
+                # 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 == "calculate":
                 prompt = ChatPromptTemplate.from_template(