Update scripts/router_chain.py

scripts/router_chain.py

@@ -1,32 +1,57 @@
 from typing import Dict, Any
 from langchain_openai import ChatOpenAI
 from langchain.prompts import ChatPromptTemplate
+from langchain.schema import StrOutputParser
 from scripts.rag_chat import build_general_qa_chain
 
 def build_router_chain(model_name=None):
     general_qa = build_general_qa_chain(model_name=model_name)
     llm = ChatOpenAI(model_name=model_name or "gpt-4o-mini", temperature=0.0)
 
+    # This prompt asks the LLM to choose which "mode" to use
+    router_prompt = ChatPromptTemplate.from_template("""
+You are a routing assistant for a chatbot. 
+Classify the following user request into one of these categories:
+- "code" for programming or debugging
+- "summarize" for summary requests
+- "calculate" for math or numeric calculations
+- "general" for general Q&A using course files
+
+Return ONLY the category word.
+
+User request: {input}
+""")
+
+    router_chain = router_prompt | llm | StrOutputParser()
+
     class Router:
         def invoke(self, input_dict: Dict[str, Any]):
-            text = input_dict.get("input", "").lower()
-            if "code" in text or "program" in text or "debug" in text:
+            category = router_chain.invoke({"input": input_dict["input"]}).strip().lower()
+
+            print(f"[ROUTER] User query routed to category: {category}")
+
+            if category == "code":
                 prompt = ChatPromptTemplate.from_template(
                     "As a coding assistant, help with this Python question.\nQuestion: {input}\nAnswer:"
                 )
-                chain = prompt | llm
-                return {"result": chain.invoke({"input": input_dict["input"]}).content}
-            elif "summarize" in text or "summary" in text:
+                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
-                return {"result": chain.invoke({"input": input_dict["input"]}).content}
-            elif "calculate" in text or any(char.isdigit() for char in text):
-                return {"result": "For calculations, please ask a specific calculation or provide more context."}
-            else:
-                # Use RAG chain
-                result = general_qa({"query": input_dict["input"]})
-                return result
+                chain = prompt | llm | StrOutputParser()
+                return {"result": chain.invoke({"input": input_dict["input"]})}
+
+            elif category == "calculate":
+                prompt = ChatPromptTemplate.from_template(
+                    "Solve the following calculation step-by-step:\n{input}"
+                )
+                chain = prompt | llm | StrOutputParser()
+                return {"result": chain.invoke({"input": input_dict["input"]})}
+
+            else:  # "general"
+                return general_qa({"query": input_dict["input"]})
 
     return Router()