Upgrade Opengpts (#361)

* Migrate pydantic * Upgrade poetry * Adapt to manage checkpoint using an AstncSaver * Adjust Tools model * Add checkpoint * Update poetry * Format * Fix tests * Modify tables * Fix gpt4o * Fix bots * Fix retrieval * Adding eugenes suggestions * Fix state handling inconsistency between different agent types * Improve doc * Update backend/pyproject.toml * lint fix * lint --------- Co-authored-by: “lgesuellip” <“[email protected]”> Co-authored-by: Eugene Yurtsev <[email protected]> Co-authored-by: Eugene Yurtsev <[email protected]>
langchain-ai · Jan 29, 2025 · 2cf3bf7 · 2cf3bf7
1 parent 541ae6f
commit 2cf3bf7
Show file tree

Hide file tree

Showing 27 changed files with 3,411 additions and 2,723 deletions.
diff --git a/README.md b/README.md
@@ -227,6 +227,16 @@ Navigate to [http://localhost:5173/](http://localhost:5173/) and enjoy!
 
 Refer to this [guide](tools/redis_to_postgres/README.md) for migrating data from Redis to Postgres.
 
+## Breaking Changes
+
+### Migration 5 - Checkpoint Management Update
+Version 5 of the database migrations introduces a significant change to how thread checkpoints are managed:
+- Transitions from a pickle-based checkpointing system to a new multi-table checkpoint management system (breaking change)
+- Aligns with LangGraph's new checkpoint architecture for better state management and persistence
+- **Important**: Historical threads/checkpoints (created before this migration) will not be accessible in the UI
+- Previous checkpoint data is preserved in the `old_checkpoints` table but cannot be accessed by the new system
+- This architectural change improves how thread state is stored and managed, enabling more reliable state persistence in LangGraph-based agents.
+
 ## Features
 
 As much as possible, we are striving for feature parity with OpenAI.
@@ -309,14 +319,14 @@ Then, those documents are passed in the system message to a separate call to the
 
 Compared to assistants, it is more structured (but less powerful). It ALWAYS looks up something - which is good if you 
 know you want to look things up, but potentially wasteful if the user is just trying to have a normal conversation. 
-Also importantly, this only looks up things once - so if it doesn’t find the right results then it will yield a bad 
+Also importantly, this only looks up things once - so if it doesn't find the right results then it will yield a bad 
 result (compared to an assistant, which could  decide to look things up again).
 
 ![](_static/rag.png)
 
 Despite this being a more simple architecture, it is good for a few reasons. First, because it is simpler it can work 
 pretty well with a wider variety of models (including lots of open source models). Second, if you have a use case where 
-you don’t NEED the flexibility of an assistant (eg you know users will be looking up information every time) then it 
+you don't NEED the flexibility of an assistant (eg you know users will be looking up information every time) then it 
 can be more focused. And third, compared to the final architecture below it can use external knowledge.
 
 RAGBot is implemented with [LangGraph](https://github.com/langchain-ai/langgraph) `StateGraph`. A `StateGraph` is a generalized graph that can model arbitrary state (i.e. `dict`), not just a `list` of messages.
@@ -325,7 +335,7 @@ RAGBot is implemented with [LangGraph](https://github.com/langchain-ai/langgraph
 
 The final architecture is dead simple - just a call to a language model, parameterized by a system message. This allows 
 the GPT to take on different personas and characters. This is clearly far less powerful than Assistants or RAGBots 
-(which have access to external sources of data/computation) - but it’s still valuable! A lot of popular GPTs are just 
+(which have access to external sources of data/computation) - but it's still valuable! A lot of popular GPTs are just 
 system messages at the end of the day, and CharacterAI is crushing it despite largely just being system messages as 
 well.
 

diff --git a/backend/README.md b/backend/README.md
@@ -1 +1,26 @@
 # backend
+
+## Database Migrations
+
+### Migration 5 - Checkpoint Management Update
+This migration introduces a significant change to thread checkpoint management:
+
+#### Changes
+- Transitions from single-table pickle storage to a robust multi-table checkpoint management system
+- Implements LangGraph's latest checkpoint architecture for improved state persistence
+- Preserves existing checkpoint data by renaming `checkpoints` table to `old_checkpoints`
+- Introduces three new tables for better checkpoint management:
+  - `checkpoints`: Core checkpoint metadata
+  - `checkpoint_blobs`: Actual checkpoint data storage (compatible with LangGraph state serialization)
+  - `checkpoint_writes`: Tracks checkpoint write operations
+- Adds runtime initialization via `ensure_setup()` in the lifespan event
+
+#### Impact
+- **Breaking Change**: Historical threads/checkpoints (pre-migration) will not be accessible in the UI
+- Previous checkpoint data remains preserved but inaccessible in the new system
+- Designed to work seamlessly with LangGraph's state persistence requirements
+
+#### Migration Details
+- **Up Migration**: Safely preserves existing data by renaming the table
+- **Down Migration**: Restores original table structure if needed
+- New checkpoint management tables are automatically created at application startup
diff --git a/backend/app/agent.py b/backend/app/agent.py
@@ -1,4 +1,3 @@
-import pickle
 from enum import Enum
 from typing import Any, Dict, Mapping, Optional, Sequence, Union
 
@@ -7,14 +6,13 @@
     ConfigurableField,
     RunnableBinding,
 )
-from langgraph.checkpoint import CheckpointAt
 from langgraph.graph.message import Messages
 from langgraph.pregel import Pregel
 
 from app.agent_types.tools_agent import get_tools_agent_executor
 from app.agent_types.xml_agent import get_xml_agent_executor
 from app.chatbot import get_chatbot_executor
-from app.checkpoint import PostgresCheckpoint
+from app.checkpoint import AsyncPostgresCheckpoint
 from app.llms import (
     get_anthropic_llm,
     get_google_llm,
@@ -74,7 +72,7 @@ class AgentType(str, Enum):
 
 DEFAULT_SYSTEM_MESSAGE = "You are a helpful assistant."
 
-CHECKPOINTER = PostgresCheckpoint(serde=pickle, at=CheckpointAt.END_OF_STEP)
+CHECKPOINTER = AsyncPostgresCheckpoint()
 
 
 def get_agent_executor(
@@ -123,7 +121,6 @@ def get_agent_executor(
         return get_tools_agent_executor(
             tools, llm, system_message, interrupt_before_action, CHECKPOINTER
         )
-
     else:
         raise ValueError("Unexpected agent type")
 
@@ -135,7 +132,7 @@ class ConfigurableAgent(RunnableBinding):
     retrieval_description: str = RETRIEVAL_DESCRIPTION
     interrupt_before_action: bool = False
     assistant_id: Optional[str] = None
-    thread_id: Optional[str] = None
+    thread_id: Optional[str] = ""
     user_id: Optional[str] = None
 
     def __init__(
@@ -145,7 +142,7 @@ def __init__(
         agent: AgentType = AgentType.GPT_35_TURBO,
         system_message: str = DEFAULT_SYSTEM_MESSAGE,
         assistant_id: Optional[str] = None,
-        thread_id: Optional[str] = None,
+        thread_id: Optional[str] = "",
         retrieval_description: str = RETRIEVAL_DESCRIPTION,
         interrupt_before_action: bool = False,
         kwargs: Optional[Mapping[str, Any]] = None,
@@ -204,7 +201,9 @@ def get_chatbot(
     if llm_type == LLMType.GPT_35_TURBO:
         llm = get_openai_llm()
     elif llm_type == LLMType.GPT_4:
-        llm = get_openai_llm(gpt_4=True)
+        llm = get_openai_llm(model="gpt-4")
+    elif llm_type == LLMType.GPT_4O:
+        llm = get_openai_llm(model="gpt-4o")
     elif llm_type == LLMType.AZURE_OPENAI:
         llm = get_openai_llm(azure=True)
     elif llm_type == LLMType.CLAUDE2:
@@ -265,7 +264,7 @@ class ConfigurableRetrieval(RunnableBinding):
     llm_type: LLMType
     system_message: str = DEFAULT_SYSTEM_MESSAGE
     assistant_id: Optional[str] = None
-    thread_id: Optional[str] = None
+    thread_id: Optional[str] = ""
     user_id: Optional[str] = None
 
     def __init__(
@@ -274,7 +273,7 @@ def __init__(
         llm_type: LLMType = LLMType.GPT_35_TURBO,
         system_message: str = DEFAULT_SYSTEM_MESSAGE,
         assistant_id: Optional[str] = None,
-        thread_id: Optional[str] = None,
+        thread_id: Optional[str] = "",
         kwargs: Optional[Mapping[str, Any]] = None,
         config: Optional[Mapping[str, Any]] = None,
         **others: Any,
@@ -319,7 +318,9 @@ def __init__(
         assistant_id=ConfigurableField(
             id="assistant_id", name="Assistant ID", is_shared=True
         ),
-        thread_id=ConfigurableField(id="thread_id", name="Thread ID", is_shared=True),
+        thread_id=ConfigurableField(
+            id="thread_id", name="Thread ID", annotation=str, is_shared=True
+        ),
     )
     .with_types(
         input_type=Dict[str, Any],
@@ -335,7 +336,7 @@ def __init__(
         system_message=DEFAULT_SYSTEM_MESSAGE,
         retrieval_description=RETRIEVAL_DESCRIPTION,
         assistant_id=None,
-        thread_id=None,
+        thread_id="",
     )
     .configurable_fields(
         agent=ConfigurableField(id="agent_type", name="Agent Type"),
@@ -348,7 +349,9 @@ def __init__(
         assistant_id=ConfigurableField(
             id="assistant_id", name="Assistant ID", is_shared=True
         ),
-        thread_id=ConfigurableField(id="thread_id", name="Thread ID", is_shared=True),
+        thread_id=ConfigurableField(
+            id="thread_id", name="Thread ID", annotation=str, is_shared=True
+        ),
         tools=ConfigurableField(id="tools", name="Tools"),
         retrieval_description=ConfigurableField(
             id="retrieval_description", name="Retrieval Description"

diff --git a/backend/app/agent_types/tools_agent.py b/backend/app/agent_types/tools_agent.py
@@ -28,7 +28,7 @@ async def _get_messages(messages):
         msgs = []
         for m in messages:
             if isinstance(m, LiberalToolMessage):
-                _dict = m.dict()
+                _dict = m.model_dump()
                 _dict["content"] = str(_dict["content"])
                 m_c = ToolMessage(**_dict)
                 msgs.append(m_c)

diff --git a/backend/app/agent_types/xml_agent.py b/backend/app/agent_types/xml_agent.py
@@ -45,7 +45,7 @@ def construct_chat_history(messages):
                 temp_messages = []
             collapsed_messages.append(message)
         elif isinstance(message, LiberalFunctionMessage):
-            _dict = message.dict()
+            _dict = message.model_dump()
             _dict["content"] = str(_dict["content"])
             m_c = FunctionMessage(**_dict)
             temp_messages.append(m_c)

diff --git a/backend/app/api/assistants.py b/backend/app/api/assistants.py
@@ -14,9 +14,11 @@
 class AssistantPayload(BaseModel):
     """Payload for creating an assistant."""
 
-    name: str = Field(..., description="The name of the assistant.")
-    config: dict = Field(..., description="The assistant config.")
-    public: bool = Field(default=False, description="Whether the assistant is public.")
+    name: Annotated[str, Field(description="The name of the assistant.")]
+    config: Annotated[dict, Field(description="The assistant config.")]
+    public: Annotated[
+        bool, Field(default=False, description="Whether the assistant is public.")
+    ]
 
 
 AssistantID = Annotated[str, Path(description="The ID of the assistant.")]
@@ -25,7 +27,7 @@ class AssistantPayload(BaseModel):
 @router.get("/")
 async def list_assistants(user: AuthedUser) -> List[Assistant]:
     """List all assistants for the current user."""
-    return await storage.list_assistants(user["user_id"])
+    return await storage.list_assistants(user.user_id)
 
 
 @router.get("/public/")
@@ -40,7 +42,7 @@ async def get_assistant(
     aid: AssistantID,
 ) -> Assistant:
     """Get an assistant by ID."""
-    assistant = await storage.get_assistant(user["user_id"], aid)
+    assistant = await storage.get_assistant(user.user_id, aid)
     if not assistant:
         raise HTTPException(status_code=404, detail="Assistant not found")
     return assistant
@@ -53,7 +55,7 @@ async def create_assistant(
 ) -> Assistant:
     """Create an assistant."""
     return await storage.put_assistant(
-        user["user_id"],
+        user.user_id,
         str(uuid4()),
         name=payload.name,
         config=payload.config,
@@ -69,7 +71,7 @@ async def upsert_assistant(
 ) -> Assistant:
     """Create or update an assistant."""
     return await storage.put_assistant(
-        user["user_id"],
+        user.user_id,
         aid,
         name=payload.name,
         config=payload.config,
@@ -83,5 +85,5 @@ async def delete_assistant(
     aid: AssistantID,
 ):
     """Delete an assistant by ID."""
-    await storage.delete_assistant(user["user_id"], aid)
+    await storage.delete_assistant(user.user_id, aid)
     return {"status": "ok"}
diff --git a/backend/app/api/runs.py b/backend/app/api/runs.py
@@ -4,14 +4,13 @@
 import langsmith.client
 from fastapi import APIRouter, BackgroundTasks, HTTPException
 from fastapi.exceptions import RequestValidationError
-from langchain.pydantic_v1 import ValidationError
 from langchain_core.messages import AnyMessage
 from langchain_core.runnables import RunnableConfig
 from langsmith.utils import tracing_is_enabled
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, ValidationError
 from sse_starlette import EventSourceResponse
 
-from app.agent import agent
+from app.agent import agent, chat_retrieval, chatbot
 from app.auth.handlers import AuthedUser
 from app.storage import get_assistant, get_thread
 from app.stream import astream_state, to_sse
@@ -34,24 +33,34 @@ async def _run_input_and_config(payload: CreateRunPayload, user_id: str):
     if not thread:
         raise HTTPException(status_code=404, detail="Thread not found")
 
-    assistant = await get_assistant(user_id, str(thread["assistant_id"]))
+    assistant = await get_assistant(user_id, str(thread.assistant_id))
     if not assistant:
         raise HTTPException(status_code=404, detail="Assistant not found")
 
     config: RunnableConfig = {
-        **assistant["config"],
+        **assistant.config,
         "configurable": {
-            **assistant["config"]["configurable"],
+            **assistant.config["configurable"],
             **((payload.config or {}).get("configurable") or {}),
             "user_id": user_id,
-            "thread_id": str(thread["thread_id"]),
-            "assistant_id": str(assistant["assistant_id"]),
+            "thread_id": str(thread.thread_id),
+            "assistant_id": str(assistant.assistant_id),
         },
     }
 
     try:
         if payload.input is not None:
-            agent.get_input_schema(config).validate(payload.input)
+            # Get the bot type from config
+            bot_type = config["configurable"].get("type", "agent")
+            # Get the correct schema based on bot type
+            if bot_type == "chat_retrieval":
+                schema = chat_retrieval.get_input_schema()
+            elif bot_type == "chatbot":
+                schema = chatbot.get_input_schema()
+            else:  # default to agent
+                schema = agent.get_input_schema()
+            # Validate against the correct schema
+            schema.model_validate(payload.input)
     except ValidationError as e:
         raise RequestValidationError(e.errors(), body=payload)
 
@@ -65,7 +74,7 @@ async def create_run(
     background_tasks: BackgroundTasks,
 ):
     """Create a run."""
-    input_, config = await _run_input_and_config(payload, user["user_id"])
+    input_, config = await _run_input_and_config(payload, user.user_id)
     background_tasks.add_task(agent.ainvoke, input_, config)
     return {"status": "ok"}  # TODO add a run id
 
@@ -76,27 +85,27 @@ async def stream_run(
     user: AuthedUser,
 ):
     """Create a run."""
-    input_, config = await _run_input_and_config(payload, user["user_id"])
+    input_, config = await _run_input_and_config(payload, user.user_id)
 
     return EventSourceResponse(to_sse(astream_state(agent, input_, config)))
 
 
 @router.get("/input_schema")
 async def input_schema() -> dict:
     """Return the input schema of the runnable."""
-    return agent.get_input_schema().schema()
+    return agent.get_input_schema().model_json_schema()
 
 
 @router.get("/output_schema")
 async def output_schema() -> dict:
     """Return the output schema of the runnable."""
-    return agent.get_output_schema().schema()
+    return agent.get_output_schema().model_json_schema()
 
 
 @router.get("/config_schema")
 async def config_schema() -> dict:
     """Return the config schema of the runnable."""
-    return agent.config_schema().schema()
+    return agent.config_schema().model_json_schema()
 
 
 if tracing_is_enabled():