wip DE checkpointing, add entities models

Author: nemozak1

src/agent/components/chains.py

@@ -1,5 +1,4 @@
 # Description: Contains the chains for the main agent system
-from langchain import hub
 from langchain_core.prompts import SystemMessagePromptTemplate, PromptTemplate, ChatPromptTemplate, MessagesPlaceholder
 from langchain_core.messages import HumanMessage, SystemMessage
 from langchain_core.output_parsers import StrOutputParser

src/checkpointer/entities.py

@@ -1,6 +1,7 @@
 from datetime import datetime
 from typing import Optional
 from pydantic import BaseModel, Field
+from src.client.obp_client import OBPClient
 
 class OpeyCheckpointEntity(BaseModel):
     """OBP Dynamic Entity representation of a LangGraph checkpoint."""
@@ -19,19 +20,15 @@ def obp_entity_name(cls) -> str:
 
     @classmethod
     def to_obp_schema(cls) -> dict:
+        json_schema = cls.model_json_schema()
         return {
             "hasPersonalEntity": True,
             cls.obp_entity_name(): {
-                "description": "LangGraph checkpoint snapshots for Opey AI conversations",
-                "required": ["thread_id", "checkpoint_id", "checkpoint_ns", "checkpoint_data"],
+                "description": cls.__doc__.strip(),
+                "required": json_schema["required"],
                 "properties": {
-                    "thread_id": {"type": "string", "description": "Thread identifier"},
-                    "checkpoint_id": {"type": "string", "description": "Unique checkpoint ID (UUID)"},
-                    "checkpoint_ns": {"type": "string", "description": "Checkpoint namespace"},
-                    "parent_checkpoint_id": {"type": "string", "description": "Parent checkpoint ID for history"},
-                    "checkpoint_data": {"type": "string", "description": "JSON-serialized checkpoint (channel_values, versions, etc.)"},
-                    "metadata": {"type": "string", "description": "JSON-serialized metadata (step, source, writes)"},
-                    "created_at": {"type": "string", "description": "ISO timestamp"}
+                    k: {"type": v.get("type", "string"), "description": v.get("description", "")}
+                    for k, v in json_schema["properties"].items()
                 }
             }
         }
@@ -54,24 +51,48 @@ def obp_entity_name(cls) -> str:
 
     @classmethod
     def to_obp_schema(cls) -> dict:
-        """OBP Dynamic Entity schema definition."""
+        json_schema = cls.model_json_schema()
         return {
             "hasPersonalEntity": True,
             cls.obp_entity_name(): {
-                "description": "Pending writes for LangGraph checkpoints",
-                "required": ["thread_id", "checkpoint_id", "checkpoint_ns", "task_id", "channel", "value", "idx"],
+                "description": cls.__doc__.strip(),
+                "required": json_schema["required"],
                 "properties": {
-                    "thread_id": {"type": "string"},
-                    "checkpoint_id": {"type": "string"},
-                    "checkpoint_ns": {"type": "string"},
-                    "task_id": {"type": "string", "description": "Task that produced this write"},
-                    "idx": {"type": "integer", "description": "Write index within task"},
-                    "channel": {"type": "string", "description": "Channel name"},
-                    "value": {"type": "string", "description": "Serialized write value"}
+                    k: {"type": v.get("type", "string"), "description": v.get("description", "")}
+                    for k, v in json_schema["properties"].items()
                 }
             }
         }
 
+class DynamicEntitiesManager:
+    """Wrapper for DynamicEntities CRUD endpoints."""
+    def __init__(self, obp_client: OBPClient, endpoint_url: str):
+        
+        self.endpoint_url = endpoint_url
+        self.client = obp_client
+        
+
+    async def create(self, entity: OpeyCheckpointEntity | OpeyCheckpointWriteEntity):
+        await self.client.async_obp_requests(
+            method="POST",
+            body=str(entity.model_dump()),
+            path=self.endpoint_url
+        )
+        
+    async def delete(self, entity_id: str):
+        await self.client.async_obp_requests(
+            method="DELETE",
+            body="",
+            path=f"{self.endpoint_url}/{entity_id}"
+        )
+    
+    async def read(self, entity_id: str) -> dict:
+        response = await self.client.async_obp_requests(
+            method="GET",
+            body="",
+            path=f"{self.endpoint_url}/{entity_id}"
+        )
+        return response
 
 opey_checkpoint_entity = {
   "hasPersonalEntity": True,
@@ -105,4 +126,4 @@ def to_obp_schema(cls) -> dict:
       "value": {"type": "string", "description": "JSON-serialized write value"}
     }
   }
-}
+}

src/checkpointer/obp_checkpoint_saver.py

@@ -0,0 +1,40 @@
+from typing import Iterator, Optional, Sequence
+from langchain_core.runnables import RunnableConfig
+from langgraph.checkpoint.base import (
+    BaseCheckpointSaver,
+    Checkpoint,
+    CheckpointMetadata,
+    CheckpointTuple,
+    ChannelVersions,
+    PendingWrite
+)
+from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
+
+from src.client import obp_client
+from src.checkpointer.entities import OpeyCheckpointEntity, OpeyCheckpointWriteEntity
+
+class OBPCheckpointSaver(BaseCheckpointSaver):
+    """Checkpoint saver using OBP Dynamic Entitites as a storage backend."""
+    
+    def __init__(self, client: OBPClient):
+        super().__init__()
+        self.client = client
+        self.is_setup = False
+        self.serde = JsonPlusSerializer()
+        
+    def setup(self) -> None:
+        """
+        Setup the Dynamic Entities in OBP if they do not exist.
+        Needs to use an admin OBP client to create the system level entities.
+        """
+        if self.is_setup:
+            return
+        
+    def put(
+        self,
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: CheckpointMetadata,
+        new_versions: ChannelVersions,
+    ) -> RunnableConfig:
+        """Store a checkpoint on OBP."""

test/checkpointer/test_entities.py

@@ -0,0 +1,29 @@
+from src.checkpointer.entities import OpeyCheckpointEntity, OpeyCheckpointWriteEntity
+
+
+def test_opey_checkpoint_entity_schema():
+    schema = OpeyCheckpointEntity.to_obp_schema()
+    
+    assert schema["hasPersonalEntity"] is True
+    assert schema["OpeyCheckpoint"]["description"] == "OBP Dynamic Entity representation of a LangGraph checkpoint."
+    assert set(schema["OpeyCheckpoint"]["required"]) == {
+        "thread_id", "checkpoint_id", "checkpoint_ns", 
+        "parent_checkpoint_id", "checkpoint_data", "metadata"
+    }
+    assert "created_at" in schema["OpeyCheckpoint"]["properties"]
+    assert schema["OpeyCheckpoint"]["properties"]["checkpoint_data"]["description"] == "Serialized checkpoint TypedDict"
+    
+
+def test_opey_checkpoint_write_entity_schema():
+    schema = OpeyCheckpointWriteEntity.to_obp_schema()
+    
+    assert schema["hasPersonalEntity"] is True
+    assert schema["OpeyCheckpointWrite"]["description"] == "OBP Dynamic Entity representation of a pending LangGraph checkpoint write."
+    assert set(schema["OpeyCheckpointWrite"]["required"]) == {
+        "thread_id", "checkpoint_id", "checkpoint_ns", 
+        "task_id", "idx", "channel", "value"
+    }
+    assert all(
+        prop in schema["OpeyCheckpointWrite"]["properties"] 
+        for prop in ["thread_id", "checkpoint_id", "task_id", "idx", "channel", "value"]
+    )

test/conftest.py

@@ -12,15 +12,15 @@
 from httpx import AsyncClient, ASGITransport
 from asgi_lifespan import LifespanManager
 
-@pytest_asyncio.fixture(scope="session", loop_scope="session")
-def anyio_backend():
-    return "asyncio"
-
-@pytest_asyncio.fixture(scope="session", loop_scope="session")
-async def client():
-    async with AsyncClient(transport=ASGITransport(app=app), base_url="http://localhost:5000") as client:
-        print("Client is ready")
-        yield client
+# @pytest_asyncio.fixture(scope="session", loop_scope="session")
+# def anyio_backend():
+#     return "asyncio"
+
+# @pytest_asyncio.fixture(scope="session", loop_scope="session")
+# async def client():
+#     async with AsyncClient(transport=ASGITransport(app=app), base_url="http://localhost:5000") as client:
+#         print("Client is ready")
+#         yield client
 
 @pytest.fixture(scope="session")
 def get_obp_consent():