src/graph.ipynb

{
 "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
}