Add langgraph studio config and refactor

agent/__init__.py

@@ -0,0 +1,3 @@
+from agent.agent import graph, graph_with_memory
+
+__all__ = ["graph", "graph_with_memory"] 

agent/agent.py

@@ -0,0 +1,50 @@
+from langgraph.graph import StateGraph
+from langgraph.checkpoint.memory import MemorySaver
+
+from agent.utils.state import AgentState
+from agent.utils.nodes import call_model, tool_node, should_continue
+
+def create_agent_graph():
+    # Create the graph
+    builder = StateGraph(AgentState)
+
+    # Add nodes
+    builder.add_node("agent", call_model)
+    builder.add_node("action", tool_node)
+
+    # Update edges
+    builder.set_entry_point("agent")
+    builder.add_conditional_edges(
+        "agent",
+        should_continue,
+    )
+    builder.add_edge("action", "agent")
+
+    # Initialize memory saver for conversation persistence
+    memory = MemorySaver()
+
+    # Compile the graph with memory
+    return builder.compile(checkpointer=memory)
+
+def create_agent_graph_without_memory():
+    # Create the graph
+    builder = StateGraph(AgentState)
+
+    # Add nodes
+    builder.add_node("agent", call_model)
+    builder.add_node("action", tool_node)
+
+    # Update edges
+    builder.set_entry_point("agent")
+    builder.add_conditional_edges(
+        "agent",
+        should_continue,
+    )
+    builder.add_edge("action", "agent")
+
+    # Compile the graph without memory
+    return builder.compile()
+
+# Create both graph variants
+graph_with_memory = create_agent_graph()
+graph = create_agent_graph_without_memory() 

agent/utils/__init__.py

@@ -0,0 +1 @@
+ 

agent/utils/nodes.py

@@ -0,0 +1,35 @@
+from langchain_openai import ChatOpenAI
+from langchain_core.messages import SystemMessage
+from langgraph.graph import END
+from langgraph.prebuilt import ToolNode
+
+from agent.utils.tools import tool_belt
+from agent.utils.state import AgentState
+
+# Initialize LLM
+llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
+model = llm.bind_tools(tool_belt)
+
+# Define system prompt
+SYSTEM_PROMPT = SystemMessage(content="""
+You are a helpful AI assistant that answers questions clearly and concisely.
+If you don't know something, simply say you don't know.
+Be engaging and professional in your responses.
+Use the retrieve_context tool when you need specific information about events and activities.
+Use the tavily_search tool for general web searches.
+""")
+
+def call_model(state: AgentState):
+    messages = [SYSTEM_PROMPT] + state["messages"]
+    response = model.invoke(messages)
+    return {"messages": [response]}
+
+# Initialize tool node
+tool_node = ToolNode(tool_belt)
+
+# Simple flow control - always go to final
+def should_continue(state):
+    last_message = state["messages"][-1]
+    if last_message.tool_calls:
+        return "action"
+    return END 

agent/utils/state.py

@@ -0,0 +1,6 @@
+from typing import Annotated, TypedDict
+from langgraph.graph.message import add_messages
+
+class AgentState(TypedDict):
+    messages: Annotated[list, add_messages]
+    context: list  # Store retrieved context 

agent/utils/tools.py

@@ -0,0 +1,28 @@
+from langchain_core.tools import tool
+from langchain_community.tools.tavily_search import TavilySearchResults
+from rag import create_rag_pipeline, add_urls_to_vectorstore
+
+# Initialize RAG pipeline
+rag_components = create_rag_pipeline(collection_name="london_events")
+
+# Add some initial URLs to the vector store
+urls = [
+    "https://www.timeout.com/london/things-to-do-in-london-this-weekend",
+    "https://www.timeout.com/london/london-events-in-march"
+]
+add_urls_to_vectorstore(
+    rag_components["vector_store"],
+    rag_components["text_splitter"],
+    urls
+)
+
+@tool
+def retrieve_context(query: str) -> list[str]:
+    """Searches the knowledge base for relevant information about events and activities. Use this when you need specific details about events."""
+    return [doc.page_content for doc in rag_components["retriever"].get_relevant_documents(query)]
+
+# Initialize Tavily search tool
+tavily_tool = TavilySearchResults(max_results=5)
+
+# Create tool belt
+tool_belt = [tavily_tool, retrieve_context] 

app.py

@@ -1,93 +1,9 @@
 import uuid
-from typing import Annotated, TypedDict, Literal
-from langchain_openai import ChatOpenAI
-from langgraph.graph import StateGraph, START, END
-from langgraph.graph.message import MessagesState, add_messages
-from langgraph.prebuilt import ToolNode
-from langgraph.checkpoint.memory import MemorySaver
-
-from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
+from langchain_core.messages import HumanMessage, AIMessage
 from langchain.schema.runnable.config import RunnableConfig
-from langchain_community.tools.tavily_search import TavilySearchResults
-from langchain.tools import Tool
-from langchain_core.tools import tool
-
 import chainlit as cl
-from rag import create_rag_pipeline, add_urls_to_vectorstore
-
-# Initialize RAG pipeline
-rag_components = create_rag_pipeline(collection_name="london_events")
-
-# Add some initial URLs to the vector store
-urls = [
-    "https://www.timeout.com/london/things-to-do-in-london-this-weekend",
-    "https://www.timeout.com/london/london-events-in-march"
-]
-add_urls_to_vectorstore(
-    rag_components["vector_store"],
-    rag_components["text_splitter"],
-    urls
-)
-
-class AgentState(TypedDict):
-    messages: Annotated[list, add_messages]
-    context: list  # Store retrieved context
-
-# Create a retrieve tool
-@tool
-def retrieve_context(query: str) -> list[str]:
-    """Searches the knowledge base for relevant information about events and activities. Use this when you need specific details about events."""
-    return [doc.page_content for doc in rag_components["retriever"].get_relevant_documents(query)]
-
-tavily_tool = TavilySearchResults(max_results=5)
-tool_belt = [tavily_tool, retrieve_context]
-
-llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
-model = llm.bind_tools(tool_belt)
-
-# Define system prompt
-SYSTEM_PROMPT = SystemMessage(content="""
-You are a helpful AI assistant that answers questions clearly and concisely.
-If you don't know something, simply say you don't know.
-Be engaging and professional in your responses.
-Use the retrieve_context tool when you need specific information about events and activities.
-Use the tavily_search tool for general web searches.
-""")
-
-def call_model(state: AgentState):
-    messages = [SYSTEM_PROMPT] + state["messages"]
-    response = model.invoke(messages)
-    return {"messages": [response]}
-
-tool_node = ToolNode(tool_belt)
-
-# Simple flow control - always go to final
-def should_continue(state):
-    last_message = state["messages"][-1]
-    if last_message.tool_calls:
-        return "action"
-    return END
-
-# Create the graph
-builder = StateGraph(AgentState)
-
-# Remove retrieve node and modify graph structure
-builder.add_node("agent", call_model)
-builder.add_node("action", tool_node)
-
-# Update edges
-builder.set_entry_point("agent")
-builder.add_conditional_edges(
-    "agent",
-    should_continue,
-)
-builder.add_edge("action", "agent")
-
-# Initialize memory saver for conversation persistence
-memory = MemorySaver()
-
-# Compile the graph with memory
-graph = builder.compile(checkpointer=memory)
+from agent import graph_with_memory as graph
+from agent.utils.state import AgentState
 
 @cl.on_chat_start
 async def on_chat_start():
@@ -109,7 +25,7 @@ async def on_chat_start():
     # Initialize empty state with auth
     try:
         await graph.ainvoke(
-            {"messages": [], "context": []},
+            AgentState(messages=[], context=[]),
             config=config
         )
     except Exception as e:
@@ -159,7 +75,7 @@ async def on_message(message: cl.Message):
     
     # Stream the response
     async for chunk in graph.astream(
-        {"messages": current_messages, "context": []},
+        AgentState(messages=current_messages, context=[]),
         config=RunnableConfig(
             configurable={
                 "thread_id": session_id,

langgraph.json

@@ -0,0 +1,7 @@
+{
+  "dependencies": ".",
+  "graphs": {
+    "agent": "agent:graph"
+  },
+  "env": ".env"
+}

rag.py

@@ -1,6 +1,6 @@
-from typing import List, TypedDict
+from typing import List
 from langchain_core.documents import Document
-from langchain_openai import ChatOpenAI, OpenAIEmbeddings
+from langchain_openai import OpenAIEmbeddings
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_qdrant import QdrantVectorStore
 from qdrant_client import QdrantClient

tools.py

@@ -0,0 +1,28 @@
+from langchain_core.tools import tool
+from langchain_community.tools.tavily_search import TavilySearchResults
+from rag import create_rag_pipeline, add_urls_to_vectorstore
+
+# Initialize RAG pipeline
+rag_components = create_rag_pipeline(collection_name="london_events")
+
+# Add some initial URLs to the vector store
+urls = [
+    "https://www.timeout.com/london/things-to-do-in-london-this-weekend",
+    "https://www.timeout.com/london/london-events-in-march"
+]
+add_urls_to_vectorstore(
+    rag_components["vector_store"],
+    rag_components["text_splitter"],
+    urls
+)
+
+@tool
+def retrieve_context(query: str) -> list[str]:
+    """Searches the knowledge base for relevant information about events and activities. Use this when you need specific details about events."""
+    return [doc.page_content for doc in rag_components["retriever"].get_relevant_documents(query)]
+
+# Initialize Tavily search tool
+tavily_tool = TavilySearchResults(max_results=5)
+
+# Create tool belt
+tool_belt = [tavily_tool, retrieve_context]