task: [wip] initial graph

src/app.py

@@ -1,10 +1,6 @@
 import chainlit as cl
-from langchain_openai import ChatOpenAI
 from langchain_core.messages import HumanMessage, SystemMessage
-from nodes.design_rag import DesignRAG
-
-# Initialize components
-design_rag = DesignRAG()
+from graph import graph
 
 # System message focused on design analysis
 SYSTEM_MESSAGE = """You are a helpful design assistant that finds and explains design examples.
@@ -18,51 +14,42 @@ For every user message, analyze their design preferences and requirements, consi
 
 @cl.on_chat_start
 async def init():
-    # Initialize LLM with callback handler inside the Chainlit context
-    llm = ChatOpenAI(
-        model="gpt-4o-mini",
-        temperature=0,
-        streaming=True,
-        callbacks=[cl.LangchainCallbackHandler()]
-    )
-    
-    # Store the LLM in the user session
-    cl.user_session.set("design_llm", llm)
+    # Store the graph in the user session
+    cl.user_session.set("graph", graph)
     
-    # init conversation history for each user
-    cl.user_session.set("conversation_history", [
-        SystemMessage(content=SYSTEM_MESSAGE)
-    ])
+    # Initialize conversation state with system message
+    initial_state = {
+        "messages": [SystemMessage(content=SYSTEM_MESSAGE)]
+    }
+    cl.user_session.set("state", initial_state)
     
     # Send welcome message
     await cl.Message(content="Welcome to ImagineUI! I'm here to help you design beautiful and functional user interfaces. What kind of design are you looking for?").send()
 
 @cl.on_message
 async def main(message: cl.Message):
-    # Get the LLM from the user session
-    llm = cl.user_session.get("design_llm")
-    conversation_history = cl.user_session.get("conversation_history")
+    # Get the graph and current state from the user session
+    graph = cl.user_session.get("graph")
+    state = cl.user_session.get("state")
     
-    # Add user message to history
-    conversation_history.append(HumanMessage(content=message.content))
+    # Add user message to state
+    state["messages"].append(HumanMessage(content=message.content))
     
-    # Get LLM's analysis of requirements
-    analysis = await llm.ainvoke(conversation_history)
+    # Process message through the graph
+    result = await graph.ainvoke(state)
     
-    # Get best design example based on full conversation
-    designs = await design_rag.query_similar_designs(
-        [msg.content for msg in conversation_history],
-        num_examples=1
-    )
-    
-    # Combine analysis with designs
-    response = f"{analysis.content}\n\nHere is the best match from the zen garden:\n\n{designs}"
+    # Update state with the result
+    state["messages"].extend(result["messages"])
     
-    # Add assistant's response to history
-    conversation_history.append(SystemMessage(content=response))
+    # Extract the last assistant message for display
+    last_message = next(
+        (msg.content for msg in reversed(result["messages"]) 
+         if isinstance(msg, SystemMessage)),
+        "I apologize, but I couldn't process your request."
+    )
     
     # Send response to user
-    await cl.Message(content=response).send()
+    await cl.Message(content=last_message).send()
 
 if __name__ == "__main__":
-    cl.run() 
+    cl.run()

src/graph.py

@@ -1,10 +1,10 @@
 from typing import Annotated
 from typing_extensions import TypedDict
-from langgraph.graph import StateGraph, START, END
 from langgraph.graph.message import add_messages
-from langgraph.prebuilt import ToolInvoker
-from nodes.designer import DesignerNode
-from langchain.tools.render import format_tool_to_openai_function
+from langgraph.prebuilt import create_react_agent
+from langchain_anthropic import ChatAnthropic
+from tools.design_retriever import design_retriever_tool
+
 
 class State(TypedDict):
     # Messages have the type "list". The `add_messages` function
@@ -12,40 +12,13 @@ class State(TypedDict):
     # (in this case, it appends messages to the list, rather than overwriting them)
     messages: Annotated[list, add_messages]
 
-def create_graph():
-    # Initialize nodes
-    designer = DesignerNode()
-    
-    # Create graph
-    graph = StateGraph(State)
-    
-    # Add designer node
-    graph.add_node("designer", designer)
-    
-    # Create tool invoker node with designer's tools
-    tools = designer.get_available_tools()
-    tool_executor = ToolInvoker(tools=tools)
-    graph.add_node("tools", tool_executor)
-    
-    # Add edges
-    graph.add_edge(START, "designer")
-    
-    # Add conditional edges based on tool calls
-    graph.add_conditional_edges(
-        "designer",
-        lambda state: "tools" if state["messages"][-1].get("tool_calls") else END,
-        {
-            "tools": "tools",
-            END: END
-        }
-    )
-    
-    # After tool execution, return to designer
-    graph.add_edge("tools", "designer")
-    
-    return graph.compile()
+model = ChatAnthropic(model="claude-3-5-sonnet-20240620", temperature=0)
+
+tools = [
+    design_retriever_tool
+]
 
-# Create the graph
-graph = create_graph()
+model_with_tools = model.bind_tools(tools)
 
+graph = create_react_agent(model_with_tools, tools=tools)
 

src/nodes/designer.py

@@ -1,122 +0,0 @@
-from typing import Dict, List
-from anthropic import AsyncAnthropic
-import json
-from langchain_core.tools import tool
-from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
-from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
-from nodes.design_rag import DesignRAG
-
-class DesignerNode:
-    """Main conversation node for discussing design requirements and retrieving examples"""
-    
-    def __init__(self):
-        self.client = AsyncAnthropic()
-        self.rag = DesignRAG()
-        
-        # Define the conversation prompt
-        self.prompt = ChatPromptTemplate.from_messages([
-            ("system", """You are an expert design assistant helping users find design inspiration.
-            Your goal is to understand their design needs and requirements through conversation.
-            
-            Guidelines:
-            1. Focus on understanding visual design requirements, not implementation
-            2. Ask clarifying questions about style, mood, and visual elements
-            3. When the user asks to see examples, use the retrieve_design_examples tool
-            4. Track both must-have requirements and nice-to-have preferences
-            5. When showing examples, explain how they match the requirements
-            
-            Available tools:
-            - retrieve_design_examples: Find relevant design examples based on conversation
-            
-            When the user asks to see examples, ALWAYS use the retrieve_design_examples tool.
-            Format tool calls using the exact function name and parameters.
-            """),
-            MessagesPlaceholder(variable_name="chat_history"),
-            ("human", "{input}"),
-        ])
-
-    @tool()
-    async def retrieve_design_examples(self, conversation: List[str], num_examples: int = 1) -> str:
-        """
-        Find and retrieve relevant design examples based on the conversation history.
-        
-        Args:
-            conversation: List of conversation messages
-            num_examples: Number of examples to retrieve (default: 1)
-            
-        Returns:
-            String containing design examples and their details
-        """
-        return await self.rag.query_similar_designs(conversation, num_examples)
-
-    def get_available_tools(self):
-        """Return list of available tools"""
-        return [self.retrieve_design_examples]
-
-    async def __call__(self, state: Dict) -> Dict:
-        """Process messages and manage design discussion"""
-        messages = state.get("messages", [])
-        
-        # Convert messages to chat history format
-        chat_history = []
-        for msg in messages[:-1]:  # Exclude the last message which is the current input
-            if isinstance(msg, dict):
-                role = msg.get("role", "user")
-                content = msg.get("content", "")
-                chat_history.append(
-                    HumanMessage(content=content) if role == "user" 
-                    else AIMessage(content=content)
-                )
-            elif isinstance(msg, BaseMessage):
-                chat_history.append(msg)
-        
-        # Get the current input message
-        current_input = messages[-1].get("content") if isinstance(messages[-1], dict) else messages[-1].content
-        
-        # Get response from Claude
-        response = await self.client.messages.create(
-            model="claude-3-haiku-20240307",
-            max_tokens=500,
-            messages=[{
-                "role": "user",
-                "content": self.prompt.format(
-                    chat_history=chat_history,
-                    input=current_input
-                )
-            }]
-        )
-        
-        response_text = response.content[0].text
-        
-        # Check if response indicates need for examples
-        should_retrieve = (
-            "retrieve_design_examples" in response_text or
-            any(phrase in current_input.lower() 
-                for phrase in ["show example", "find design", "get example"])
-        )
-        
-        if should_retrieve:
-            # Create tool call message
-            state["messages"].append({
-                "role": "assistant",
-                "content": response_text,
-                "tool_calls": [{
-                    "type": "function",
-                    "function": {
-                        "name": "retrieve_design_examples",
-                        "arguments": json.dumps({
-                            "conversation": [msg.get("content", msg) if isinstance(msg, dict) else msg 
-                                          for msg in messages],
-                            "num_examples": 1
-                        })
-                    }
-                }]
-            })
-        else:
-            # Regular response without tool calls
-            state["messages"].append({
-                "role": "assistant",
-                "content": response_text
-            })
-        
-        return state

src/tools/design_retriever.py

@@ -1,25 +1,10 @@
-from typing import Dict, Optional
-from langchain.tools import BaseTool
-from chains.design_rag import DesignRAG
-from pydantic import Field
-import json
+from nodes.design_rag import DesignRAG
+from langgraph.graph import MessagesState
+
+def design_retriever_tool(state: MessagesState, num_examples: int = 2):
+    """
+        Retrieves similar designs based on style requirements
+        Name: query_similar_designs
+    """
+    return DesignRAG.query_similar_designs(state["messages"], num_examples)
 
-class DesignRetrieverTool(BaseTool):
-    """Tool for retrieving similar designs based on requirements."""
-    
-    name: str = "design_retriever"
-    description: str = "Retrieves similar designs based on style requirements"
-    rag: DesignRAG = Field(description="Design RAG system for retrieving similar designs")
-    
-    def __init__(self, rag: DesignRAG):
-        """Initialize the tool with a DesignRAG instance."""
-        super().__init__(rag=rag)
-    
-    def _run(self, requirements: Dict, num_examples: int = 3) -> str:
-        """Sync version - not used but required by BaseTool"""
-        raise NotImplementedError("Use async version")
-    
-    async def _arun(self, requirements: Dict, num_examples: int = 3) -> str:
-        """Retrieve similar designs based on requirements"""
-        print(f"Retrieving {num_examples} similar designs")
-        return await self.rag.query_similar_designs(requirements, num_examples)