Update scripts/router_chain.py

scripts/router_chain.py

@@ -1,33 +1,32 @@
-# scripts/router_simple.py
-from typing import Dict, Any
-from langchain.chat_models import ChatOpenAI
-from langchain.prompts import ChatPromptTemplate
-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)
-
-    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:
-                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:
-                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
-
-    return Router()
+from typing import Dict, Any
+from langchain.chat_models import ChatOpenAI
+from langchain.prompts import ChatPromptTemplate
+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)
+
+    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:
+                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:
+                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
+
+    return Router()