langgraph_agent.py

import json
import os
from enum import Enum
from typing import Any, Dict, List, Optional, TypedDict

import pandas as pd
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import END, START, StateGraph
from langgraph.prebuilt import ToolNode
from pydantic import BaseModel, Field


# Enums for query types
class QueryType(str, Enum):
    STRUCTURED = "structured"
    UNSTRUCTURED = "unstructured"
    OUT_OF_SCOPE = "out_of_scope"
    RECOMMEND_QUERY = "recommend_query"


class AnalysisType(str, Enum):
    QUANTITATIVE = "quantitative"
    QUALITATIVE = "qualitative"
    OUT_OF_SCOPE = "out_of_scope"


# State definition
class AgentState(TypedDict):
    messages: List[Any]
    query_type: Optional[str]
    analysis_result: Optional[Dict[str, Any]]
    user_profile: Optional[Dict[str, Any]]
    session_context: Optional[Dict[str, Any]]
    recommendations: Optional[List[str]]


# User profile model
class UserProfile(BaseModel):
    interests: List[str] = Field(default_factory=list)
    query_history: List[str] = Field(default_factory=list)
    preferences: Dict[str, Any] = Field(default_factory=dict)
    expertise_level: str = "beginner"


# Dataset management
class DatasetManager:
    _instance = None
    _df = None

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(DatasetManager, cls).__new__(cls)
        return cls._instance

    def get_dataset(self) -> pd.DataFrame:
        if self._df is None:
            from datasets import load_dataset

            dataset = load_dataset(
                "bitext/Bitext-customer-support-llm-chatbot-training-dataset"
            )
            self._df = pd.DataFrame(dataset["train"])
        return self._df


# Tools for structured queries (quantitative analysis)
@tool
def get_category_distribution() -> Dict[str, int]:
    """Get the distribution of categories in the dataset."""
    df = DatasetManager().get_dataset()
    return df["category"].value_counts().to_dict()


@tool
def get_intent_distribution() -> Dict[str, int]:
    """Get the distribution of intents in the dataset."""
    df = DatasetManager().get_dataset()
    return df["intent"].value_counts().to_dict()


@tool
def get_dataset_stats() -> Dict[str, Any]:
    """Get basic statistics about the dataset."""
    df = DatasetManager().get_dataset()
    return {
        "total_records": len(df),
        "unique_categories": len(df["category"].unique()),
        "unique_intents": len(df["intent"].unique()),
        "columns": df.columns.tolist(),
    }


@tool
def get_examples_by_category(category: str, n: int = 5) -> List[Dict[str, Any]]:
    """Get examples from a specific category."""
    df = DatasetManager().get_dataset()
    filtered_df = df[df["category"].str.lower() == category.lower()]
    if filtered_df.empty:
        return []
    return filtered_df.head(n).to_dict("records")


@tool
def get_examples_by_intent(intent: str, n: int = 5) -> List[Dict[str, Any]]:
    """Get examples from a specific intent."""
    df = DatasetManager().get_dataset()
    filtered_df = df[df["intent"].str.lower() == intent.lower()]
    if filtered_df.empty:
        return []
    return filtered_df.head(n).to_dict("records")


@tool
def search_conversations(query: str, n: int = 5) -> List[Dict[str, Any]]:
    """Search for conversations containing specific keywords."""
    df = DatasetManager().get_dataset()
    mask = df["customer"].str.contains(query, case=False, na=False) | df[
        "agent"
    ].str.contains(query, case=False, na=False)
    filtered_df = df[mask]
    return filtered_df.head(n).to_dict("records")


# Tools for unstructured queries (qualitative analysis)
@tool
def get_category_summary(category: str) -> Dict[str, Any]:
    """Get a summary of conversations in a specific category."""
    df = DatasetManager().get_dataset()
    filtered_df = df[df["category"].str.lower() == category.lower()]
    if filtered_df.empty:
        return {"error": f"No data found for category: {category}"}

    return {
        "category": category,
        "count": len(filtered_df),
        "unique_intents": filtered_df["intent"].nunique(),
        "intents": filtered_df["intent"].value_counts().to_dict(),
        "sample_conversations": filtered_df.head(3).to_dict("records"),
    }


@tool
def get_intent_summary(intent: str) -> Dict[str, Any]:
    """Get a summary of conversations for a specific intent."""
    df = DatasetManager().get_dataset()
    filtered_df = df[df["intent"].str.lower() == intent.lower()]
    if filtered_df.empty:
        return {"error": f"No data found for intent: {intent}"}

    return {
        "intent": intent,
        "count": len(filtered_df),
        "categories": filtered_df["category"].value_counts().to_dict(),
        "sample_conversations": filtered_df.head(3).to_dict("records"),
    }


# Memory tools
@tool
def update_user_profile(
    interests: List[str], preferences: Dict[str, Any], expertise_level: str = "beginner"
) -> Dict[str, Any]:
    """Update the user's profile with new information."""
    return {
        "interests": interests,
        "preferences": preferences,
        "expertise_level": expertise_level,
        "updated": True,
    }


# Define tool lists for different agents
structured_tools = [
    get_category_distribution,
    get_intent_distribution,
    get_dataset_stats,
    get_examples_by_category,
    get_examples_by_intent,
    search_conversations,
]

unstructured_tools = [
    get_category_summary,
    get_intent_summary,
    search_conversations,
    get_examples_by_category,
    get_examples_by_intent,
]

memory_tools = [update_user_profile]


class DataAnalystAgent:
    def __init__(self, api_key: str, model_name: str = None):
        # Determine if using Nebius or OpenAI based on API key source
        is_nebius = os.environ.get("NEBIUS_API_KEY") == api_key

        if is_nebius:
            # Configure for Nebius API
            self.llm = ChatOpenAI(
                api_key=api_key,
                model=model_name or "Qwen/Qwen3-30B-A3B",
                base_url="https://api.studio.nebius.com/v1",
                temperature=0,
            )
        else:
            # Configure for OpenAI API
            self.llm = ChatOpenAI(
                api_key=api_key, model=model_name or "gpt-4o", temperature=0
            )

        self.memory = MemorySaver()
        self.graph = self._build_graph()

    def _build_graph(self) -> StateGraph:
        """Build the LangGraph workflow."""
        builder = StateGraph(AgentState)

        # Add nodes
        builder.add_node("classifier", self._classify_query)
        builder.add_node("structured_agent", self._structured_agent)
        builder.add_node("unstructured_agent", self._unstructured_agent)
        builder.add_node("structured_tools", ToolNode(structured_tools))
        builder.add_node("unstructured_tools", ToolNode(unstructured_tools))
        builder.add_node("summarizer", self._update_summary)
        builder.add_node("recommender", self._recommend_queries)
        builder.add_node("out_of_scope", self._handle_out_of_scope)

        # Add edges
        builder.add_edge(START, "classifier")

        # Conditional edges from classifier
        builder.add_conditional_edges(
            "classifier",
            self._route_query,
            {
                "structured": "structured_agent",
                "unstructured": "unstructured_agent",
                "out_of_scope": "out_of_scope",
                "recommend_query": "recommender",
            },
        )

        # Tool routing for structured agent
        builder.add_conditional_edges(
            "structured_agent",
            self._should_use_tools,
            {"tools": "structured_tools", "end": "summarizer"},
        )

        # Tool routing for unstructured agent
        builder.add_conditional_edges(
            "unstructured_agent",
            self._should_use_tools,
            {"tools": "unstructured_tools", "end": "summarizer"},
        )

        # From tools back to respective agents
        builder.add_edge("structured_tools", "structured_agent")
        builder.add_edge("unstructured_tools", "unstructured_agent")

        # End edges
        builder.add_edge("summarizer", END)
        builder.add_edge("out_of_scope", END)
        builder.add_edge("recommender", END)

        return builder.compile(checkpointer=self.memory)

    def _classify_query(self, state: AgentState) -> AgentState:
        """Classify the user query into different types."""
        last_message = state["messages"][-1]
        user_query = last_message.content.lower()

        # Simple keyword-based classification for better reliability
        # Check for recommendation requests first
        if any(
            word in user_query
            for word in [
                "what should i",
                "what to query",
                "recommend",
                "suggest",
                "advise",
                "what next",
                "what can i ask",
            ]
        ):
            query_type = "recommend_query"

        # Check for out of scope queries
        elif any(
            word in user_query
            for word in [
                "weather",
                "news",
                "sports",
                "politics",
                "cooking",
                "travel",
                "music",
                "movies",
                "games",
                "programming",
                "code",
            ]
        ) and not any(
            word in user_query
            for word in ["category", "intent", "customer", "support", "data", "records"]
        ):
            query_type = "out_of_scope"

        # Check for unstructured/qualitative queries
        elif any(
            word in user_query
            for word in [
                "summarize",
                "summary",
                "patterns",
                "insights",
                "analysis",
                "analyze",
                "themes",
                "trends",
                "what patterns",
                "understand",
            ]
        ):
            query_type = "unstructured"

        # Default to structured for data-related queries
        else:
            query_type = "structured"

        # Double-check with LLM for edge cases, but use simpler prompt
        if query_type == "out_of_scope":
            simple_prompt = f"""
            Is this question about customer support data analysis? 
            Question: "{last_message.content}"
            
            Answer only "yes" or "no".
            """

            try:
                response = self.llm.invoke([HumanMessage(content=simple_prompt)])
                if "yes" in response.content.lower():
                    query_type = "structured"  # Override if actually about data
            except Exception:
                pass  # Keep original classification if LLM fails

        state["query_type"] = query_type
        return state

    def _route_query(self, state: AgentState) -> str:
        """Route to appropriate agent based on classification."""
        return state["query_type"]

    def _structured_agent(self, state: AgentState) -> AgentState:
        """Handle structured/quantitative queries."""

        system_prompt = """
        You are a data analyst that MUST use tools to answer questions about 
        customer support data. You have access to these tools:
        
        - get_category_distribution: Get category counts
        - get_intent_distribution: Get intent counts  
        - get_dataset_stats: Get basic dataset statistics
        - get_examples_by_category: Get examples from a category
        - get_examples_by_intent: Get examples from an intent
        - search_conversations: Search for conversations with keywords
        
        IMPORTANT: Always use the appropriate tool to get real data. 
        Do NOT make up or guess answers. Use tools to get actual numbers.
        
        For questions about:
        - "How many categories" or "category distribution" → use get_category_distribution
        - "How many intents" or "intent distribution" → use get_intent_distribution  
        - "Total records" or "dataset size" → use get_dataset_stats
        - "Examples of X" → use get_examples_by_category or get_examples_by_intent
        - "Search for X" → use search_conversations
        """

        llm_with_tools = self.llm.bind_tools(structured_tools)
        messages = [SystemMessage(content=system_prompt)] + state["messages"]
        response = llm_with_tools.invoke(messages)

        state["messages"].append(response)
        return state

    def _unstructured_agent(self, state: AgentState) -> AgentState:
        """Handle unstructured/qualitative queries."""

        system_prompt = """
        You are a data analyst that MUST use tools to provide insights about 
        customer support data. You have access to these tools:
        
        - get_category_summary: Get detailed summary of a category
        - get_intent_summary: Get detailed summary of an intent
        - search_conversations: Search conversations for patterns
        - get_examples_by_category: Get examples to analyze patterns
        - get_examples_by_intent: Get examples to analyze patterns
        
        IMPORTANT: Always use the appropriate tool to get real data.
        Do NOT make up or guess insights. Use tools to get actual data first.
        
        For questions about:
        - "Summarize X category" → use get_category_summary
        - "Analyze X intent" → use get_intent_summary
        - "Patterns in X" → use get_examples_by_category or search_conversations
        """

        llm_with_tools = self.llm.bind_tools(unstructured_tools)
        messages = [SystemMessage(content=system_prompt)] + state["messages"]
        response = llm_with_tools.invoke(messages)

        state["messages"].append(response)
        return state

    def _should_use_tools(self, state: AgentState) -> str:
        """Determine if the agent should use tools or end."""
        last_message = state["messages"][-1]

        # Check if LLM made tool calls
        if hasattr(last_message, "tool_calls") and last_message.tool_calls:
            return "tools"

        # If no tool calls but this is the first response from agent,
        # force tool usage for data questions
        messages = state["messages"]
        human_messages = [msg for msg in messages if isinstance(msg, HumanMessage)]

        if len(human_messages) >= 1:
            last_human_msg = human_messages[-1].content.lower()

            # Check if this looks like a data question that needs tools
            needs_tools = any(
                word in last_human_msg
                for word in [
                    "how many",
                    "show me",
                    "examples",
                    "distribution",
                    "categories",
                    "intents",
                    "records",
                    "statistics",
                    "stats",
                    "count",
                    "total",
                    "billing",
                    "refund",
                    "payment",
                    "technical",
                    "support",
                ]
            )

            # Count AI messages - if this is first AI response and needs tools, force it
            ai_messages = [msg for msg in messages if not isinstance(msg, HumanMessage)]
            if needs_tools and len(ai_messages) <= 1:
                return "tools"

        return "end"

    def _update_summary(self, state: AgentState) -> AgentState:
        """Update user profile/summary based on the interaction."""
        user_profile = state.get("user_profile", {})
        last_human_message = None

        # Find the last human message
        for msg in reversed(state["messages"]):
            if isinstance(msg, HumanMessage):
                last_human_message = msg
                break

        if last_human_message:
            # Extract information about user interests
            system_prompt = """
            Based on the user's question, extract information about their 
            interests and update their profile. Consider:
            - What categories/intents they're interested in
            - Their level of technical detail preference
            - Types of analysis they prefer
            
            Return a JSON with:
            {
                "interests": ["list of topics they seem interested in"],
                "preferences": {"any preferences about analysis style"},
                "expertise_level": "beginner/intermediate/advanced"
            }
            
            If no clear information can be extracted, return empty lists/dicts.
            """

            messages = [
                SystemMessage(content=system_prompt),
                HumanMessage(content=f"User question: {last_human_message.content}"),
            ]

            try:
                response = self.llm.invoke(messages)
                profile_update = json.loads(response.content)

                # Merge with existing profile
                if not user_profile:
                    user_profile = {
                        "interests": [],
                        "preferences": {},
                        "expertise_level": "beginner",
                        "query_history": [],
                    }

                # Update interests (avoid duplicates)
                new_interests = profile_update.get("interests", [])
                existing_interests = user_profile.get("interests", [])
                user_profile["interests"] = list(
                    set(existing_interests + new_interests)
                )

                # Update preferences
                user_profile["preferences"].update(
                    profile_update.get("preferences", {})
                )

                # Update expertise level if provided
                if profile_update.get("expertise_level"):
                    user_profile["expertise_level"] = profile_update["expertise_level"]

                # Add to query history
                if "query_history" not in user_profile:
                    user_profile["query_history"] = []
                user_profile["query_history"].append(last_human_message.content)

                # Keep only last 10 queries
                user_profile["query_history"] = user_profile["query_history"][-10:]

            except (json.JSONDecodeError, Exception):
                # If parsing fails, just add to query history
                if not user_profile:
                    user_profile = {"query_history": []}
                if "query_history" not in user_profile:
                    user_profile["query_history"] = []
                user_profile["query_history"].append(last_human_message.content)
                user_profile["query_history"] = user_profile["query_history"][-10:]

        state["user_profile"] = user_profile
        return state

    def _recommend_queries(self, state: AgentState) -> AgentState:
        """Recommend next queries based on conversation history and user profile."""
        user_profile = state.get("user_profile", {})
        query_history = user_profile.get("query_history", [])
        interests = user_profile.get("interests", [])

        # Get dataset info for context
        df = DatasetManager().get_dataset()
        categories = df["category"].unique().tolist()
        intents = df["intent"].unique()[:20].tolist()

        system_prompt = f"""
        You are a query recommendation assistant. Based on the user's conversation 
        history and interests, suggest relevant follow-up questions they could ask 
        about the customer support dataset.
        
        User's query history: {query_history}
        User's interests: {interests}
        
        Available categories: {categories}
        Sample intents: {intents}
        
        Suggest 3-5 relevant questions the user might want to ask next. Consider:
        - Natural follow-ups to their previous questions
        - Related categories or intents they haven't explored
        - Different types of analysis (if they've only done quantitative, 
          suggest qualitative and vice versa)
        
        Be conversational and explain why each suggestion might be interesting.
        Start with "Based on your previous queries, you might want to..."
        """

        messages = [SystemMessage(content=system_prompt)]

        # Add conversation context
        if state["messages"]:
            messages.extend(state["messages"])
        else:
            messages.append(HumanMessage(content="What should I query next?"))

        response = self.llm.invoke(messages)
        state["messages"].append(response)

        return state

    def _handle_out_of_scope(self, state: AgentState) -> AgentState:
        """Handle queries that are out of scope."""
        response = AIMessage(
            content="I'm sorry, but I can only answer questions about the customer "
            "support dataset. Please ask questions about categories, intents, "
            "conversation examples, or data statistics."
        )
        state["messages"].append(response)
        return state

    def invoke(self, message: str, thread_id: str) -> Dict[str, Any]:
        """Invoke the agent with a message and thread ID."""
        config = {"configurable": {"thread_id": thread_id}}

        # Create input state
        input_state = {"messages": [HumanMessage(content=message)]}

        # Invoke the graph
        result = self.graph.invoke(input_state, config)

        return result

    def get_conversation_history(self, thread_id: str) -> List[Dict[str, Any]]:
        """Get conversation history for a thread."""
        config = {"configurable": {"thread_id": thread_id}}

        try:
            # Get the current state
            state = self.graph.get_state(config)
            if state and state.values.get("messages"):
                return [
                    {
                        "role": (
                            "human" if isinstance(msg, HumanMessage) else "assistant"
                        ),
                        "content": msg.content,
                    }
                    for msg in state.values["messages"]
                ]
        except Exception:
            pass

        return []

    def get_user_profile(self, thread_id: str) -> Dict[str, Any]:
        """Get user profile for a thread."""
        config = {"configurable": {"thread_id": thread_id}}

        try:
            state = self.graph.get_state(config)
            if state and state.values.get("user_profile"):
                return state.values["user_profile"]
        except Exception:
            pass

        return {}