Spaces:
Sleeping
Sleeping
| # agent_workflow.py | |
| import logging | |
| from typing import Dict, List, Any, Annotated, TypedDict | |
| from langchain_openai import ChatOpenAI | |
| from langchain_core.documents import Document | |
| from langchain_core.messages import SystemMessage, HumanMessage, AIMessage | |
| from langchain_community.tools.tavily_search import TavilySearchResults | |
| from langgraph.graph import StateGraph, START, END | |
| from langgraph.graph.message import add_messages | |
| # Logging configuration | |
| logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') | |
| logger = logging.getLogger(__name__) | |
| class AgentState(TypedDict): | |
| """Agent state for the workflow""" | |
| messages: Annotated[list, add_messages] | |
| context: List[Document] | |
| next_tool: str | |
| question: str | |
| retrieved_contexts: List[Document] | |
| context_count: int | |
| class AgentWorkflow: | |
| """Agent workflow with intelligent routing logic""" | |
| def __init__(self, rag_tool, tavily_max_results: int = 5): | |
| """ Initialize the agent workflow """ | |
| self.rag_tool = rag_tool | |
| self.tavily_tool = TavilySearchResults(max_results=tavily_max_results) | |
| # LLMs for routing and evaluation | |
| self.router_llm = ChatOpenAI(model="gpt-4o-mini", temperature=0, max_tokens=50) | |
| self.evaluator_llm = ChatOpenAI(model="gpt-4o-mini", temperature=0) | |
| self.final_llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.7) | |
| # Compile the workflow | |
| self.compiled_workflow = self._build_workflow() | |
| def evaluate_response_quality(self, question: str, response: str) -> bool: | |
| """ Evaluates if the response is satisfactory """ | |
| prompt = f"""Evaluate if this response to "{question}" is UNSATISFACTORY: | |
| "{response}" | |
| UNSATISFACTORY CRITERIA (if ANY ONE is present, the response is UNSATISFACTORY): | |
| 1. Contains "consult experts", "specialized training", "I'm sorry" | |
| 2. Doesn't provide concrete steps for "how to" questions | |
| 3. Gives general advice rather than specific methods | |
| 4. Redirects the user without directly answering | |
| Quick example: | |
| Q: "How do I train my dog to sit?" | |
| UNSATISFACTORY: "Consult a professional trainer." | |
| SATISFACTORY: "1. Use treats... 2. Be consistent..." | |
| Reply only "UNSATISFACTORY" or "SATISFACTORY". | |
| When in doubt, choose "UNSATISFACTORY". | |
| """ | |
| evaluation = self.evaluator_llm.invoke([SystemMessage(content=prompt)]) | |
| result = evaluation.content.strip().upper() | |
| is_satisfactory = "UNSATISFACTORY" not in result | |
| logger.info(f"[Evaluation] Response rated: {'SATISFACTORY' if is_satisfactory else 'UNSATISFACTORY'}") | |
| return is_satisfactory | |
| def _build_workflow(self): | |
| """Builds and compiles the agent workflow""" | |
| # 1. Node for intelligent routing | |
| def smart_router(state): | |
| """Determines if the question is about dogs or not""" | |
| messages = state["messages"] | |
| last_message = [msg for msg in messages if isinstance(msg, HumanMessage)][-1] | |
| question = last_message.content | |
| # Prompt using reverse logic - asking if it's NOT related to dogs | |
| router_prompt = f"""Evaluate if this question is UNRELATED to dogs, puppies, or canine care: | |
| Question: "{question}" | |
| INDICATORS OF NON-DOG QUESTIONS (if ANY ONE is present, mark as "NOT_DOG_RELATED"): | |
| 1. Questions about weather, time, locations, or general information | |
| 2. Questions about other animals (cats, birds, etc.) | |
| 3. Questions about technology, politics, or human activities | |
| 4. Any question that doesn't explicitly mention or imply dogs/puppies/canines | |
| Example check: | |
| Q: "What is the weather in Paris today?" | |
| This is NOT_DOG_RELATED (about weather) | |
| Q: "How do I train my puppy to sit?" | |
| This is DOG_RELATED (explicitly about puppy training) | |
| Reply ONLY with "NOT_DOG_RELATED" or "DOG_RELATED". | |
| When in doubt, choose "NOT_DOG_RELATED". | |
| """ | |
| router_response = self.router_llm.invoke([SystemMessage(content=router_prompt)]) | |
| result = router_response.content.strip().upper() | |
| is_dog_related = "NOT_DOG_RELATED" not in result | |
| logger.info(f"[Smart Router] Question {'' if is_dog_related else 'NOT '}related to dogs") | |
| # If the question is not related to dogs, go directly to out_of_scope | |
| if not is_dog_related: | |
| return { | |
| "next_tool": "out_of_scope", | |
| "question": question | |
| } | |
| # If the question is related to dogs, go to the RAG tool | |
| return { | |
| "next_tool": "rag_tool", | |
| "question": question | |
| } | |
| # 2. Node for out-of-scope questions | |
| def out_of_scope(state): | |
| """Informs that the assistant only answers questions about dogs""" | |
| out_of_scope_message = AIMessage( | |
| content="I'm sorry, but I specialize only in canine care and puppy education. I cannot answer this question as it is outside my area of expertise. Feel free to ask me any questions about dogs and puppies!" | |
| ) | |
| return { | |
| "messages": [out_of_scope_message], | |
| "next_tool": "final_response" | |
| } | |
| # 3. Node for using the RAG tool | |
| def use_rag_tool(state): | |
| """Uses the RAG tool for dog-related questions""" | |
| question = state["question"] | |
| # Call the RAG tool directly | |
| rag_result = self.rag_tool.invoke(question) | |
| rag_response = rag_result["messages"][0].content | |
| context = rag_result.get("context", []) | |
| sources_info = rag_result.get("sources_info", []) | |
| total_chunks = rag_result.get("total_chunks", 0) | |
| # Evaluate the quality of the response | |
| is_satisfactory = self.evaluate_response_quality(question, rag_response) | |
| # Format detailed source information | |
| sources_text = "" | |
| if sources_info: | |
| sources_text = f"*Based on {total_chunks} chunk(s):*\n" | |
| for source in sources_info: | |
| sources_text += f"- *Chunk {source['chunk_number']} - {source['source']} (Page: {source['page']})*\n" | |
| else: | |
| sources_text = "*Source: Livre \"Puppies for Dummies\"*" | |
| # Create an AI message with the response and detailed sources | |
| response_message = AIMessage(content=f"[Using RAG tool] - {sources_text}\n{rag_response}") | |
| # If the response is not satisfactory, prepare to use Tavily | |
| next_tool = "final_response" if is_satisfactory else "need_tavily" | |
| return { | |
| "messages": [response_message], | |
| "context": context, | |
| "sources_info": sources_info, | |
| "next_tool": next_tool, | |
| "retrieved_contexts": context, | |
| "context_count": len(context) | |
| } | |
| # 4. Node for using the Tavily tool | |
| def use_tavily_tool(state): | |
| """Uses the Tavily tool as a fallback for dog-related questions""" | |
| question = state["question"] | |
| # Call Tavily | |
| tavily_result = self.tavily_tool.invoke(question) | |
| # Format the sources and prepare content for LLM | |
| sources_text = "" | |
| sources_content = "" | |
| has_useful_results = False | |
| if tavily_result and len(tavily_result) > 0: | |
| sources_text = f"*Based on {len(tavily_result[:3])} internet source(s):*\n" | |
| for i, result in enumerate(tavily_result[:3], 1): | |
| title = result.get('title', 'Unknown Source') | |
| url = result.get('url', '') | |
| content = result.get('content', '') | |
| if content and len(content.strip()) > 50: | |
| has_useful_results = True | |
| # Format source in italics | |
| domain = url.split('/')[2] if url and '/' in url else 'Web' | |
| sources_text += f"- *Source {i} - {domain}: {title}*\n" | |
| # Collect content for LLM processing | |
| sources_content += f"Source {i} ({title}): {content[:300]}...\n\n" | |
| if not has_useful_results: | |
| # No useful results found | |
| dont_know_message = AIMessage( | |
| content=f"[Using Tavily tool] - *No reliable internet sources found for this question.*\n\nI couldn't find specific information about '{question}' in my knowledge base or through online searches. This might be a specialized topic that requires expertise from professionals in the field of canine education." | |
| ) | |
| return { | |
| "messages": [dont_know_message], | |
| "next_tool": "final_response" | |
| } | |
| # Generate a proper response using LLM based on the sources | |
| response_prompt = f"""Based on the following internet sources, provide a clear and helpful answer to the question: "{question}" | |
| {sources_content} | |
| Instructions: | |
| - Provide a comprehensive answer based on the sources above | |
| - Focus on practical, actionable information | |
| - If the sources contain contradictory information, mention the different perspectives | |
| - Keep the response clear and well-structured | |
| - Do not mention the sources in your response (they will be displayed separately) | |
| """ | |
| try: | |
| llm_response = self.final_llm.invoke([SystemMessage(content=response_prompt)]) | |
| generated_answer = llm_response.content | |
| except Exception as e: | |
| logger.error(f"Error generating Tavily response: {e}") | |
| generated_answer = "I found some relevant information but couldn't process it properly." | |
| # Create the final formatted message | |
| response_message = AIMessage(content=f"[Using Tavily tool] - {sources_text}\n{generated_answer}") | |
| return { | |
| "messages": [response_message], | |
| "next_tool": "final_response" | |
| } | |
| # 5. Node for cases where no source has a satisfactory answer | |
| def say_dont_know(state): | |
| """Responds when no source has useful information""" | |
| question = state["question"] | |
| dont_know_message = AIMessage(content=f"I'm sorry, but I couldn't find specific information about '{question}' in my knowledge base or through online searches. This might be a specialized topic that requires expertise from professionals in the field of canine education.") | |
| return { | |
| "messages": [dont_know_message], | |
| "next_tool": "final_response" | |
| } | |
| # 6. Node for generating the final response | |
| def generate_final_response(state): | |
| """Generates a final response based on tool results""" | |
| messages = state["messages"] | |
| original_question = state["question"] | |
| # Find tool messages | |
| tool_responses = [msg.content for msg in messages if isinstance(msg, AIMessage)] | |
| # If no tool messages, return a default response | |
| if not tool_responses: | |
| return {"messages": [AIMessage(content="I couldn't find information about your dog-related question.")]} | |
| # Take the last tool message as the main content | |
| tool_content = tool_responses[-1] | |
| # If the tool message already contains detailed sources, return it as-is | |
| if "[Using RAG tool]" in tool_content or "[Using Tavily tool]" in tool_content: | |
| # Already contains detailed sources, return as-is | |
| return {"messages": [AIMessage(content=tool_content)]} | |
| # Use an LLM to generate a coherent final response but preserve source markers | |
| system_prompt = f"""Here are the search results for the dog-related question: "{original_question}" | |
| {tool_content} | |
| Formulate a clear, helpful, and concise response based ONLY on these results. | |
| IMPORTANT: If the search results start with "[Using RAG tool]" or "[Using Tavily tool]", keep these markers exactly as they are at the beginning of your response. | |
| If the search results contain useful information, include it in your response rather than saying "I don't know". | |
| Say "I don't know" only if the search results contain no useful information. | |
| """ | |
| response = self.final_llm.invoke([SystemMessage(content=system_prompt)]) | |
| return {"messages": [response]} | |
| # 7. Routing function | |
| def route_to_next_tool(state): | |
| next_tool = state["next_tool"] | |
| if next_tool == "rag_tool": | |
| return "use_rag_tool" | |
| elif next_tool == "out_of_scope": | |
| return "out_of_scope" | |
| elif next_tool == "tavily_tool": | |
| return "use_tavily_tool" | |
| elif next_tool == "need_tavily": | |
| return "use_tavily_tool" | |
| elif next_tool == "say_dont_know": | |
| return "say_dont_know" | |
| elif next_tool == "final_response": | |
| return "generate_response" | |
| else: | |
| return "generate_response" | |
| # 8. Building the LangGraph | |
| workflow = StateGraph(AgentState) | |
| # Adding nodes | |
| workflow.add_node("smart_router", smart_router) | |
| workflow.add_node("out_of_scope", out_of_scope) | |
| workflow.add_node("use_rag_tool", use_rag_tool) | |
| workflow.add_node("use_tavily_tool", use_tavily_tool) | |
| workflow.add_node("say_dont_know", say_dont_know) | |
| workflow.add_node("generate_response", generate_final_response) | |
| # Connections | |
| workflow.add_edge(START, "smart_router") | |
| workflow.add_conditional_edges("smart_router", route_to_next_tool) | |
| workflow.add_edge("out_of_scope", "generate_response") | |
| workflow.add_conditional_edges("use_rag_tool", route_to_next_tool) | |
| workflow.add_conditional_edges("use_tavily_tool", route_to_next_tool) | |
| workflow.add_edge("say_dont_know", "generate_response") | |
| workflow.add_edge("generate_response", END) | |
| # Compile the graph | |
| return workflow.compile() | |
| def process_question(self, question: str): | |
| """ Process a question with the agent workflow """ | |
| # Invoke the workflow | |
| result = self.compiled_workflow.invoke({ | |
| "messages": HumanMessage(content=question), | |
| "context": [], | |
| "next_tool": "", | |
| "question": "", | |
| "retrieved_contexts": [], | |
| "context_count": 0 | |
| }) | |
| return result | |
| def get_final_response(self, result): | |
| """Extract the final response from the agent result with source information.""" | |
| messages = result.get("messages", []) | |
| if not messages: | |
| return "No response available." | |
| # Get the last AI message | |
| last_message = None | |
| for msg in reversed(messages): | |
| if hasattr(msg, 'content') and msg.content: | |
| last_message = msg | |
| break | |
| if not last_message: | |
| return "No valid response found." | |
| response_content = last_message.content | |
| # Extract and store source information in result for main.py to use | |
| if "Tavily" in response_content and "Source" in response_content: | |
| # Extract Tavily sources from the response content | |
| tavily_sources = [] | |
| lines = response_content.split('\n') | |
| for line in lines: | |
| if line.strip().startswith('- *Source') and ':' in line: | |
| # Parse line like "- *Source 1 - domain.com: Title*" | |
| try: | |
| # Extract source number, domain, and title | |
| source_part = line.split('- *Source')[1].split('*')[0] | |
| if ' - ' in source_part and ':' in source_part: | |
| parts = source_part.split(' - ', 1) | |
| source_num = parts[0].strip() | |
| domain_title = parts[1] | |
| if ':' in domain_title: | |
| domain, title = domain_title.split(':', 1) | |
| tavily_sources.append({ | |
| 'source_num': source_num, | |
| 'domain': domain.strip(), | |
| 'title': title.strip() | |
| }) | |
| except: | |
| continue | |
| # Store Tavily sources in result | |
| result['tavily_sources'] = tavily_sources | |
| return response_content | |