src/sales_assistant/agent_main/agent_runner.py

import os
import uuid
from typing import Dict, Any, Optional, List
from langchain_core.messages import HumanMessage, AIMessage
from langgraph.graph import MessagesState
from langgraph.checkpoint.memory import MemorySaver
from langchain_core.tracers.langchain import LangChainTracer
from langchain_core.callbacks import CallbackManager
from dotenv import load_dotenv
from pydantic import BaseModel, Field, validator
from .agent_graph import create_agent_graph, AgentConfig

load_dotenv()

class ConversationConfig(BaseModel):
    """Enhanced configuration with validation for GPT-5-mini conversations."""
    session_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
    user_id: str = Field(default="default_user")
    langsmith_project: str = Field(default_factory=lambda: os.getenv("LANGSMITH_PROJECT", "sales-assistant"))
    enable_langsmith: bool = Field(default=True)
    model_name: str = Field(default_factory=lambda: os.getenv("MODEL_NAME", "gpt-5-mini"))
    model_provider: str = Field(default_factory=lambda: os.getenv("MODEL_PROVIDER", "openai"))
    max_turns: int = Field(default=50, description="Maximum conversation turns")
    recursion_limit: int = Field(default_factory=lambda: int(os.getenv("LANGRAPH_RECURSION_LIMIT", "100")), description="Maximum recursion limit for LangGraph")
    
    @validator('model_name')
    def validate_model(cls, v):
        if v != "gpt-5-mini":
            raise ValueError("Only gpt-5-mini is supported")
        return v
    
    @validator('model_provider')
    def validate_provider(cls, v):
        if v != "openai":
            raise ValueError("Only openai provider is supported")
        return v

def setup_langsmith_tracing(config: ConversationConfig) -> Optional[CallbackManager]:
    """
    Setup LangSmith tracing using built-in LangChain tracers.
    
    Args:
        config: Configuration for the conversation
        
    Returns:
        CallbackManager with LangChain tracer if enabled
    """
    if not config.enable_langsmith or not os.getenv("LANGSMITH_API_KEY"):
        return None
    
    try:
        # Use built-in environment variable setup
        os.environ["LANGCHAIN_TRACING_V2"] = "true"
        os.environ["LANGCHAIN_PROJECT"] = config.langsmith_project
        os.environ["LANGCHAIN_ENDPOINT"] = os.getenv("LANGSMITH_ENDPOINT", "https://api.smith.langchain.com")
        os.environ["LANGCHAIN_API_KEY"] = os.getenv("LANGSMITH_API_KEY")
        
        # Use built-in LangChain tracer
        tracer = LangChainTracer(project_name=config.langsmith_project)
        callback_manager = CallbackManager([tracer])
        
        return callback_manager
        
    except Exception as e:
        print(f"Warning: Could not setup LangSmith tracing: {e}")
        return None

def create_agent_runner(config: ConversationConfig = None) -> tuple:
    """
    Create agent runner with built-in LangGraph memory and state management.
    
    Args:
        config: Optional conversation configuration
        
    Returns:
        Tuple of (compiled_graph, checkpointer, callback_manager, thread_id)
    """
    if config is None:
        config = ConversationConfig()
    
    # Setup built-in LangSmith tracing
    callback_manager = setup_langsmith_tracing(config)
    
    # Use built-in MemorySaver with validation
    checkpointer = MemorySaver()
    
    # Create agent config
    agent_config = AgentConfig(
        model_name=config.model_name,
        enable_checkpointing=True,
        recursion_limit=config.recursion_limit
    )
    
    # Create the agent graph with built-in components
    compiled_graph = create_agent_graph(checkpointer, agent_config)
    
    thread_id = config.session_id
    
    return compiled_graph, checkpointer, callback_manager, thread_id

def run_conversation_turn(
    compiled_graph, 
    thread_id: str, 
    user_input: str,
    callback_manager: Optional[CallbackManager] = None,
    recursion_limit: int = 100
) -> str:
    """
    Run conversation turn with built-in state management and error handling.
    
    Args:
        compiled_graph: The compiled LangGraph with built-in state
        thread_id: Thread ID for conversation persistence  
        user_input: User's input message
        callback_manager: Optional callback manager for tracing
        recursion_limit: Maximum number of recursion steps (default: 100)
        
    Returns:
        Agent's response as a string
    """
    try:
        # Create user message with built-in message types
        user_message = HumanMessage(content=user_input)
        
        # Use built-in thread configuration with recursion limit
        config = {
            "configurable": {"thread_id": thread_id},
            "callbacks": callback_manager.handlers if callback_manager else None,
            "recursion_limit": recursion_limit
        }
        
        # Invoke with built-in state management and error handling
        result = compiled_graph.invoke(
            {"messages": [user_message]}, 
            config=config
        )
        
        # Extract response using built-in message handling
        if result and "messages" in result and result["messages"]:
            assistant_message = result["messages"][-1]
            return assistant_message.content if hasattr(assistant_message, 'content') else str(assistant_message)
        
        return "I apologize, but I couldn't generate a response. Please try again."
        
    except Exception as e:
        error_message = f"Error processing message: {str(e)}"
        print(error_message)
        return error_message

def run_interactive_loop(config: ConversationConfig = None) -> None:
    """
    Run interactive loop with built-in state persistence and validation.
    
    Args:
        config: Optional conversation configuration
    """
    if config is None:
        config = ConversationConfig()
    
    # Create agent runner with built-in components
    compiled_graph, checkpointer, callback_manager, thread_id = create_agent_runner(config)
    
    print("\n🤖 Sales Assistant initialized!")
    print(f"Using model: {config.model_name} (Provider: {config.model_provider})")
    print("💬 Type your questions about products or 'quit' to exit")
    print(f"📊 Session ID: {config.session_id}")
    print(f"🧵 Thread ID: {thread_id}")
    if callback_manager:
        print(f"📈 LangSmith tracking enabled - Project: {config.langsmith_project}")
    print("-" * 50)
    
    turn_count = 0
    while turn_count < config.max_turns:
        try:
            user_input = input("\n👤 You: ").strip()
            
            if user_input.lower() in ['quit', 'exit', 'bye']:
                print("👋 Goodbye!")
                break
            
            if not user_input:
                continue
            
            print("🔄 Processing...")
            response = run_conversation_turn(
                compiled_graph, 
                thread_id, 
                user_input, 
                callback_manager,
                recursion_limit=config.recursion_limit
            )
            
            print(f"\n🤖 Assistant: {response}")
            turn_count += 1
            
        except KeyboardInterrupt:
            print("\n👋 Goodbye!")
            break
        except Exception as e:
            print(f"\n❌ Error: {str(e)}")

def get_conversation_history(compiled_graph, thread_id: str) -> List[Dict[str, Any]]:
    """
    Get conversation history using built-in state management.
    
    Args:
        compiled_graph: The compiled graph with checkpointer
        thread_id: Thread ID for conversation
        
    Returns:
        List of messages with built-in validation
    """
    try:
        config = {"configurable": {"thread_id": thread_id}}
        # Use built-in state retrieval
        state = compiled_graph.get_state(config)
        messages = state.values.get("messages", [])
        
        # Convert to structured format with built-in message handling
        history = []
        for msg in messages:
            if hasattr(msg, "type") and hasattr(msg, "content"):
                history.append({
                    "type": msg.type,
                    "content": msg.content,
                    "id": getattr(msg, "id", str(uuid.uuid4()))
                })
        return history
    except Exception as e:
        print(f"Error getting conversation history: {e}")
        return []

def clear_conversation_history(compiled_graph, thread_id: str) -> None:
    """
    Clear the conversation history in LangGraph's checkpointer.
    
    Args:
        compiled_graph: The compiled graph with checkpointer
        thread_id: Thread ID for conversation
    """
    try:
        config = {"configurable": {"thread_id": thread_id}}
        # Note: LangGraph's MemorySaver doesn't have a direct clear method
        print(f"Note: To clear history, restart with a new thread_id")
    except Exception as e:
        print(f"Error clearing conversation history: {e}")


# Example usage function
def demo_agent_runner():
    """
    Demo function to show how to use the agent runner with LangGraph's built-in state management.
    """
    # Create configuration
    config = ConversationConfig(
        session_id="demo_session",
        user_id="demo_user",
        langsmith_project="sales-assistant-demo"
    )
    
    # Create agent runner with LangGraph's built-in memory
    compiled_graph, checkpointer, langsmith_client, thread_id = create_agent_runner(config)
    
    # Run a few example interactions
    test_queries = [
        "What products do you have from Samsung?",
        "Show me the cheapest camara available from angekis", 
        "Can you find me Samsung outdoor tv-s?"
    ]
    
    print("🚀 Running demo conversations...")
    
    for query in test_queries:
        print(f"\n👤 Demo Query: {query}")
        response = run_conversation_turn(
            compiled_graph, 
            thread_id, 
            query, 
            langsmith_client,
            recursion_limit=config.recursion_limit
        )
        print(f"🤖 Response: {response}")
    
    # Show conversation history
    print("\n📝 Conversation History:")
    history = get_conversation_history(compiled_graph, thread_id)
    for i, msg in enumerate(history, 1):
        print(f"{i}. [{msg['type']}]: {msg['content'][:100]}...")


if __name__ == "__main__":
    # Run with configuration from environment variables
    config = ConversationConfig()
    run_interactive_loop(config)