Merge pull request #31 from langchain-ai/nc/26feb/checkpoint

Make checkpoint an object
langchain-ai · Feb 26, 2025 · d2b1b80 · d2b1b80
2 parents 6d54b52 + d0655d1
commit d2b1b80
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 9 deletions.
diff --git a/openapi.json b/openapi.json
@@ -2,7 +2,7 @@
   "openapi": "3.1.0",
   "info": {
     "title": "Agent Protocol",
-    "version": "0.1.4"
+    "version": "0.1.5"
   },
   "tags": [
     {
@@ -2196,13 +2196,29 @@
         ],
         "title": "Thread"
       },
-      "ThreadState": {
+      "ThreadCheckpoint": {
         "properties": {
           "checkpoint_id": {
             "type": "string",
             "format": "uuid",
             "title": "Checkpoint Id",
             "description": "The ID of the checkpoint."
+          }
+        },
+        "type": "object",
+        "additionalProperties": true,
+        "required": [
+          "checkpoint_id"
+        ],
+        "title": "ThreadCheckpoint",
+        "description": "Structured identifier for a thread checkpoint, ie. an entry in the thread's history."
+      },
+      "ThreadState": {
+        "properties": {
+          "checkpoint": {
+            "$ref": "#/components/schemas/ThreadCheckpoint",
+            "title": "Checkpoint",
+            "description": "The identifier for this checkpoint."
           },
           "values": {
             "type": "object",
@@ -2215,13 +2231,17 @@
               "$ref": "#/components/schemas/Message"
             },
             "title": "Messages",
-            "description": "The current Messages of the thread. If messages are contained in Thread.values, implementations should remove them from values when returning messages. When this key isn't present it means the thread/agent doesn't support messages."
+            "description": "The current messages of the thread. If messages are contained in Thread.values, implementations should remove them from values when returning messages. When this key isn't present it means the thread/agent doesn't support messages."
+          },
+          "metadata": {
+            "type": "object",
+            "title": "Metadata",
+            "description": "The checkpoint metadata."
           }
         },
         "type": "object",
         "required": [
-          "thread_id",
-          "checkpoint_id",
+          "checkpoint",
           "values"
         ],
         "title": "ThreadState"
@@ -2256,6 +2276,11 @@
       },
       "ThreadPatch": {
         "properties": {
+          "checkpoint": {
+            "$ref": "#/components/schemas/ThreadCheckpoint",
+            "title": "Checkpoint",
+            "description": "The identifier of the checkpoint to branch from. Ignored for metadata-only patches. If not provided, defaults to the latest checkpoint."
+          },
           "metadata": {
             "type": "object",
             "title": "Metadata",

diff --git a/server/ap_server/main.py b/server/ap_server/main.py
@@ -9,7 +9,7 @@
 
 app = FastAPI(
     title="Agent Protocol",
-    version="0.1.3",
+    version="0.1.5",
 )
 
 app.include_router(agents.router)

diff --git a/server/ap_server/models.py b/server/ap_server/models.py
@@ -249,6 +249,15 @@ class ThreadSearchRequest(BaseModel):
     )
 
 
+class ThreadCheckpoint(BaseModel):
+    model_config = ConfigDict(
+        extra="allow",
+    )
+    checkpoint_id: UUID = Field(
+        ..., description="The ID of the checkpoint.", title="Checkpoint Id"
+    )
+
+
 class IfExists(Enum):
     raise_ = "raise"
     do_nothing = "do_nothing"
@@ -465,20 +474,28 @@ class Thread(BaseModel):
 
 
 class ThreadState(BaseModel):
-    checkpoint_id: UUID = Field(
-        ..., description="The ID of the checkpoint.", title="Checkpoint Id"
+    checkpoint: ThreadCheckpoint = Field(
+        ..., description="The identifier for this checkpoint.", title="Checkpoint"
     )
     values: Dict[str, Any] = Field(
         ..., description="The current state of the thread.", title="Values"
     )
     messages: Optional[List[Message]] = Field(
         None,
-        description="The current Messages of the thread. If messages are contained in Thread.values, implementations should remove them from values when returning messages. When this key isn't present it means the thread/agent doesn't support messages.",
+        description="The current messages of the thread. If messages are contained in Thread.values, implementations should remove them from values when returning messages. When this key isn't present it means the thread/agent doesn't support messages.",
         title="Messages",
     )
+    metadata: Optional[Dict[str, Any]] = Field(
+        None, description="The checkpoint metadata.", title="Metadata"
+    )
 
 
 class ThreadPatch(BaseModel):
+    checkpoint: Optional[ThreadCheckpoint] = Field(
+        None,
+        description="The identifier of the checkpoint to branch from. Ignored for metadata-only patches. If not provided, defaults to the latest checkpoint.",
+        title="Checkpoint",
+    )
     metadata: Optional[Dict[str, Any]] = Field(
         None,
         description="Metadata to merge with existing thread metadata.",