# 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