# Install required packages
pip install langgraph langchain openai python-dotenv

# Create environment file
echo "OPENAI_API_KEY=your_openai_api_key_here" > .env

import os
from typing import TypedDict, Annotated
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolExecutor
from langchain_core.tools import tool
import operator
from dotenv import load_dotenv

load_dotenv()

# Define the state structure
class AgentState(TypedDict):
    messages: Annotated[list, operator.add]
    user_info: dict
    escalation_needed: bool
    conversation_summary: str

# Initialize the LLM
llm = ChatOpenAI(model="gpt-4", temperature=0.1)

# Define tools for the agent
@tool
def search_knowledge_base(query: str) -> str:
    """Search the company knowledge base for information."""
    # Simulate knowledge base search
    knowledge_base = {
        "refund": "Refunds are processed within 5-7 business days. Please provide your order number.",
        "shipping": "Standard shipping takes 3-5 business days. Express shipping is 1-2 business days.",
        "account": "To reset your password, click 'Forgot Password' on the login page.",
        "product": "Our products come with a 1-year warranty. Contact support for warranty claims."
    }

    for key, value in knowledge_base.items():
        if key.lower() in query.lower():
            return value
    return "I couldn't find specific information about that. Let me escalate this to a human agent."

@tool
def get_order_status(order_number: str) -> str:
    """Get the status of a customer order."""
    # Simulate order lookup
    orders = {
        "ORD123": "Shipped - Expected delivery: Tomorrow",
        "ORD456": "Processing - Will ship within 24 hours",
        "ORD789": "Delivered - Delivered yesterday at 2:30 PM"
    }
    return orders.get(order_number, "Order not found. Please check the order number.")

# Create tool executor
tools = [search_knowledge_base, get_order_status]
tool_executor = ToolExecutor(tools)

def should_escalate(state: AgentState) -> str:
    """Determine if the conversation should be escalated to a human."""
    last_message = state["messages"][-1].content if state["messages"] else ""

    escalation_triggers = [
        "speak to manager", "human agent", "not satisfied",
        "complaint", "angry", "frustrated", "cancel subscription"
    ]

    if any(trigger in last_message.lower() for trigger in escalation_triggers):
        return "escalate"

    if state.get("escalation_needed", False):
        return "escalate"

    return "continue"

def agent_node(state: AgentState):
    """Main agent processing node."""
    messages = state["messages"]

    # System prompt for the customer support agent
    system_prompt = """You are a helpful customer support agent. Your goal is to:
    1. Provide accurate information from the knowledge base
    2. Use tools to look up order information when needed
    3. Escalate to human agents when appropriate
    4. Maintain a friendly and professional tone

    If you cannot resolve an issue or the customer requests human assistance,
    set escalation_needed to True."""

    # Prepare messages for the LLM
    llm_messages = [SystemMessage(content=system_prompt)] + messages

    # Get response from LLM
    response = llm.invoke(llm_messages)

    # Check if escalation is needed based on the response
    escalation_needed = "escalate" in response.content.lower() or state.get("escalation_needed", False)

    return {
        "messages": [response],
        "escalation_needed": escalation_needed,
        "conversation_summary": f"Agent responded to: {messages[-1].content[:100]}..."
    }

def escalation_node(state: AgentState):
    """Handle escalation to human agents."""
    escalation_message = AIMessage(
        content="I'm connecting you with a human agent who can better assist you. "
                "Please hold for a moment while I transfer your conversation."
    )

    return {
        "messages": [escalation_message],
        "escalation_needed": True,
        "conversation_summary": "Conversation escalated to human agent"
    }

def tool_node(state: AgentState):
    """Execute tools when needed."""
    last_message = state["messages"][-1]

    # Simple tool routing based on message content
    if "order" in last_message.content.lower() and any(char.isdigit() for char in last_message.content):
        # Extract order number (simplified)
        order_num = "".join([word for word in last_message.content.split() if word.startswith("ORD")])
        if order_num:
            result = get_order_status.invoke({"order_number": order_num})
        else:
            result = "Please provide your order number (format: ORD123)"
    else:
        result = search_knowledge_base.invoke({"query": last_message.content})

    tool_response = AIMessage(content=result)

    return {
        "messages": [tool_response],
        "conversation_summary": f"Tool executed for: {last_message.content[:50]}..."
    }

# Build the graph
workflow = StateGraph(AgentState)

# Add nodes
workflow.add_node("agent", agent_node)
workflow.add_node("tools", tool_node)
workflow.add_node("escalate", escalation_node)

# Set entry point
workflow.set_entry_point("agent")

# Add conditional edges
workflow.add_conditional_edges(
    "agent",
    should_escalate,
    {
        "continue": "tools",
        "escalate": "escalate"
    }
)

# Add edges
workflow.add_edge("tools", "agent")
workflow.add_edge("escalate", END)

# Compile the graph
app = workflow.compile()

# Example usage
def run_customer_support():
    """Run the customer support agent."""
    print("🤖 Customer Support Agent Started")
    print("Type 'quit' to exit\n")

    # Initialize state
    state = {
        "messages": [],
        "user_info": {},
        "escalation_needed": False,
        "conversation_summary": ""
    }

    while True:
        user_input = input("You: ")
        if user_input.lower() == 'quit':
            break

        # Add user message to state
        user_message = HumanMessage(content=user_input)
        state["messages"].append(user_message)

        # Run the graph
        result = app.invoke(state)

        # Get the latest AI response
        ai_response = result["messages"][-1].content
        print(f"Agent: {ai_response}\n")

        # Update state for next iteration
        state = result

        # Check if escalated
        if result.get("escalation_needed"):
            print("🔄 Conversation has been escalated to a human agent.")
            break

if __name__ == "__main__":
    run_customer_support()

from typing import List
import asyncio

class ResearchState(TypedDict):
    query: str
    research_results: List[str]
    analysis: str
    final_report: str
    current_agent: str

# Specialized research agents
def researcher_agent(state: ResearchState):
    """Agent that conducts initial research."""
    query = state["query"]

    # Simulate research (in real implementation, this would call APIs)
    research_prompt = f"""
    Research the topic: {query}
    Provide 3-5 key findings with sources.
    Focus on recent developments and factual information.
    """

    response = llm.invoke([SystemMessage(content=research_prompt)])

    return {
        "research_results": [response.content],
        "current_agent": "researcher"
    }

def analyst_agent(state: ResearchState):
    """Agent that analyzes research findings."""
    research_data = "\n".join(state["research_results"])

    analysis_prompt = f"""
    Analyze the following research findings:
    {research_data}

    Provide:
    1. Key insights and patterns
    2. Potential implications
    3. Areas needing further investigation
    """

    response = llm.invoke([SystemMessage(content=analysis_prompt)])

    return {
        "analysis": response.content,
        "current_agent": "analyst"
    }

def report_writer_agent(state: ResearchState):
    """Agent that creates the final report."""
    research = "\n".join(state["research_results"])
    analysis = state["analysis"]

    report_prompt = f"""
    Create a comprehensive report based on:

    Research Findings:
    {research}

    Analysis:
    {analysis}

    Structure the report with:
    1. Executive Summary
    2. Key Findings
    3. Analysis and Insights
    4. Recommendations
    5. Conclusion
    """

    response = llm.invoke([SystemMessage(content=report_prompt)])

    return {
        "final_report": response.content,
        "current_agent": "report_writer"
    }

# Build multi-agent workflow
def create_research_workflow():
    workflow = StateGraph(ResearchState)

    # Add agents as nodes
    workflow.add_node("researcher", researcher_agent)
    workflow.add_node("analyst", analyst_agent)
    workflow.add_node("report_writer", report_writer_agent)

    # Define the flow
    workflow.set_entry_point("researcher")
    workflow.add_edge("researcher", "analyst")
    workflow.add_edge("analyst", "report_writer")
    workflow.add_edge("report_writer", END)

    return workflow.compile()

# Usage example
def run_research_system():
    app = create_research_workflow()

    initial_state = {
        "query": "Impact of AI on software development in 2024",
        "research_results": [],
        "analysis": "",
        "final_report": "",
        "current_agent": ""
    }

    result = app.invoke(initial_state)

    print("🔍 Research Complete!")
    print(f"📊 Final Report:\n{result['final_report']}")

if __name__ == "__main__":
    run_research_system()

# Get conversation history
def inspect_conversation_state(app, state):
    """Inspect the current state of the conversation."""
    print("📋 Current State:")
    print(f"Messages: {len(state.get('messages', []))}")
    print(f"Escalation needed: {state.get('escalation_needed', False)}")
    print(f"Summary: {state.get('conversation_summary', 'None')}")

# Rewind to a previous state
def rewind_conversation(app, state, steps_back=1):
    """Rewind the conversation to a previous state."""
    if len(state["messages"]) >= steps_back * 2:  # User + AI message pairs
        state["messages"] = state["messages"][:-steps_back * 2]
        print(f"⏪ Rewound {steps_back} steps")
    return state

def human_approval_node(state: AgentState):
    """Node that requires human approval before proceeding."""
    last_message = state["messages"][-1].content

    print(f"🤔 Agent wants to: {last_message}")
    approval = input("Approve this action? (y/n): ")

    if approval.lower() == 'y':
        return {"approved": True}
    else:
        return {
            "approved": False,
            "messages": [AIMessage(content="Action was not approved. How would you like me to proceed?")]
        }

import json
from datetime import datetime

class ConversationMemory:
    def __init__(self, file_path="conversation_memory.json"):
        self.file_path = file_path
        self.memory = self.load_memory()

    def load_memory(self):
        try:
            with open(self.file_path, 'r') as f:
                return json.load(f)
        except FileNotFoundError:
            return {"conversations": [], "user_preferences": {}}

    def save_memory(self):
        with open(self.file_path, 'w') as f:
            json.dump(self.memory, f, indent=2)

    def add_conversation(self, state):
        conversation = {
            "timestamp": datetime.now().isoformat(),
            "messages": [msg.content for msg in state["messages"]],
            "summary": state.get("conversation_summary", ""),
            "escalated": state.get("escalation_needed", False)
        }
        self.memory["conversations"].append(conversation)
        self.save_memory()

    def get_user_context(self, user_id):
        # Return relevant context for the user
        recent_conversations = self.memory["conversations"][-5:]  # Last 5 conversations
        return {
            "recent_topics": [conv["summary"] for conv in recent_conversations],
            "preferences": self.memory["user_preferences"].get(user_id, {})
        }

# Usage in agent
memory = ConversationMemory()

def memory_aware_agent(state: AgentState):
    # Get user context
    context = memory.get_user_context("current_user")

    # Include context in system prompt
    system_prompt = f"""
    You are a customer support agent.

    User Context:
    - Recent topics: {context['recent_topics']}
    - Preferences: {context['preferences']}

    Use this context to provide personalized assistance.
    """

    # Process as before...
    response = llm.invoke([SystemMessage(content=system_prompt)] + state["messages"])

    # Save conversation
    memory.add_conversation(state)

    return {"messages": [response]}

def robust_agent_node(state: AgentState):
    """Agent node with comprehensive error handling."""
    try:
        # Main agent logic
        response = llm.invoke(state["messages"])
        return {"messages": [response]}

    except Exception as e:
        # Log error and provide fallback response
        print(f"❌ Error in agent: {e}")

        fallback_response = AIMessage(
            content="I'm experiencing some technical difficulties. "
                   "Let me connect you with a human agent who can assist you."
        )

        return {
            "messages": [fallback_response],
            "escalation_needed": True
        }

# Use async for better performance
async def async_agent_node(state: AgentState):
    """Asynchronous agent processing."""
    tasks = []

    # Process multiple operations concurrently
    if needs_knowledge_search(state):
        tasks.append(search_knowledge_base.ainvoke({"query": get_query(state)}))

    if needs_order_lookup(state):
        tasks.append(get_order_status.ainvoke({"order_number": get_order_number(state)}))

    # Wait for all tasks to complete
    results = await asyncio.gather(*tasks)

    # Combine results and generate response
    combined_info = "\n".join(results)
    response = await llm.ainvoke([
        SystemMessage(content=f"Based on this information: {combined_info}"),
        *state["messages"]
    ])

    return {"messages": [response]}

import unittest

class TestCustomerSupportAgent(unittest.TestCase):
    def setUp(self):
        self.app = workflow.compile()

    def test_basic_inquiry(self):
        """Test basic customer inquiry handling."""
        state = {
            "messages": [HumanMessage(content="What's your refund policy?")],
            "user_info": {},
            "escalation_needed": False,
            "conversation_summary": ""
        }

        result = self.app.invoke(state)

        # Assert that a response was generated
        self.assertTrue(len(result["messages"]) > 1)
        self.assertIn("refund", result["messages"][-1].content.lower())

    def test_escalation_trigger(self):
        """Test that escalation triggers work correctly."""
        state = {
            "messages": [HumanMessage(content="I want to speak to a manager!")],
            "user_info": {},
            "escalation_needed": False,
            "conversation_summary": ""
        }

        result = self.app.invoke(state)

        # Assert that escalation was triggered
        self.assertTrue(result["escalation_needed"])

if __name__ == "__main__":
    unittest.main()