Upgrade Opengpts

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().get_instance()
 
 
 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,7 @@ 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.AZURE_OPENAI:
         llm = get_openai_llm(azure=True)
     elif llm_type == LLMType.CLAUDE2:
@@ -265,7 +262,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 +271,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 +316,7 @@ 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 +332,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 +345,7 @@ 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"

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)

backend/app/api/assistants.py

@@ -14,9 +14,9 @@
 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 +25,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 +40,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 +53,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 +69,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 +83,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"}

backend/app/api/runs.py

@@ -4,7 +4,7 @@
 import langsmith.client
 from fastapi import APIRouter, BackgroundTasks, HTTPException
 from fastapi.exceptions import RequestValidationError
-from langchain.pydantic_v1 import ValidationError
+from pydantic import ValidationError
 from langchain_core.messages import AnyMessage
 from langchain_core.runnables import RunnableConfig
 from langsmith.utils import tracing_is_enabled
@@ -34,24 +34,24 @@ 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)
+            agent.get_input_schema(config).model_validate(payload.input)
     except ValidationError as e:
         raise RequestValidationError(e.errors(), body=payload)
 
@@ -65,7 +65,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 +76,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():

backend/app/api/threads.py

@@ -17,9 +17,8 @@
 
 class ThreadPutRequest(BaseModel):
     """Payload for creating a thread."""
-
-    name: str = Field(..., description="The name of the thread.")
-    assistant_id: str = Field(..., description="The ID of the assistant to use.")
+    name: Annotated[str, Field(description="The name of the thread.")]
+    assistant_id: Annotated[str, Field(description="The ID of the assistant to use.")]
 
 
 class ThreadPostRequest(BaseModel):
@@ -32,7 +31,7 @@ class ThreadPostRequest(BaseModel):
 @router.get("/")
 async def list_threads(user: AuthedUser) -> List[Thread]:
     """List all threads for the current user."""
-    return await storage.list_threads(user["user_id"])
+    return await storage.list_threads(user.user_id)
 
 
 @router.get("/{tid}/state")
@@ -41,14 +40,14 @@ async def get_thread_state(
     tid: ThreadID,
 ):
     """Get state for a thread."""
-    thread = await storage.get_thread(user["user_id"], tid)
+    thread = await storage.get_thread(user.user_id, tid)
     if not thread:
         raise HTTPException(status_code=404, detail="Thread not found")
-    assistant = await storage.get_assistant(user["user_id"], thread["assistant_id"])
+    assistant = await storage.get_assistant(user.user_id, thread.assistant_id)
     if not assistant:
         raise HTTPException(status_code=400, detail="Thread has no assistant")
     return await storage.get_thread_state(
-        user_id=user["user_id"],
+        user_id=user.user_id,
         thread_id=tid,
         assistant=assistant,
     )
@@ -61,16 +60,16 @@ async def add_thread_state(
     payload: ThreadPostRequest,
 ):
     """Add state to a thread."""
-    thread = await storage.get_thread(user["user_id"], tid)
+    thread = await storage.get_thread(user.user_id, tid)
     if not thread:
         raise HTTPException(status_code=404, detail="Thread not found")
-    assistant = await storage.get_assistant(user["user_id"], thread["assistant_id"])
+    assistant = await storage.get_assistant(user.user_id, thread.assistant_id)
     if not assistant:
         raise HTTPException(status_code=400, detail="Thread has no assistant")
     return await storage.update_thread_state(
         payload.config or {"configurable": {"thread_id": tid}},
         payload.values,
-        user_id=user["user_id"],
+        user_id=user.user_id,
         assistant=assistant,
     )
 
@@ -81,14 +80,14 @@ async def get_thread_history(
     tid: ThreadID,
 ):
     """Get all past states for a thread."""
-    thread = await storage.get_thread(user["user_id"], tid)
+    thread = await storage.get_thread(user.user_id, tid)
     if not thread:
         raise HTTPException(status_code=404, detail="Thread not found")
-    assistant = await storage.get_assistant(user["user_id"], thread["assistant_id"])
+    assistant = await storage.get_assistant(user.user_id, thread.assistant_id)
     if not assistant:
         raise HTTPException(status_code=400, detail="Thread has no assistant")
     return await storage.get_thread_history(
-        user_id=user["user_id"],
+        user_id=user.user_id,
         thread_id=tid,
         assistant=assistant,
     )
@@ -100,7 +99,7 @@ async def get_thread(
     tid: ThreadID,
 ) -> Thread:
     """Get a thread by ID."""
-    thread = await storage.get_thread(user["user_id"], tid)
+    thread = await storage.get_thread(user.user_id, tid)
     if not thread:
         raise HTTPException(status_code=404, detail="Thread not found")
     return thread
@@ -113,7 +112,7 @@ async def create_thread(
 ) -> Thread:
     """Create a thread."""
     return await storage.put_thread(
-        user["user_id"],
+        user.user_id,
         str(uuid4()),
         assistant_id=thread_put_request.assistant_id,
         name=thread_put_request.name,
@@ -128,7 +127,7 @@ async def upsert_thread(
 ) -> Thread:
     """Update a thread."""
     return await storage.put_thread(
-        user["user_id"],
+        user.user_id,
         tid,
         assistant_id=thread_put_request.assistant_id,
         name=thread_put_request.name,
@@ -141,5 +140,5 @@ async def delete_thread(
     tid: ThreadID,
 ):
     """Delete a thread by ID."""
-    await storage.delete_thread(user["user_id"], tid)
+    await storage.delete_thread(user.user_id, tid)
     return {"status": "ok"}