task: refactors as LangGraph

src/graph.ipynb

@@ -1,317 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import os\n",
-    "import getpass\n",
-    "\n",
-    "\n",
-    "os.environ[\"OPENAI_API_KEY\"] = getpass.getpass(\"OpenAI API Key:\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Add tools later"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Loaded 82 design documents\n",
-      "Testing RAG retriever with requirements:\n",
-      "\n",
-      "Retrieved Designs:\n",
-      "----------------------------------------\n",
-      "Generated query: \"vintage classic easy to use grandmother love design\"\n",
-      "Design 180:\n",
-      "Description: This design employs a vintage newspaper aesthetic with a classic serif typography that evokes an old-world charm, utilizing sepia-toned paper backgrounds to enhance its nostalgic feel. The layout is text-heavy with a deliberate obfuscation, reflecting a layered collage effect. Its balanced placement keeps the focus central, inviting closer inspection and interaction.\n",
-      "Categories: Vintage, Nostalgic, Typography, Collage, Editorial\n",
-      "Visual Characteristics: Sepia tone, Serif typography, Textured background, Layered elements, Central focus\n",
-      "URL: https://csszengarden.com/180\n",
-      "\n",
-      "Design 182:\n",
-      "Description: The design creatively utilizes a retro theme with vinyl records as the prominent visual element to evoke a sense of nostalgia and classic style, complemented by a muted green color palette that brings harmony and balance. Handwritten and vintage-style typography enhance the retro aesthetic, while background illustrations and decorative elements like stars add whimsy and depth to the composition.\n",
-      "Categories: Retro, Nostalgic, Music-themed, Decorative, Vintage\n",
-      "Visual Characteristics: Vinyl Records, Muted Green Palette, Handwritten Typography, Background Illustrations, Decorative Elements\n",
-      "URL: https://csszengarden.com/182\n",
-      "\n",
-      "Design 194:\n",
-      "Description: This design exudes a minimalist elegance with a muted, earthy color palette and a clean layout, embodying a sense of calm and sophistication. The subtle use of textures and classic serif typography enhances the refined aesthetic, while the centered alignment and generous spacing contribute to a relaxed readability. The incorporation of a delicate floral illustration adds a touch of organic charm, making the design feel both timeless and inviting.\n",
-      "Categories: Minimalism, Elegant, Organic, Sophisticated, Classic\n",
-      "Visual Characteristics: Muted Color Palette, Serif Typography, Centered Layout, Generous Spacing, Floral Illustration\n",
-      "URL: https://csszengarden.com/194\n",
-      "\n",
-      "Design 212:\n",
-      "Description: The design features a retro aesthetic using a muted color palette of browns and creams, creating a nostalgic and vintage feel. The asymmetrical layout and bold typography contribute to the visual hierarchy, guiding the viewer through the content effortlessly. Illustrations with a mid-century modern style add character, merging traditional design elements with contemporary functionality.\n",
-      "Categories: Retro, Typography, Illustration, Vintage Style, Educational\n",
-      "Visual Characteristics: Muted Color Palette, Asymmetrical Layout, Bold Typography, Retro Illustrations, Functional Design\n",
-      "URL: https://csszengarden.com/212\n"
-     ]
-    }
-   ],
-   "source": [
-    "#from tools.design_retriever import DesignRetrieverTool\n",
-    "from chains.design_rag import DesignRAG\n",
-    "\n",
-    "# Initialize DesignRAG and create the tool\n",
-    "design_rag = DesignRAG()\n",
-    "#design_retriever = DesignRetrieverTool(rag=design_rag)\n",
-    "\n",
-    "test_requirements = {\n",
-    "    \"I want a design that is vintage and classic, something easy to use that a grandmother would love\"\n",
-    "    }\n",
-    "\n",
-    "# Test the retriever\n",
-    "async def test_rag():\n",
-    "    print(\"Testing RAG retriever with requirements:\")\n",
-    "    print(\"\\nRetrieved Designs:\")\n",
-    "    print(\"----------------------------------------\")\n",
-    "    \n",
-    "    results = await design_rag.query_similar_designs(test_requirements, 2)\n",
-    "    print(results)\n",
-    "\n",
-    "# Run the test\n",
-    "await test_rag()\n"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Pick a model good for chat and tools"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "RunnableBinding(bound=ChatOpenAI(client=<openai.resources.chat.completions.completions.Completions object at 0x1245518d0>, async_client=<openai.resources.chat.completions.completions.AsyncCompletions object at 0x124548e50>, root_client=<openai.OpenAI object at 0x1108f9310>, root_async_client=<openai.AsyncOpenAI object at 0x115d92090>, model_name='gpt-4o', temperature=0.0, model_kwargs={}, openai_api_key=SecretStr('**********'), streaming=True), kwargs={'tools': [{'type': 'function', 'function': {'name': 'design_retriever', 'description': 'Retrieves similar designs based on style requirements', 'parameters': {'properties': {'requirements': {'type': 'object'}, 'num_examples': {'default': 3, 'type': 'integer'}}, 'required': ['requirements'], 'type': 'object'}}}]}, config={}, config_factories=[])"
-      ]
-     },
-     "execution_count": 11,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "from langchain_openai import ChatOpenAI\n",
-    "\n",
-    "model = ChatOpenAI(\n",
-    "    model=\"gpt-4o\", \n",
-    "    temperature=0,\n",
-    "    streaming=True\n",
-    ")\n",
-    "\n",
-    "model.bind_tools(tool_belt)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Initialize state\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 12,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from typing import TypedDict, Annotated\n",
-    "from langgraph.graph.message import add_messages\n",
-    "\n",
-    "class AgentState(TypedDict):\n",
-    "  messages: Annotated[list, add_messages]"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Set up the nodes and graph\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 16,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from langgraph.prebuilt import ToolNode\n",
-    "from langgraph.graph import StateGraph, END\n",
-    "from langchain_core.messages import HumanMessage, SystemMessage\n",
-    "\n",
-    "system_message = SystemMessage(content=\"\"\"You are a helpful design assistant that can retrieve and analyze design examples. \n",
-    "When a user describes their design preferences or requirements, use the design_retriever tool to find relevant examples.\n",
-    "\n",
-    "Always use the design_retriever tool when:\n",
-    "- A user describes specific design requirements\n",
-    "- A user asks to see similar designs\n",
-    "- You need to find design inspiration based on user preferences\n",
-    "\n",
-    "Format the requirements as a dictionary with these keys:\n",
-    "- style_description: Brief description of desired visual style\n",
-    "- key_elements: List of important visual elements\n",
-    "- color_scheme: Description of colors\n",
-    "- layout_preferences: Layout requirements\n",
-    "- mood: Desired emotional impact\n",
-    "\"\"\")\n",
-    "\n",
-    "def call_model(state):\n",
-    "  messages = [system_message] + state[\"messages\"]\n",
-    "  response = model.invoke(messages)\n",
-    "  return {\"messages\" : [response]}\n",
-    "\n",
-    "tool_node = ToolNode(tool_belt)\n",
-    "\n",
-    "uncompiled_graph = StateGraph(AgentState)\n",
-    "\n",
-    "uncompiled_graph.add_node(\"agent\", call_model)\n",
-    "uncompiled_graph.add_node(\"action\", tool_node)\n",
-    "uncompiled_graph.set_entry_point(\"agent\")\n",
-    "\n",
-    "\n",
-    "def should_continue(state):\n",
-    "  last_message = state[\"messages\"][-1]\n",
-    "\n",
-    "  if last_message.tool_calls:\n",
-    "    return \"action\"\n",
-    "\n",
-    "  return END\n",
-    "\n",
-    "uncompiled_graph.add_conditional_edges(\n",
-    "  \"agent\",\n",
-    "  should_continue\n",
-    ")\n",
-    "uncompiled_graph.add_edge(\"action\", \"agent\")\n",
-    "\n",
-    "graph = uncompiled_graph.compile()\n",
-    "\n",
-    "#formatted chain\n",
-    "\n",
-    "def convert_inputs(input_object):\n",
-    "  return {\"messages\" : [HumanMessage(content=input_object[\"question\"])]}\n",
-    "\n",
-    "def parse_output(input_state):\n",
-    "  return input_state[\"messages\"][-1].content\n",
-    "\n",
-    "graph_chain = convert_inputs | graph | parse_output\n",
-    "\n"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Try it out!"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Receiving update from node: 'agent'\n",
-      "[AIMessage(content=\"Hello! I'm here and ready to help you with any design needs or questions you might have. How can I assist you today?\", additional_kwargs={}, response_metadata={'finish_reason': 'stop', 'model_name': 'gpt-4o-2024-08-06', 'system_fingerprint': 'fp_f9f4fb6dbf'}, id='run-4edce0b5-fdec-4d5d-a4a6-92430faca51a-0')]\n",
-      "\n",
-      "\n",
-      "\n"
-     ]
-    }
-   ],
-   "source": [
-    "from langchain_core.messages import HumanMessage\n",
-    "\n",
-    "async for chunk in graph.astream({\"messages\" : [HumanMessage(content=\"Hello, how are you?\")]}, stream_mode=\"updates\"):\n",
-    "    for node, values in chunk.items():\n",
-    "        print(f\"Receiving update from node: '{node}'\")\n",
-    "        print(values[\"messages\"])\n",
-    "        print(\"\\n\\n\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Let's see if the RAG tool works."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 18,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Receiving update from node: 'agent'\n",
-      "[AIMessage(content=\"To find a design that matches your description, I'll use the design_retriever tool. Here are the requirements based on your description:\\n\\n- style_description: Monochromatic with subtle accents\\n- key_elements: Grid-based layout, clear hierarchy\\n- color_scheme: Monochromatic with subtle accent colors\\n- layout_preferences: Grid-based\\n- mood: Professional and sophisticated\\n\\nLet's find some examples for you.\", additional_kwargs={}, response_metadata={'finish_reason': 'stop', 'model_name': 'gpt-4o-2024-08-06', 'system_fingerprint': 'fp_f9f4fb6dbf'}, id='run-8fa2e4af-671c-4c75-82fd-a7b3d6237e54-0')]\n",
-      "\n",
-      "\n",
-      "\n"
-     ]
-    }
-   ],
-   "source": [
-    "# Create a test message\n",
-    "from langchain_core.messages import HumanMessage\n",
-    "\n",
-    "test_message = HumanMessage(\n",
-    "    content=\"\"\"I want to see a design matching this description: \n",
-    "    I want it to use a monochromatic color scheme with subtle accent colors. \n",
-    "    The layout should be grid-based with clear hierarchy. \n",
-    "    The overall mood should be professional and sophisticated.\"\"\"\n",
-    ")\n",
-    "\n",
-    "async for chunk in graph.astream({\"messages\" : [test_message]}, stream_mode=\"updates\"):\n",
-    "    for node, values in chunk.items():\n",
-    "        print(f\"Receiving update from node: '{node}'\")\n",
-    "        print(values[\"messages\"])\n",
-    "        print(\"\\n\\n\")"
-   ]
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": ".venv",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.11.11"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 2
-}

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
 
 class State(TypedDict):
     # Messages have the type "list". The `add_messages` function
@@ -12,7 +12,40 @@ class State(TypedDict):
     # (in this case, it appends messages to the list, rather than overwriting them)
     messages: Annotated[list, add_messages]
 
-
-graph_builder = StateGraph(State)
+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()
+
+# Create the graph
+graph = create_graph()
 
 

src/nodes/analyzer.py

@@ -0,0 +1 @@
+

src/nodes/designer.py

@@ -0,0 +1,122 @@
+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