Spaces:
Sleeping
Sleeping
File size: 1,559 Bytes
dfda80f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 |
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()
|