created basic agent

codinglabsong · codinglabsong · commit d471ca7fa003 · 2025-08-19T12:28:29.000-07:00
diff --git a/.gitignore b/.gitignore
@@ -56,3 +56,6 @@ build/
 
 # Trash bin
 trash/
+
+# Database
+db.sqlite3
diff --git a/0.2.27 b/0.2.27
diff --git a/api/.pre-commit-config.yaml b/api/.pre-commit-config.yaml
@@ -1,8 +1,9 @@
 repos:
   - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: v0.5.0
+    rev: v0.6.8
     hooks:
       - id: ruff
+        args: ["--fix"]
       - id: ruff-format
 
   - repo: https://github.com/psf/black
diff --git a/api/llm/agent.py b/api/llm/agent.py
@@ -0,0 +1,85 @@
+import sqlite3
+from typing import List, Optional
+
+from dotenv import load_dotenv
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langgraph.checkpoint.sqlite import SqliteSaver
+from langgraph.prebuilt import create_react_agent
+
+load_dotenv()
+
+# Global agent instance
+_agent_executor = None
+_memory = None
+
+
+def get_agent():
+    """Get or create the agent instance."""
+    global _agent_executor, _memory
+
+    if _agent_executor is None:
+        llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash")
+
+        # Build persistent checkpointer
+        con = sqlite3.connect("db.sqlite3", check_same_thread=False)
+        _memory = SqliteSaver(con)
+
+        # Create agent
+        _agent_executor = create_react_agent(
+            llm,
+            [],
+            prompt="You are a helpful AI image editing assistant. \
+                You help users with image editing tasks and provide guidance \
+                    on how to modify their images.",
+            checkpointer=_memory,
+        )
+
+    return _agent_executor
+
+
+def chat_with_agent(
+    message: str, user_id: str = "default", selected_images: Optional[List[str]] = None
+) -> str:
+    """
+    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 names (optional)
+
+    Returns:
+        The agent's response as a string
+    """
+    agent = get_agent()
+
+    # Prepare the message with context
+    full_message = message
+    if selected_images and len(selected_images) > 0:
+        image_context = f" Selected images: {', '.join(selected_images)}."
+        full_message = message + image_context
+
+    # Configure thread ID for conversation continuity
+    config = {"configurable": {"thread_id": user_id}}
+
+    # Get response from agent
+    response = agent.invoke(
+        {"messages": [{"role": "user", "content": full_message}]}, config=config
+    )
+
+    # Extract the last message from the agent
+    if response and "messages" in response and len(response["messages"]) > 0:
+        last_message = response["messages"][-1]
+        # Handle both AIMessage objects and dictionaries
+        if hasattr(last_message, "content"):
+            return last_message.content
+        elif isinstance(last_message, dict) and "content" in last_message:
+            return last_message["content"]
+
+    return "I'm sorry, I couldn't process your request. Please try again."
+
+
+if __name__ == "__main__":
+    # Test the agent
+    response = chat_with_agent("Hello! How can you help me with image editing?")
+    print(response)
diff --git a/api/pyproject.toml b/api/pyproject.toml
@@ -19,6 +19,7 @@ dependencies = [
     "langchain-core",
     "langgraph>0.2.27",
     "langchain[google-genai]",
+    "langgraph-checkpoint-sqlite",
 ]
 
 [project.optional-dependencies]
@@ -42,6 +43,7 @@ target-version = ["py310"]
 line-length = 88
 target-version = "py310"
 fix = true
+unsafe-fixes = true
 
 [tool.ruff.lint]
 select = ["E", "F", "I"]
diff --git a/api/server/main.py b/api/server/main.py
@@ -3,6 +3,8 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 
+from llm.agent import chat_with_agent
+
 app = FastAPI(
     title="AI Image Editor API",
     description="API for AI-powered image editing assistant",
@@ -21,84 +23,6 @@ class ChatResponse(BaseModel):
     status: str = "success"
 
 
-# Template responses for different types of messages
-TEMPLATE_RESPONSES = {
-    "greeting": [
-        "Hello! I'm your AI image editing assistant. How can I help you today?",
-        "Hi there! What would you like to work on?",
-        "Welcome! I'm your AI assistant ready to help you transform your images.",
-    ],
-    "image_selection": [
-        "I can see you've selected some images. What would you like to do with them?",
-        "Great choice! Those images look interesting. What kind of editing you need?",
-        "Perfect! I can help you edit those selected images. What's your vision?",
-    ],
-    "editing_request": [
-        "I understand you want to edit your images. Let me help you with that!",
-        "Great! I can assist you with image editing. What specific changes you need?",
-        "Excellent! I'm ready to help you transform your images.",
-    ],
-    "general_help": [
-        "Just let me know what you'd like to do!",
-        "Feel free to ask me anything about image editing.",
-    ],
-    "upload": [
-        "I see you've uploaded an image! What would you like to do with it?",
-        "Great! I can help you edit that uploaded image.",
-        "Perfect! I'm ready to work with your uploaded image. What's your vision?",
-    ],
-}
-
-
-def get_template_response(
-    message: str, selected_images: Optional[List[str]] = None
-) -> str:
-    """Get a template response based on the message content and context."""
-    message_lower = message.lower()
-
-    # Check for greetings
-    if any(word in message_lower for word in ["hello", "hi", "hey", "greetings"]):
-        import random
-
-        return random.choice(TEMPLATE_RESPONSES["greeting"])
-
-    # Check for image uploads
-    if "uploaded" in message_lower or "📷" in message:
-        import random
-
-        return random.choice(TEMPLATE_RESPONSES["upload"])
-
-    # Check for image selection context
-    if selected_images and len(selected_images) > 0:
-        import random
-
-        base_response = random.choice(TEMPLATE_RESPONSES["image_selection"])
-        image_names = ", ".join(selected_images)
-        return f"{base_response} I can see you've selected: {image_names}."
-
-    # Check for editing requests
-    if any(
-        word in message_lower
-        for word in [
-            "edit",
-            "change",
-            "modify",
-            "transform",
-            "enhance",
-            "filter",
-            "effect",
-        ]
-    ):
-        import random
-
-        return random.choice(TEMPLATE_RESPONSES["editing_request"])
-
-    # Default response
-    import random
-
-    return random.choice(TEMPLATE_RESPONSES["general_help"])
-
-
 @app.get("/")
 async def root():
     return {"message": "AI Image Editor API is running!"}
@@ -121,8 +45,13 @@ async def chat_endpoint(request: ChatRequest):
         ChatResponse with AI response and status.
     """
     try:
-        # Get template response based on message and context
-        response = get_template_response(request.message, request.selected_images)
+        # Use the LLM agent to get a response
+        user_id = request.user_id or "default"
+        response = chat_with_agent(
+            message=request.message,
+            user_id=user_id,
+            selected_images=request.selected_images,
+        )
 
         return ChatResponse(response=response, status="success")
 
