task: working graph

src/app.py

@@ -28,25 +28,24 @@ async def init():
 
 @cl.on_message
 async def main(message: cl.Message):
-    # Get the graph and current state from the user session
+    # Get the graph from the user session
     graph = cl.user_session.get("graph")
+    
+    # Get current state
     state = cl.user_session.get("state")
     
-    # Add user message to state
+    # Add the new user message to the state
     state["messages"].append(HumanMessage(content=message.content))
     
     # Process message through the graph
     result = await graph.ainvoke(state)
+    print("Here's the result: ", result)
     
     # Update state with the result
     state["messages"].extend(result["messages"])
     
     # 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."
-    )
+    last_message = result["messages"][-1].content
     
     # Send response to user
     await cl.Message(content=last_message).send()

src/graph.ipynb

@@ -0,0 +1,104 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from typing import Annotated\n",
+    "from typing_extensions import TypedDict\n",
+    "from langgraph.graph.message import add_messages\n",
+    "from langgraph.prebuilt import create_react_agent\n",
+    "from langchain.tools.render import format_tool_to_openai_function\n",
+    "from nodes.design_rag import DesignRAG\n",
+    "from langchain_openai import ChatOpenAI\n",
+    "from langgraph.prebuilt import ToolNode\n",
+    "from tools.design_retriever import design_retriever_tool\n",
+    "\n",
+    "class State(TypedDict):\n",
+    "    # Messages have the type \"list\". The `add_messages` function\n",
+    "    # in the annotation defines how this state key should be updated\n",
+    "    # (in this case, it appends messages to the list, rather than overwriting them)\n",
+    "    messages: Annotated[list, add_messages]\n",
+    "\n",
+    "model = ChatOpenAI(model=\"gpt-4o-mini\", temperature=0)\n",
+    "\n",
+    "tools = [\n",
+    "    design_retriever_tool\n",
+    "]\n",
+    "\n",
+    "tool_node = ToolNode(tools=tools)\n",
+    "\n",
+    "model_with_tools = model.bind_tools(tools)\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[{'name': 'design_retriever_tool',\n",
+       "  'args': {'state': {'messages': [{'content': 'Can you show me two designs with a comic book style?',\n",
+       "      'type': 'human'}]},\n",
+       "   'num_examples': 2},\n",
+       "  'id': 'call_6GZ0WpBwYGzmioC0p9IyyZ7h',\n",
+       "  'type': 'tool_call'}]"
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "model_with_tools.invoke(\"Can you show me two designs with a comic book style?\").tool_calls"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "<coroutine object RunnableCallable.ainvoke at 0x110215fc0>"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "tool_node.ainvoke({\"messages\": [model_with_tools.invoke(\"Can you show me a design with a comic book style?\")]})"
+   ]
+  }
+ ],
+ "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

@@ -20,5 +20,8 @@ tools = [
 
 model_with_tools = model.bind_tools(tools)
 
-graph = create_react_agent(model_with_tools, tools=tools)
+graph = create_react_agent(
+    model_with_tools,
+    tools=tools
+)
 

src/nodes/design_rag.py

@@ -29,7 +29,7 @@ class DesignRAG:
         # Create retriever with tracing
         self.retriever = self.vector_store.as_retriever(
             search_type="similarity",
-            search_kwargs={"k": 1},
+            search_kwargs={"k": 4},
             tags=["design_retriever"]  # Add tags for tracing
         )
         
@@ -50,25 +50,19 @@ class DesignRAG:
             # Load all metadata files
             for design_dir in designs_dir.glob("**/metadata.json"):
                 try:
+                    print(f"Processing design: {design_dir.parent.name}")
                     with open(design_dir, "r") as f:
                         metadata = json.load(f)
                     
                     # Create document text from metadata with safe gets
                     text = f"""
-                    Design {metadata.get('id', 'unknown')}:
-                    Description: {metadata.get('description', 'No description available')}
+                    Title: {metadata.get('title', 'Untitled')}
+                    Author: {metadata.get('author', 'Unknown')}
+                    Description: {metadata.get('description', {}).get('summary', 'No description available')}
                     Categories: {', '.join(metadata.get('categories', []))}
                     Visual Characteristics: {', '.join(metadata.get('visual_characteristics', []))}
+                    Artistic Style: {metadata.get('artistic_context', 'None Provided').get('style_influences', "None specified")}
                     """
-                    
-                    # Load associated CSS
-                    '''
-                    css_path = design_dir.parent / "style.css"
-                    if css_path.exists():
-                        with open(css_path, "r") as f:
-                            css = f.read()
-                        text += f"\nCSS:\n{css}"
-                    '''
 
                     # Create Document object with minimal metadata
                     documents.append(
@@ -76,12 +70,14 @@ class DesignRAG:
                             page_content=text.strip(),
                             metadata={
                                 "id": metadata.get('id', 'unknown'),
-                                "path": str(design_dir.parent)
+                                "path": str(design_dir.parent),
+                                "title": metadata.get('title', 'Untitled'),
+                                "author": metadata.get('author', 'Unknown')
                             }
                         )
                     )
                 except Exception as e:
-                    print(f"Error processing design {design_dir}: {e}")
+                    print(f"Error processing design {design_dir.parent.name}: {str(e)}")
                     continue
             
             if not documents:
@@ -140,17 +136,48 @@ class DesignRAG:
         docs = self.retriever.get_relevant_documents(
             query_response.content, 
             k=num_examples,
-            callbacks=[ConsoleCallbackHandler()]  # Use standard callback instead
+            callbacks=[ConsoleCallbackHandler()]
         )
         
-        # Format examples
+        # Format examples with improved readability
         examples = []
         for doc in docs:
             design_id = doc.metadata.get("id", "unknown")
+            title = doc.metadata.get("title", "Untitled")
+            author = doc.metadata.get("author", "Unknown")
+            
+            # Parse the content into sections
             content_lines = doc.page_content.strip().split("\n")
-            examples.append(
-                "\n".join(line.strip() for line in content_lines if line.strip()) +
-                f"\nURL: https://csszengarden.com/{design_id}"
-            )
+            sections = {}
+            current_section = None
+            
+            for line in content_lines:
+                line = line.strip()
+                if not line:
+                    continue
+                if ":" in line:
+                    current_section, value = line.split(":", 1)
+                    sections[current_section.strip()] = value.strip()
+            
+            # Format the example with clear sections
+            example = f"""
+Design: {title}
+By: {author}
+
+Description:
+{sections.get('Description', 'No description available')}
+
+Categories:
+{sections.get('Categories', 'No categories available')}
+
+Visual Characteristics:
+{sections.get('Visual Characteristics', 'No characteristics available')}
+
+Artistic Style:
+{sections.get('Artistic Style', 'No style information available')}
+
+View at: https://csszengarden.com/{design_id}
+"""
+            examples.append(example.strip())
         
-        return "\n\n".join(examples) 
+        return "\n\n" + "="*50 + "\n\n".join(examples) 

src/tools/design_retriever.py

@@ -1,10 +1,16 @@
 from nodes.design_rag import DesignRAG
 from langgraph.graph import MessagesState
+from langchain_core.messages import SystemMessage
 
-def design_retriever_tool(state: MessagesState, num_examples: int = 2):
+async 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)
+    rag = DesignRAG()  # Create instance
+    
+    result = await rag.query_similar_designs(state["messages"], num_examples)
+    print("Here's the result: ", result)
+
+    return SystemMessage(content=result)