session metadata (#48)

* connection timeout, better url concatentation, improve docs, improve session_id checks

* add /get session

* add Session to examples

Author: danielchalef

examples/async.py

@@ -17,6 +17,7 @@
     NotFoundError,
     ZepClient,
 )
+from zep_python.models import Session
 
 
 async def main() -> None:
@@ -25,6 +26,39 @@ async def main() -> None:
     async with ZepClient(base_url, api_key) as client:
         session_id = uuid.uuid4().hex
 
+        #
+        # Create session
+        #
+        print(f"Creating session: {session_id}")
+        try:
+            session = Session(session_id=session_id, metadata={"foo": "bar"})
+            result = await client.aadd_session(session)
+            print(result)
+        except APIError as e:
+            print(f"Unable to create session {session_id} got error: {e}")
+
+        #
+        # Update session metadata
+        #
+        print(f"Creating session: {session_id}")
+        try:
+            # The new metadata values will be merged with the existing metadata
+            session = Session(session_id=session_id, metadata={"bar": "foo"})
+            result = await client.aadd_session(session)
+            print(result)
+        except APIError as e:
+            print(f"Unable to create session {session_id} got error: {e}")
+
+        #
+        # Get session
+        #
+        print(f"Getting session: {session_id}")
+        try:
+            session = await client.aget_session(session_id)
+            print(session.dict())
+        except NotFoundError:
+            print("Session not found")
+
         print(f"\n1---getMemory for Session: {session_id}")
         try:
             memory = await client.aget_memory(session_id)

examples/sync.py

@@ -16,6 +16,7 @@
     NotFoundError,
     ZepClient,
 )
+from zep_python.models import Session
 
 
 def main() -> None:
@@ -25,6 +26,39 @@ def main() -> None:
         # Example usage
         session_id = uuid.uuid4().hex
 
+        #
+        # Create session
+        #
+        print(f"Creating session: {session_id}")
+        try:
+            session = Session(session_id=session_id, metadata={"foo": "bar"})
+            result = client.add_session(session)
+            print(result)
+        except APIError as e:
+            print(f"Unable to create session {session_id} got error: {e}")
+
+        #
+        # Update session metadata
+        #
+        print(f"Creating session: {session_id}")
+        try:
+            # The new metadata values will be merged with the existing metadata
+            session = Session(session_id=session_id, metadata={"bar": "foo"})
+            result = client.add_session(session)
+            print(result)
+        except APIError as e:
+            print(f"Unable to create session {session_id} got error: {e}")
+
+        #
+        # Get session
+        #
+        print(f"Getting session: {session_id}")
+        try:
+            session = client.get_session(session_id)
+            print(session.dict())
+        except NotFoundError:
+            print("Session not found")
+
         #
         # Get memory
         #

tests/client_test.py

@@ -6,12 +6,12 @@
 from pytest_httpx import HTTPXMock
 
 from zep_python import NotFoundError
-from zep_python.models import Memory, MemorySearchPayload, Message
+from zep_python.models import Memory, MemorySearchPayload, Message, Session
 from zep_python.zep_client import ZepClient
 
 api_base_url = "http://localhost/api/v1"
 
-mock_response = {
+mock_messages = {
     "messages": [
         {
             "uuid": "msg-uuid",
@@ -60,15 +60,15 @@ def filter_unset_fields(d: Dict) -> Dict:
 
 
 def validate_memory(memory: Memory):
-    assert len(memory.messages) == len(mock_response["messages"])
+    assert len(memory.messages) == len(mock_messages["messages"])
 
     for i in range(len(memory.messages)):
-        assert memory.messages[i].uuid == mock_response["messages"][i]["uuid"]
-        assert memory.messages[i].role == mock_response["messages"][i]["role"]
-        assert memory.messages[i].content == mock_response["messages"][i]["content"]
-        assert memory.messages[i].metadata == mock_response["messages"][i]["metadata"]
+        assert memory.messages[i].uuid == mock_messages["messages"][i]["uuid"]
+        assert memory.messages[i].role == mock_messages["messages"][i]["role"]
+        assert memory.messages[i].content == mock_messages["messages"][i]["content"]
+        assert memory.messages[i].metadata == mock_messages["messages"][i]["metadata"]
 
-    assert filter_unset_fields(memory.dict()) == mock_response
+    assert filter_unset_fields(memory.dict()) == mock_messages
 
 
 @pytest.mark.asyncio
@@ -87,7 +87,7 @@ async def test_set_authorization_header(httpx_mock):
 async def test_aget_memory(httpx_mock: HTTPXMock):
     session_id = str(uuid4())
 
-    httpx_mock.add_response(status_code=200, json=mock_response)
+    httpx_mock.add_response(status_code=200, json=mock_messages)
 
     async with ZepClient(base_url=api_base_url) as client:
         memory = await client.aget_memory(session_id)
@@ -132,7 +132,7 @@ async def test_aget_memory_missing_values(httpx_mock: HTTPXMock):
 def test_get_memory(httpx_mock: HTTPXMock):
     session_id = str(uuid4())
 
-    httpx_mock.add_response(status_code=200, json=mock_response)
+    httpx_mock.add_response(status_code=200, json=mock_messages)
 
     with ZepClient(base_url=api_base_url) as client:
         memory = client.get_memory(session_id)
@@ -267,3 +267,94 @@ def test_search_memory(httpx_mock: HTTPXMock):
         assert search_results[0].message["role"] == "user"
         assert search_results[0].message["content"] == "Test message"
         assert search_results[0].dist == 0.9
+
+
+# Predefined session response
+mock_session = {
+    "uuid": str(uuid4()),
+    "created_at": "2020-12-31T23:59:59",
+    "updated_at": "2021-01-01T00:00:00",
+    "deleted_at": None,
+    "session_id": "abc123",
+    "metadata": {},
+}
+
+
+def validate_session(session):
+    # Validate the session object here
+    assert Session.parse_obj(mock_session) == session
+
+
+@pytest.mark.asyncio
+async def test_get_session(httpx_mock: HTTPXMock):
+    session_id = "abc123"
+
+    httpx_mock.add_response(status_code=200, json=mock_session)
+
+    client = ZepClient(base_url=api_base_url)
+    session = client.get_session(session_id)
+
+    validate_session(session)
+
+
+def test_get_session_missing_id():
+    client = ZepClient(base_url=api_base_url)
+
+    with pytest.raises(ValueError):
+        client.get_session(session_id=None)  # type: ignore
+
+
+@pytest.mark.asyncio
+async def test_aget_session(httpx_mock: HTTPXMock):
+    session_id = "abc123"
+
+    httpx_mock.add_response(status_code=200, json=mock_session)
+
+    async with ZepClient(base_url=api_base_url) as client:
+        session = await client.aget_session(session_id)
+
+        validate_session(session)
+
+
+@pytest.mark.asyncio
+async def test_aget_session_missing_id(httpx_mock: HTTPXMock):
+    async with ZepClient(base_url=api_base_url) as client:
+        with pytest.raises(ValueError):
+            _ = await client.aget_session(session_id=None)  # type: ignore
+
+
+def test_add_session(httpx_mock: HTTPXMock):
+    session = Session(**mock_session)
+
+    httpx_mock.add_response(status_code=200, text="OK")
+
+    client = ZepClient(base_url=api_base_url)
+    result = client.add_session(session)
+
+    assert result == "OK"
+
+
+def test_add_session_missing_session():
+    client = ZepClient(base_url=api_base_url)
+
+    with pytest.raises(ValueError):
+        client.add_session(session=None)  # type: ignore
+
+
+@pytest.mark.asyncio
+async def test_aadd_session(httpx_mock: HTTPXMock):
+    session = Session(**mock_session)
+
+    httpx_mock.add_response(status_code=200, text="OK")
+
+    async with ZepClient(base_url=api_base_url) as client:
+        result = await client.aadd_session(session)
+
+        assert result == "OK"
+
+
+@pytest.mark.asyncio
+async def test_aadd_session_missing_session(httpx_mock: HTTPXMock):
+    async with ZepClient(base_url=api_base_url) as client:
+        with pytest.raises(ValueError):
+            _ = await client.aadd_session(session=None)  # type: ignore

zep_python/models.py

@@ -5,6 +5,39 @@
 from pydantic import BaseModel, Field
 
 
+class Session(BaseModel):
+    """
+    Represents a session object with a unique identifier, metadata,
+    and other attributes.
+
+    Attributes
+    ----------
+    uuid : Optional[str]
+        A unique identifier for the session.
+        This is generated server-side and is not expected to be present on creation.
+    created_at : str
+        The timestamp when the session was created.
+        Generated by the server.
+    updated_at : str
+        The timestamp when the session was last updated.
+        Generated by the server.
+    deleted_at : Optional[datetime]
+        The timestamp when the session was deleted.
+        Generated by the server.
+    session_id : str
+        The unique identifier of the session.
+    metadata : Dict[str, Any]
+        The metadata associated with the session.
+    """
+
+    uuid: Optional[str]
+    created_at: Optional[str]
+    updated_at: Optional[str]
+    deleted_at: Optional[str]
+    session_id: str
+    metadata: Dict[str, Any]
+
+
 class Summary(BaseModel):
     """
     Represents a summary of a conversation.