veryfinal.py

"""
Enhanced Multi-LLM Agent System - CORRECTED VERSION
Fixes the issue where questions are returned as answers
"""

import os
import time
import random
import operator
from typing import List, Dict, Any, TypedDict, Annotated
from dotenv import load_dotenv

from langchain_core.tools import tool
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_community.document_loaders import WikipediaLoader
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
from langchain_groq import ChatGroq

load_dotenv()

# Enhanced system prompt for proper question-answering
ENHANCED_SYSTEM_PROMPT = (
    "You are a helpful assistant tasked with answering questions using available tools. "
    "Follow these guidelines:\n"
    "1. Read the question carefully and understand what is being asked\n"
    "2. Use available tools when you need external information\n"
    "3. Provide accurate, specific answers based on the information you find\n"
    "4. For numbers: don't use commas or units unless specified\n"
    "5. For strings: don't use articles or abbreviations, write digits in plain text\n"
    "6. Always end with 'FINAL ANSWER: [YOUR ANSWER]' where [YOUR ANSWER] is concise\n"
    "7. Never repeat the question as your answer\n"
    "8. If you cannot find the answer, state 'Information not available'\n"
)

# ---- Tool Definitions ----
@tool
def multiply(a: int, b: int) -> int:
    """Multiply two integers and return the product."""
    return a * b

@tool
def add(a: int, b: int) -> int:
    """Add two integers and return the sum."""
    return a + b

@tool
def subtract(a: int, b: int) -> int:
    """Subtract the second integer from the first and return the difference."""
    return a - b

@tool
def divide(a: int, b: int) -> float:
    """Divide the first integer by the second and return the quotient."""
    if b == 0:
        raise ValueError("Cannot divide by zero.")
    return a / b

@tool
def modulus(a: int, b: int) -> int:
    """Return the remainder when dividing the first integer by the second."""
    return a % b

@tool
def optimized_web_search(query: str) -> str:
    """Perform web search using TavilySearchResults."""
    try:
        time.sleep(random.uniform(0.7, 1.5))
        search_tool = TavilySearchResults(max_results=3)
        docs = search_tool.invoke({"query": query})
        return "\n\n---\n\n".join(
            f"<Doc url='{d.get('url','')}'>{d.get('content','')[:800]}</Doc>"
            for d in docs
        )
    except Exception as e:
        return f"Web search failed: {e}"

@tool
def optimized_wiki_search(query: str) -> str:
    """Perform Wikipedia search and return content."""
    try:
        time.sleep(random.uniform(0.3, 1))
        docs = WikipediaLoader(query=query, load_max_docs=2).load()
        return "\n\n---\n\n".join(
            f"<Doc src='{d.metadata.get('source','Wikipedia')}'>{d.page_content[:1000]}</Doc>"
            for d in docs
        )
    except Exception as e:
        return f"Wikipedia search failed: {e}"

# ---- Enhanced Agent State ----
class EnhancedAgentState(TypedDict):
    """State structure for the enhanced agent system."""
    messages: Annotated[List[HumanMessage | AIMessage], operator.add]
    query: str
    agent_type: str
    final_answer: str
    perf: Dict[str, Any]
    agno_resp: str

# ---- Enhanced Multi-LLM System ----
class HybridLangGraphMultiLLMSystem:
    """Enhanced question-answering system with proper response handling."""
    
    def __init__(self):
        """Initialize the enhanced multi-LLM system."""
        self.tools = [
            multiply, add, subtract, divide, modulus,
            optimized_web_search, optimized_wiki_search
        ]
        self.graph = self._build_graph()

    def _llm(self, model_name: str) -> ChatGroq:
        """Create a Groq LLM instance."""
        return ChatGroq(
            model=model_name,
            temperature=0,
            api_key=os.getenv("GROQ_API_KEY")
        )

    def _build_graph(self) -> StateGraph:
        """Build the LangGraph state machine with proper response handling."""
        # Initialize LLMs
        llama8_llm = self._llm("llama3-8b-8192")
        llama70_llm = self._llm("llama3-70b-8192")
        deepseek_llm = self._llm("deepseek-chat")

        def router(st: EnhancedAgentState) -> EnhancedAgentState:
            """Route queries to appropriate LLM based on content analysis."""
            q = st["query"].lower()
            
            # Enhanced routing logic
            if any(keyword in q for keyword in ["calculate", "compute", "math", "multiply", "add", "subtract", "divide"]):
                t = "llama70"  # Use more powerful model for calculations
            elif any(keyword in q for keyword in ["search", "find", "lookup", "wikipedia", "information about"]):
                t = "search_enhanced"  # Use search-enhanced processing
            elif "deepseek" in q or any(keyword in q for keyword in ["analyze", "reasoning", "complex"]):
                t = "deepseek"
            elif "llama-8" in q:
                t = "llama8"
            elif len(q.split()) > 20:  # Complex queries
                t = "llama70"
            else:
                t = "llama8"  # Default for simple queries
                
            return {**st, "agent_type": t}

        def llama8_node(st: EnhancedAgentState) -> EnhancedAgentState:
            """Process query with Llama-3 8B model."""
            t0 = time.time()
            try:
                # Create enhanced prompt with context
                enhanced_query = f"""
                Question: {st["query"]}
                
                Please provide a direct, accurate answer to this question. Do not repeat the question.
                """
                
                sys = SystemMessage(content=ENHANCED_SYSTEM_PROMPT)
                res = llama8_llm.invoke([sys, HumanMessage(content=enhanced_query)])
                
                # Extract and clean the answer
                answer = res.content.strip()
                if "FINAL ANSWER:" in answer:
                    answer = answer.split("FINAL ANSWER:")[-1].strip()
                
                return {**st,
                        "final_answer": answer,
                        "perf": {"time": time.time() - t0, "prov": "Groq-Llama3-8B"}}
            except Exception as e:
                return {**st, "final_answer": f"Error: {e}", "perf": {"error": str(e)}}

        def llama70_node(st: EnhancedAgentState) -> EnhancedAgentState:
            """Process query with Llama-3 70B model."""
            t0 = time.time()
            try:
                # Create enhanced prompt with context
                enhanced_query = f"""
                Question: {st["query"]}
                
                Please provide a direct, accurate answer to this question. Do not repeat the question.
                """
                
                sys = SystemMessage(content=ENHANCED_SYSTEM_PROMPT)
                res = llama70_llm.invoke([sys, HumanMessage(content=enhanced_query)])
                
                # Extract and clean the answer
                answer = res.content.strip()
                if "FINAL ANSWER:" in answer:
                    answer = answer.split("FINAL ANSWER:")[-1].strip()
                
                return {**st,
                        "final_answer": answer,
                        "perf": {"time": time.time() - t0, "prov": "Groq-Llama3-70B"}}
            except Exception as e:
                return {**st, "final_answer": f"Error: {e}", "perf": {"error": str(e)}}

        def deepseek_node(st: EnhancedAgentState) -> EnhancedAgentState:
            """Process query with DeepSeek model."""
            t0 = time.time()
            try:
                # Create enhanced prompt with context
                enhanced_query = f"""
                Question: {st["query"]}
                
                Please provide a direct, accurate answer to this question. Do not repeat the question.
                """
                
                sys = SystemMessage(content=ENHANCED_SYSTEM_PROMPT)
                res = deepseek_llm.invoke([sys, HumanMessage(content=enhanced_query)])
                
                # Extract and clean the answer
                answer = res.content.strip()
                if "FINAL ANSWER:" in answer:
                    answer = answer.split("FINAL ANSWER:")[-1].strip()
                
                return {**st,
                        "final_answer": answer,
                        "perf": {"time": time.time() - t0, "prov": "Groq-DeepSeek"}}
            except Exception as e:
                return {**st, "final_answer": f"Error: {e}", "perf": {"error": str(e)}}

        def search_enhanced_node(st: EnhancedAgentState) -> EnhancedAgentState:
            """Process query with search enhancement."""
            t0 = time.time()
            
            try:
                # Determine search strategy
                query = st["query"]
                search_results = ""
                
                if any(keyword in query.lower() for keyword in ["wikipedia", "wiki"]):
                    search_results = optimized_wiki_search.invoke({"query": query})
                else:
                    search_results = optimized_web_search.invoke({"query": query})
                
                # Create comprehensive prompt with search results
                enhanced_query = f"""
                Original Question: {query}
                
                Search Results:
                {search_results}
                
                Based on the search results above, provide a direct answer to the original question.
                Extract the specific information requested. Do not repeat the question.
                """
                
                sys = SystemMessage(content=ENHANCED_SYSTEM_PROMPT)
                res = llama70_llm.invoke([sys, HumanMessage(content=enhanced_query)])
                
                # Extract and clean the answer
                answer = res.content.strip()
                if "FINAL ANSWER:" in answer:
                    answer = answer.split("FINAL ANSWER:")[-1].strip()
                
                return {**st,
                        "final_answer": answer,
                        "perf": {"time": time.time() - t0, "prov": "Search-Enhanced-Llama70"}}
            except Exception as e:
                return {**st, "final_answer": f"Error: {e}", "perf": {"error": str(e)}}

        # Build graph
        g = StateGraph(EnhancedAgentState)
        g.add_node("router", router)
        g.add_node("llama8", llama8_node)
        g.add_node("llama70", llama70_node)
        g.add_node("deepseek", deepseek_node)
        g.add_node("search_enhanced", search_enhanced_node)
        
        g.set_entry_point("router")
        g.add_conditional_edges("router", lambda s: s["agent_type"], {
            "llama8": "llama8",
            "llama70": "llama70",
            "deepseek": "deepseek",
            "search_enhanced": "search_enhanced"
        })
        
        for node in ["llama8", "llama70", "deepseek", "search_enhanced"]:
            g.add_edge(node, END)
            
        return g.compile(checkpointer=MemorySaver())

    def process_query(self, q: str) -> str:
        """Process a query and return the final answer."""
        state = {
            "messages": [HumanMessage(content=q)],
            "query": q,
            "agent_type": "",
            "final_answer": "",
            "perf": {},
            "agno_resp": ""
        }
        cfg = {"configurable": {"thread_id": f"qa_{hash(q)}"}}
        
        try:
            out = self.graph.invoke(state, cfg)
            answer = out.get("final_answer", "").strip()
            
            # Ensure we don't return the question as the answer
            if answer == q or answer.startswith(q):
                return "Information not available"
            
            return answer if answer else "No answer generated"
        except Exception as e:
            return f"Error processing query: {e}"

def build_graph(provider: str | None = None) -> StateGraph:
    """Build and return the graph for the enhanced agent system."""
    return HybridLangGraphMultiLLMSystem().graph

if __name__ == "__main__":
    # Test the system
    qa_system = HybridLangGraphMultiLLMSystem()
    
    test_questions = [
        "What is 25 multiplied by 17?",
        "Who was the first president of the United States?",
        "Find information about artificial intelligence on Wikipedia"
    ]
    
    for question in test_questions:
        print(f"Question: {question}")
        answer = qa_system.process_query(question)
        print(f"Answer: {answer}")
        print("-" * 50)