task: begins notebook to demonstrate rag retrieval

src/graph.ipynb

@@ -0,0 +1,249 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "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": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Loaded 82 design documents\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",
+    "\n",
+    "tool_belt = [design_retriever]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Pick a model good for chat and tools"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "RunnableBinding(bound=ChatOpenAI(client=<openai.resources.chat.completions.completions.Completions object at 0x10d8b4750>, async_client=<openai.resources.chat.completions.completions.AsyncCompletions object at 0x10d8c1110>, root_client=<openai.OpenAI object at 0x10b08ef10>, root_async_client=<openai.AsyncOpenAI object at 0x10d8b4910>, model_name='gpt-4o-mini', temperature=0.0, model_kwargs={}, openai_api_key=SecretStr('**********')), kwargs={'tools': []}, config={}, config_factories=[])"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from langchain_openai import ChatOpenAI\n",
+    "\n",
+    "model = ChatOpenAI(model=\"gpt-4o-mini\", temperature=0)\n",
+    "\n",
+    "model.bind_tools(tool_belt)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Initialize state\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "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": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "Graph(nodes={'__start__': Node(id='__start__', name='__start__', data=<class 'langchain_core.utils.pydantic.LangGraphInput'>, metadata=None), 'agent': Node(id='agent', name='agent', data=agent(tags=None, recurse=True, explode_args=False, func_accepts_config=False, func_accepts={}), metadata=None), 'action': Node(id='action', name='action', data=tools(tags=None, recurse=True, explode_args=False, func_accepts_config=True, func_accepts={'store': ('__pregel_store', None)}, tools_by_name={}, tool_to_state_args={}, tool_to_store_arg={}, handle_tool_errors=True, messages_key='messages'), metadata=None), '__end__': Node(id='__end__', name='__end__', data=<class 'langchain_core.utils.pydantic.LangGraphOutput'>, metadata=None)}, edges=[Edge(source='__start__', target='agent', data=None, conditional=False), Edge(source='action', target='agent', data=None, conditional=False), Edge(source='agent', target='action', data=None, conditional=True), Edge(source='agent', target='__end__', data=None, conditional=True)])"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from langgraph.prebuilt import ToolNode\n",
+    "from langgraph.graph import StateGraph, END\n",
+    "\n",
+    "def call_model(state):\n",
+    "  messages = 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"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Try it out!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'messages': [HumanMessage(content='Hello, how are you?', additional_kwargs={}, response_metadata={}, id='1af14df1-568c-460c-997f-11d59e70a3b3'),\n",
+       "  AIMessage(content=\"Hello! I'm just a computer program, so I don't have feelings, but I'm here and ready to help you. How can I assist you today?\", additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 31, 'prompt_tokens': 13, 'total_tokens': 44, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_709714d124', 'finish_reason': 'stop', 'logprobs': None}, id='run-d532debb-85a6-4730-9786-43543f1cfd43-0', usage_metadata={'input_tokens': 13, 'output_tokens': 31, 'total_tokens': 44, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})]}"
+      ]
+     },
+     "execution_count": 11,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from langchain_core.messages import HumanMessage\n",
+    "\n",
+    "graph.invoke({\"messages\" : [HumanMessage(content=\"Hello, how are you?\")]})"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Let's see if the RAG tool works."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "\n",
+    "# Test the tool\n",
+    "test_requirements = {\n",
+    "    \"style_description\": \"Modern and minimalist\",\n",
+    "    \"key_elements\": [\"clean lines\", \"white space\", \"typography\"],\n",
+    "    \"color_scheme\": \"Monochromatic with subtle accent colors\",\n",
+    "    \"layout_preferences\": \"Grid-based layout with clear hierarchy\",\n",
+    "    \"mood\": \"Professional and sophisticated\"\n",
+    "}\n",
+    "\n",
+    "# Create a test message\n",
+    "from langchain_core.messages import HumanMessage\n",
+    "\n",
+    "test_message = HumanMessage(\n",
+    "    content=\"\"\"I need a design that's modern and minimalist, with clean lines and plenty of white space. \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",
+    "# Invoke the graph with the test message\n",
+    "response = graph.invoke({\"messages\": [test_message]})\n",
+    "print(response)"
+   ]
+  }
+ ],
+ "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
+}