llm/agent.py

import logging
import os
from datetime import datetime
from typing import List, Optional

from dotenv import load_dotenv
from langchain_google_genai import ChatGoogleGenerativeAI
from langgraph.prebuilt import create_react_agent

from llm.connection_manager import get_checkpointer
from llm.prompt import system_message
from llm.tools import initialize_tools
from llm.utils import cleanup_old_tool_results, get_tool_result

load_dotenv()

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Global agent instance
_agent_executor = None
# Counter for periodic cleanup
_request_count = 0


def _get_agent():
    """Get or create the agent instance."""
    global _agent_executor

    if _agent_executor is None:
        # Build LLM
        print("[AGENT] building LLM")
        llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash")

        # Build tools
        print("[AGENT] initializing tools")
        tools = initialize_tools()

        # Create agent with fresh checkpointer
        print("[AGENT] creating agent")
        _agent_executor = create_react_agent(
            llm,
            tools=tools,
            prompt=system_message,
            checkpointer=get_checkpointer(),
        )

    return _agent_executor


def _build_message_with_context(message: str, selected_images: Optional[List[dict]], user_id: str) -> str:
    """Build the full message with image context if provided."""
    if not selected_images or len(selected_images) == 0:
        return message

    image_context = "\n\nSelected Images:\n"
    for i, img in enumerate(selected_images, 1):
        image_context += f"{i}. {img.get('title', 'Untitled')} (ID: {img.get('id', 'unknown')})\n"
        image_context += f"   Type: {img.get('type', 'unknown')}\n"
        image_context += f"   Description: {img.get('description', 'No description')}\n"
        if img.get("url"):
            image_context += f"   URL: {img.get('url')}\n"
        image_context += "\n"

    return message + image_context + f"\n\nUser ID: {user_id}"


def _extract_agent_response(response) -> str:
    """Extract the agent's response text from the response object."""
    if not response or "messages" not in response or len(response["messages"]) == 0:
        return "I'm sorry, I couldn't process your request. Please try again."

    last_message = response["messages"][-1]

    # Handle None or unexpected message types
    if last_message is None:
        return "I'm sorry, I couldn't process your request. Please try again."

    # Handle both AIMessage objects and dictionaries
    if hasattr(last_message, "content"):
        content = last_message.content
        if content is None:
            return "I'm sorry, I couldn't process your request. Please try again."
        return content
    elif isinstance(last_message, dict) and "content" in last_message:
        content = last_message["content"]
        if content is None:
            return "I'm sorry, I couldn't process your request. Please try again."
        return content

    return "I'm sorry, I couldn't process your request. Please try again."


def _generate_presigned_url(user_id: str, image_id: str) -> Optional[str]:
    """Generate a presigned URL for an image."""
    import boto3

    s3_client = boto3.client(
        "s3",
        region_name=os.environ.get("AWS_REGION", "us-east-1"),
        aws_access_key_id=os.environ.get("AWS_ACCESS_KEY_ID"),
        aws_secret_access_key=os.environ.get("AWS_SECRET_ACCESS_KEY"),
    )

    bucket_name = os.environ.get("AWS_S3_BUCKET_NAME")
    if not bucket_name:
        print("[AGENT] AWS_S3_BUCKET_NAME not set")
        return None

    try:
        s3_key = f"users/{user_id}/images/{image_id}"
        print(f"[AGENT] Generating presigned URL for S3 key: {s3_key}")

        presigned_url = s3_client.generate_presigned_url(
            "get_object",
            Params={"Bucket": bucket_name, "Key": s3_key},
            ExpiresIn=7200,  # 2 hours
        )
        print(f"[AGENT] Generated presigned URL: {presigned_url[:50]}...")
        return presigned_url

    except Exception as e:
        print(f"[AGENT] Error generating presigned URL: {e}")
        return None


def _process_generated_image(user_id: str, tool_result: dict) -> Optional[dict]:
    """Process a generated image tool result and return image data."""
    image_id = tool_result.get("image_id")
    title = tool_result.get("title", "Generated Image")
    prompt = tool_result.get("prompt", "Based on your request")

    if not image_id:
        print("[AGENT] No image_id found in tool result")
        return None

    print(f"[AGENT] Processing generated image with ID: {image_id}")

    # Generate presigned URL
    presigned_url = _generate_presigned_url(user_id, image_id)
    if not presigned_url:
        return None

    # Create image data structure using data from tool result
    generated_image_data = {
        "id": image_id,
        "url": presigned_url,
        "title": title,
        "description": f"AI-generated image: {prompt}",
        "timestamp": datetime.now().isoformat(),
        "type": "generated",
    }

    print(f"[AGENT] Created generated_image_data: {generated_image_data}")
    return generated_image_data


def _process_tool_results(user_id: str) -> Optional[dict]:
    """Process any tool results for the user and return generated image data if found."""
    print(f"[AGENT] Checking for tool results for user {user_id}")
    tool_result = get_tool_result(user_id, "generate_image")

    if tool_result:
        print(f"[AGENT] Found tool result: {tool_result}")
        return _process_generated_image(user_id, tool_result)
    else:
        print(f"[AGENT] No tool result found for user {user_id}")

    return None


def chat_with_agent(
    message: str,
    client_ip: str,
    user_id: str = "default",
    selected_images: Optional[List[dict]] = None,
) -> tuple[str, Optional[dict]]:
    """
    Send a message to the agent and get a response.

    Args:
        message: The user's message
        user_id: Unique identifier for the user/thread
        selected_images: List of selected image objects (optional)
        client_ip: IP address of the client
    Returns:
        Tuple of (agent_response, generated_image_data)
    """
    global _request_count

    # Periodic cleanup every 10 requests
    _request_count += 1
    if _request_count % 10 == 0:
        print(f"[AGENT] Running periodic cleanup (request #{_request_count})")
        cleanup_old_tool_results()

    print(f"[AGENT] Starting chat_with_agent - user_id: {user_id}, message: {message[:100]}...")
    agent = _get_agent()

    # Prepare the message with context
    print("[AGENT] building message with context")
    full_message = _build_message_with_context(message, selected_images, user_id)

    # Configure thread ID for conversation continuity
    config = {"configurable": {"thread_id": user_id, "client_ip": client_ip}}

    # Get response from agent
    print(f"[AGENT] Invoking agent with config: {config}")
    response = agent.invoke({"messages": [{"role": "user", "content": full_message}]}, config=config)
    print(f"[AGENT] Agent response received: {type(response)}")

    # Extract the agent's response
    agent_response = _extract_agent_response(response)
    print(f"[AGENT] Extracted agent response: {agent_response[:100]}...")

    # Check for tool results and process generated images
    generated_image_data = _process_tool_results(user_id)

    print(f"[AGENT] Returning response - agent_response length: {len(agent_response)}, generated_image_data: {generated_image_data is not None}")
    return agent_response, generated_image_data


if __name__ == "__main__":
    # Test the agent
    response = chat_with_agent("Hello! How can you help me with image editing?", "127.0.0.1")
    print(response)