nodes.py

import os
import json
from langchain_groq import ChatGroq
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage, ToolMessage
from langchain_core.prompts import ChatPromptTemplate
from state import AgentState

# Load keys if not already loaded
if "GROQ_API_KEY" not in os.environ:
    import config
    config.load_keys()

# Initialize LLM with the requested model
llm = ChatGroq(model_name="openai/gpt-oss-120b", temperature=0)

# Initialize Tool
tavily_tool = TavilySearchResults(max_results=3)

# --- 1. Clarity Agent ---
def clarity_agent(state: AgentState):
    """
    Analyzes if the user's query is clear and specific.
    """
    messages = state["messages"]
    
    system_prompt = (
        "You are a Clarity Agent. Your job is to determine if the user's latest query is clear "
        "and specific enough to conduct research, CONSIDERING THE CONVERSATION CONTEXT.\n"
        "1. If the query mentions a company (e.g. Apple) explicitly, it is 'clear'.\n"
        "2. If the query is implicit (e.g. 'tell me about it', 'what is the price') BUT the company was mentioned in previous messages, it is 'clear'.\n"
        "3. ONLY mark as 'needs_clarification' if the query is vague AND no company context exists in history.\n"
        "Output ONLY a JSON object: {\"status\": \"clear\"} or {\"status\": \"needs_clarification\", \"message\": \"<question to user>\"}."
    )
    
    # Pass full history to allow resolving references
    response = llm.invoke([SystemMessage(content=system_prompt)] + messages)
    
    try:
        content = response.content.strip()
        if "```json" in content:
            content = content.split("```json")[1].split("```")[0]
        elif "{" not in content:
            return {"clarity_status": "clear"}
            
        data = json.loads(content)
        status = data.get("status", "clear")
        msg = data.get("message", None)
        
        updates = {"clarity_status": status}
        if status == "needs_clarification" and msg:
            updates["messages"] = [AIMessage(content=msg)]
            
        return updates
    except Exception as e:
        print(f"Clarity parsing error: {e}")
        return {"clarity_status": "clear"}

# --- 2. Research Agent ---
def research_agent(state: AgentState):
    """
    Generates a search query and uses Tavily to find info.
    """
    messages = state["messages"]
    
    # Improved query generation prompt
    query_gen_prompt = (
        "You are a Research Assistant. Your task is to generate ONE effective search query for Tavily "
        "to find the most relevant and up-to-date information for the user's request. "
        "Consider the entire conversation context if available.\n"
        "Output ONLY the search query string. Do not include quotes or prefixes."
    )
    
    # We pass the full history to let LLM understand context (e.g., "What about *their* CEO?")
    query_msg = llm.invoke([SystemMessage(content=query_gen_prompt)] + messages)
    search_query = query_msg.content.strip()
    
    # Clean up quotes if present
    if search_query.startswith('"') and search_query.endswith('"'):
        search_query = search_query[1:-1]
    
    try:
        print(f"  [Research] Searching for: {search_query}")
        # Explicitly invoking the tool
        results = tavily_tool.invoke(search_query) # returns list of dicts or str
        
        # Serialize results slightly for better consumption
        if isinstance(results, list):
            findings_str = json.dumps(results, indent=2)
            # Heuristic: If we got results, high confidence
            confidence = 8.0 if len(results) > 0 else 2.0
        else:
            findings_str = str(results)
            confidence = 5.0 # Unsure format
            
    except Exception as e:
        findings_str = f"Search failed: {e}"
        confidence = 0.0
        
    current_attempts = state.get("attempt_count", 0)
    
    return {
        "research_findings": findings_str,
        "confidence_score": confidence,
        "attempt_count": current_attempts + 1,
        # Log the action in messages
        "messages": [AIMessage(content=f"Conducted research on: {search_query}")]
    }

# --- 3. Validator Agent ---
def validator_agent(state: AgentState):
    """
    Reviews research quality.
    """
    prompt = (
        "You are a Validator Agent. Review the following research findings and determine if they are sufficient "
        "to answer the user's latest question.\n\n"
        f"Findings: {state.get('research_findings', '')}\n\n"
        "Output ONLY a JSON: {\"result\": \"sufficient\"} or {\"result\": \"insufficient\"}."
    )
    
    # Only need the prompt, not full history, to validate findings content?
    # Better to give history so it knows WHAT was asked.
    response = llm.invoke([SystemMessage(content=prompt)] + state["messages"])
    
    try:
        content = response.content.strip()
        if "```json" in content:
            content = content.split("```json")[1].split("```")[0]
        data = json.loads(content)
        result = data.get("result", "sufficient")
    except:
        result = "sufficient"
        
    return {"validation_result": result}

# --- 4. Synthesis Agent ---
def synthesis_agent(state: AgentState):
    """
    Creates a coherent summary.
    """
    prompt = (
        "You are a Synthesis Agent. Answer the user's question using the provided research findings.\n"
        "Maintain a helpful and professional tone. Cite sources if available in the findings.\n\n"
        f"Research Findings: {state.get('research_findings', '')}"
    )
    
    response = llm.invoke([SystemMessage(content=prompt)] + state["messages"])
    return {"messages": [response]}