From 1504681710c7c3e9f93a7ef80bf98621e7e5e5c6 Mon Sep 17 00:00:00 2001 From: Dan Jones Date: Wed, 18 Mar 2026 13:02:02 +0000 Subject: [PATCH 1/4] Update API format and support batch inputs --- src/engram/__init__.py | 2 ++ src/engram/_models/__init__.py | 2 ++ src/engram/_models/memory.py | 15 ++++++++++++--- src/engram/_serialization/_builders.py | 21 ++++++++++----------- 4 files changed, 26 insertions(+), 14 deletions(-) diff --git a/src/engram/__init__.py b/src/engram/__init__.py index 5e6eee2..354a895 100644 --- a/src/engram/__init__.py +++ b/src/engram/__init__.py @@ -5,6 +5,7 @@ Memory, MessageContent, PreExtractedContent, + PreExtractedItem, RetrievalConfig, Run, RunStatus, @@ -40,6 +41,7 @@ "Memory", "MessageContent", "PreExtractedContent", + "PreExtractedItem", "RetrievalConfig", "Run", "RunStatus", diff --git a/src/engram/_models/__init__.py b/src/engram/_models/__init__.py index 06656aa..7799517 100644 --- a/src/engram/_models/__init__.py +++ b/src/engram/_models/__init__.py @@ -4,6 +4,7 @@ Memory, MessageContent, PreExtractedContent, + PreExtractedItem, RetrievalConfig, SearchResults, StringContent, @@ -21,6 +22,7 @@ "Memory", "MessageContent", "PreExtractedContent", + "PreExtractedItem", "RetrievalConfig", "Run", "RunStatus", diff --git a/src/engram/_models/memory.py b/src/engram/_models/memory.py index 6b21f58..13f3969 100644 --- a/src/engram/_models/memory.py +++ b/src/engram/_models/memory.py @@ -7,7 +7,16 @@ @dataclass(slots=True) class PreExtractedContent: - """Pre-extracted content that bypasses the extraction pipeline.""" + """Pre-extracted content that skips the extraction step continues through the pipeline as-is. + Each individual item represents a separate memory. + """ + + items: list[PreExtractedItem] + + +@dataclass(slots=True) +class PreExtractedItem: + """A single pre-extracted memory.""" content: str topic: str @@ -15,9 +24,9 @@ class PreExtractedContent: @dataclass(slots=True) class StringContent: - """String content that bypasses the extraction pipeline.""" + """String content to extract memories from.""" - content: str + content: str | list[str] @dataclass(slots=True) diff --git a/src/engram/_serialization/_builders.py b/src/engram/_serialization/_builders.py index abcd0c5..ef98c99 100644 --- a/src/engram/_serialization/_builders.py +++ b/src/engram/_serialization/_builders.py @@ -21,21 +21,20 @@ def _serialize_tool_call(tc: ToolCallInput) -> dict[str, Any]: return out -def _serialize_content(content: AddContent) -> dict[str, Any]: +def _serialize_input(content: AddContent) -> dict[str, Any]: """Build the content envelope with the type discriminator.""" if isinstance(content, str): - return {"type": "string", "content": content} + return {"string": {"content": [content]}} if isinstance(content, StringContent): - return {"type": "string", "content": content.content} + if isinstance(content.content, list): + return {"string": {"content": content.content}} + else: + return {"string": {"content": [content.content]}} if isinstance(content, PreExtractedContent): - return { - "type": "pre_extracted", - "content": content.content, - "topic": content.topic, - } + items = [{"content": item.content, "topic": item.topic} for item in content.items] + return {"pre_extracted": {"items": items}} if isinstance(content, list): return { - "type": "conversation", "conversation": {"messages": content}, } if isinstance(content, ConversationContent): @@ -63,7 +62,7 @@ def _serialize_conversation_content(content: ConversationContent) -> dict[str, A conversation["created_at"] = content.created_at if content.updated_at is not None: conversation["updated_at"] = content.updated_at - return {"type": "conversation", "conversation": conversation} + return {"conversation": conversation} def build_add_body( @@ -73,7 +72,7 @@ def build_add_body( conversation_id: str | None, group: str | None, ) -> dict[str, Any]: - body: dict[str, Any] = {"content": _serialize_content(content)} + body: dict[str, Any] = {"input": _serialize_input(content)} if user_id is not None: body["user_id"] = user_id if conversation_id is not None: From e54a7da3d2a5c430a3a5aaba9738ad38fafb5631 Mon Sep 17 00:00:00 2001 From: Dan Jones Date: Wed, 18 Mar 2026 13:02:21 +0000 Subject: [PATCH 2/4] Update tests --- tests/test_client_async.py | 59 ++++++++++++++++++++++++++++++++--- tests/test_client_sync.py | 62 +++++++++++++++++++++++++++++++------ tests/test_imports.py | 3 ++ tests/test_serialization.py | 33 +++++++++----------- 4 files changed, 125 insertions(+), 32 deletions(-) diff --git a/tests/test_client_async.py b/tests/test_client_async.py index 6942072..8834766 100644 --- a/tests/test_client_async.py +++ b/tests/test_client_async.py @@ -9,6 +9,7 @@ ConversationContent, MessageContent, PreExtractedContent, + PreExtractedItem, RetrievalConfig, StringContent, ToolCallFuncInput, @@ -117,7 +118,7 @@ async def test_add_str() -> None: async def test_add_pre_extracted() -> None: client = _make_client(body={"run_id": "r2", "status": "pending"}) result = await client.memories.add( - PreExtractedContent(content="fact", topic="topic"), + PreExtractedContent(items=[PreExtractedItem(content="fact", topic="topic")]), user_id="u1", ) assert result.run_id == "r2" @@ -146,12 +147,61 @@ def handler(request: httpx.Request) -> httpx.Response: await client.memories.add("hello", user_id="u1", group="g1") body = json.loads(captured[0].content) assert body == { - "content": {"type": "string", "content": "hello"}, + "input": {"string": {"content": ["hello"]}}, "user_id": "u1", "group": "g1", } +@pytest.mark.asyncio +async def test_add_multiple_strings() -> None: + captured: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured.append(request) + return httpx.Response(200, json={"run_id": "r1", "status": "pending"}) + + client = _make_client_with_handler(handler) + await client.memories.add(StringContent(content=["first", "second"]), user_id="u1") + body = json.loads(captured[0].content) + assert body == { + "input": {"string": {"content": ["first", "second"]}}, + "user_id": "u1", + } + + +@pytest.mark.asyncio +async def test_add_multiple_pre_extracted_items() -> None: + captured: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured.append(request) + return httpx.Response(200, json={"run_id": "r1", "status": "pending"}) + + client = _make_client_with_handler(handler) + await client.memories.add( + PreExtractedContent( + items=[ + PreExtractedItem(content="fact one", topic="topic_a"), + PreExtractedItem(content="fact two", topic="topic_b"), + ] + ), + user_id="u1", + ) + body = json.loads(captured[0].content) + assert body == { + "input": { + "pre_extracted": { + "items": [ + {"content": "fact one", "topic": "topic_a"}, + {"content": "fact two", "topic": "topic_b"}, + ] + } + }, + "user_id": "u1", + } + + @pytest.mark.asyncio async def test_add_string_content() -> None: client = _make_client(body={"run_id": "r4", "status": "pending"}) @@ -171,7 +221,7 @@ def handler(request: httpx.Request) -> httpx.Response: await client.memories.add(StringContent(content="hello"), user_id="u1", group="g1") body = json.loads(captured[0].content) assert body == { - "content": {"type": "string", "content": "hello"}, + "input": {"string": {"content": ["hello"]}}, "user_id": "u1", "group": "g1", } @@ -215,8 +265,7 @@ def handler(request: httpx.Request) -> httpx.Response: conversation_id="c1", ) body = json.loads(captured[0].content) - assert body["content"]["type"] == "conversation" - conv = body["content"]["conversation"] + conv = body["input"]["conversation"] assert conv["metadata"] == {"session_id": "s1"} assert conv["messages"][1]["tool_calls"] == [ {"id": "tc1", "type": "function", "function": {"name": "search", "arguments": "{}"}} diff --git a/tests/test_client_sync.py b/tests/test_client_sync.py index fe6586b..c8bc790 100644 --- a/tests/test_client_sync.py +++ b/tests/test_client_sync.py @@ -9,6 +9,7 @@ ConversationContent, MessageContent, PreExtractedContent, + PreExtractedItem, RetrievalConfig, StringContent, ToolCallFuncInput, @@ -118,7 +119,7 @@ def test_add_str() -> None: def test_add_pre_extracted() -> None: client = _make_client(body={"run_id": "r2", "status": "pending"}) result = client.memories.add( - PreExtractedContent(content="fact", topic="topic"), + PreExtractedContent(items=[PreExtractedItem(content="fact", topic="topic")]), user_id="u1", ) assert result.run_id == "r2" @@ -145,7 +146,7 @@ def handler(request: httpx.Request) -> httpx.Response: client.memories.add("hello", user_id="u1", group="g1") body = json.loads(captured[0].content) assert body == { - "content": {"type": "string", "content": "hello"}, + "input": {"string": {"content": ["hello"]}}, "user_id": "u1", "group": "g1", } @@ -163,14 +164,58 @@ def handler(request: httpx.Request) -> httpx.Response: client.memories.add(messages, conversation_id="c1") body = json.loads(captured[0].content) assert body == { - "content": { - "type": "conversation", - "conversation": {"messages": messages}, - }, + "input": {"conversation": {"messages": messages}}, "conversation_id": "c1", } +def test_add_multiple_strings() -> None: + captured: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured.append(request) + return httpx.Response(200, json={"run_id": "r1", "status": "pending"}) + + client = _make_client_with_handler(handler) + client.memories.add(StringContent(content=["first", "second"]), user_id="u1") + body = json.loads(captured[0].content) + assert body == { + "input": {"string": {"content": ["first", "second"]}}, + "user_id": "u1", + } + + +def test_add_multiple_pre_extracted_items() -> None: + captured: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured.append(request) + return httpx.Response(200, json={"run_id": "r1", "status": "pending"}) + + client = _make_client_with_handler(handler) + client.memories.add( + PreExtractedContent( + items=[ + PreExtractedItem(content="fact one", topic="topic_a"), + PreExtractedItem(content="fact two", topic="topic_b"), + ] + ), + user_id="u1", + ) + body = json.loads(captured[0].content) + assert body == { + "input": { + "pre_extracted": { + "items": [ + {"content": "fact one", "topic": "topic_a"}, + {"content": "fact two", "topic": "topic_b"}, + ] + } + }, + "user_id": "u1", + } + + def test_add_string_content() -> None: client = _make_client(body={"run_id": "r4", "status": "pending"}) result = client.memories.add(StringContent(content="hello"), user_id="u1") @@ -188,7 +233,7 @@ def handler(request: httpx.Request) -> httpx.Response: client.memories.add(StringContent(content="hello"), user_id="u1", group="g1") body = json.loads(captured[0].content) assert body == { - "content": {"type": "string", "content": "hello"}, + "input": {"string": {"content": ["hello"]}}, "user_id": "u1", "group": "g1", } @@ -230,8 +275,7 @@ def handler(request: httpx.Request) -> httpx.Response: conversation_id="c1", ) body = json.loads(captured[0].content) - assert body["content"]["type"] == "conversation" - conv = body["content"]["conversation"] + conv = body["input"]["conversation"] assert conv["metadata"] == {"session_id": "s1"} assert conv["messages"][1]["tool_calls"] == [ {"id": "tc1", "type": "function", "function": {"name": "search", "arguments": "{}"}} diff --git a/tests/test_imports.py b/tests/test_imports.py index 0a8e8e4..8bd559a 100644 --- a/tests/test_imports.py +++ b/tests/test_imports.py @@ -14,6 +14,7 @@ def test_public_imports() -> None: Memory, MessageContent, PreExtractedContent, + PreExtractedItem, RetrievalConfig, Run, RunStatus, @@ -37,6 +38,7 @@ def test_public_imports() -> None: assert isinstance(RunStatus, type) assert isinstance(SearchResults, type) assert isinstance(PreExtractedContent, type) + assert isinstance(PreExtractedItem, type) assert isinstance(RetrievalConfig, type) assert isinstance(CommittedOperation, type) assert isinstance(CommittedOperations, type) @@ -61,6 +63,7 @@ def test_public_imports() -> None: "Memory", "MessageContent", "PreExtractedContent", + "PreExtractedItem", "RetrievalConfig", "Run", "RunStatus", diff --git a/tests/test_serialization.py b/tests/test_serialization.py index 3fe1ab2..19d1fa4 100644 --- a/tests/test_serialization.py +++ b/tests/test_serialization.py @@ -2,6 +2,7 @@ ConversationContent, MessageContent, PreExtractedContent, + PreExtractedItem, RetrievalConfig, StringContent, ToolCallCustomInput, @@ -28,7 +29,7 @@ def test_build_add_body_str() -> None: conversation_id=None, group=None, ) - assert body == {"content": {"type": "string", "content": "hello world"}} + assert body == {"input": {"string": {"content": ["hello world"]}}} def test_build_add_body_str_with_options() -> None: @@ -39,7 +40,7 @@ def test_build_add_body_str_with_options() -> None: group="g1", ) assert body == { - "content": {"type": "string", "content": "hello"}, + "input": {"string": {"content": ["hello"]}}, "user_id": "u1", "conversation_id": "c1", "group": "g1", @@ -48,13 +49,13 @@ def test_build_add_body_str_with_options() -> None: def test_build_add_body_pre_extracted() -> None: body = build_add_body( - PreExtractedContent(content="fact", topic="topic"), + PreExtractedContent(items=[PreExtractedItem(content="fact", topic="topic")]), user_id=None, conversation_id=None, group=None, ) assert body == { - "content": {"type": "pre_extracted", "content": "fact", "topic": "topic"}, + "input": {"pre_extracted": {"items": [{"content": "fact", "topic": "topic"}]}}, } @@ -70,10 +71,7 @@ def test_build_add_body_conversation() -> None: group=None, ) assert body == { - "content": { - "type": "conversation", - "conversation": {"messages": messages}, - }, + "input": {"conversation": {"messages": messages}}, "user_id": "u1", "conversation_id": "c1", } @@ -86,7 +84,7 @@ def test_build_add_body_string_content() -> None: conversation_id=None, group=None, ) - assert body == {"content": {"type": "string", "content": "hello world"}} + assert body == {"input": {"string": {"content": ["hello world"]}}} def test_build_add_body_string_content_with_options() -> None: @@ -97,7 +95,7 @@ def test_build_add_body_string_content_with_options() -> None: group="g1", ) assert body == { - "content": {"type": "string", "content": "hello"}, + "input": {"string": {"content": ["hello"]}}, "user_id": "u1", "conversation_id": "c1", "group": "g1", @@ -116,8 +114,7 @@ def test_build_add_body_conversation_content() -> None: group=None, ) assert body == { - "content": { - "type": "conversation", + "input": { "conversation": { "messages": [ {"role": "user", "content": "hi"}, @@ -143,7 +140,7 @@ def test_build_add_body_conversation_content_with_metadata() -> None: conversation_id=None, group=None, ) - conv = body["content"]["conversation"] + conv = body["input"]["conversation"] assert conv["metadata"] == {"session_id": "s1"} assert conv["created_at"] == "2024-01-01T00:00:00Z" assert conv["updated_at"] == "2024-01-02T00:00:00Z" @@ -157,7 +154,7 @@ def test_build_add_body_conversation_content_with_message_timestamps() -> None: conversation_id=None, group=None, ) - msg = body["content"]["conversation"]["messages"][0] + msg = body["input"]["conversation"]["messages"][0] assert msg["created_at"] == "2024-01-01T00:00:00Z" assert "tool_call_metadata" not in msg @@ -179,7 +176,7 @@ def test_build_add_body_conversation_content_with_tool_calls() -> None: conversation_id=None, group=None, ) - msg = body["content"]["conversation"]["messages"][0] + msg = body["input"]["conversation"]["messages"][0] assert msg["tool_calls"] == [ {"id": "tc1", "type": "function", "function": {"name": "search", "arguments": '{"q":"x"}'}} ] @@ -204,7 +201,7 @@ def test_build_add_body_conversation_content_with_custom_tool_calls() -> None: conversation_id=None, group=None, ) - msg = body["content"]["conversation"]["messages"][0] + msg = body["input"]["conversation"]["messages"][0] assert msg["tool_calls"] == [ {"id": "tc2", "type": "custom", "custom": {"name": "my_tool", "input": "some input"}} ] @@ -218,7 +215,7 @@ def test_build_add_body_conversation_content_with_tool_role() -> None: conversation_id=None, group=None, ) - msg = body["content"]["conversation"]["messages"][0] + msg = body["input"]["conversation"]["messages"][0] assert msg["role"] == "tool" assert msg["tool_call_id"] == "tc1" assert msg["name"] == "search" @@ -233,7 +230,7 @@ def test_build_add_body_conversation_content_with_developer_role() -> None: conversation_id=None, group=None, ) - msg = body["content"]["conversation"]["messages"][0] + msg = body["input"]["conversation"]["messages"][0] assert msg["role"] == "developer" assert msg["content"] == "You are a helpful assistant." From 51054b22f058a69fe249424f8ca2ddd7c5f00a04 Mon Sep 17 00:00:00 2001 From: Dan Jones Date: Thu, 19 Mar 2026 12:33:04 +0000 Subject: [PATCH 3/4] Rename content to input throughout --- src/engram/__init__.py | 16 +++++----- src/engram/_models/__init__.py | 20 ++++++------ src/engram/_models/memory.py | 20 ++++++------ src/engram/_resources/memories.py | 10 +++--- src/engram/_serialization/_builders.py | 44 +++++++++++++------------- tests/test_client_async.py | 26 +++++++-------- tests/test_client_sync.py | 26 +++++++-------- tests/test_imports.py | 24 +++++++------- tests/test_serialization.py | 44 +++++++++++++------------- 9 files changed, 115 insertions(+), 115 deletions(-) diff --git a/src/engram/__init__.py b/src/engram/__init__.py index 354a895..9fbb4b5 100644 --- a/src/engram/__init__.py +++ b/src/engram/__init__.py @@ -1,16 +1,16 @@ from ._models import ( CommittedOperation, CommittedOperations, - ConversationContent, + ConversationInput, Memory, - MessageContent, - PreExtractedContent, + MessageInput, + PreExtractedInput, PreExtractedItem, RetrievalConfig, Run, RunStatus, SearchResults, - StringContent, + StringInput, ToolCallCustomInput, ToolCallFuncInput, ToolCallInput, @@ -34,19 +34,19 @@ "CommittedOperation", "CommittedOperations", "ConnectionError", - "ConversationContent", + "ConversationInput", "EngramClient", "EngramError", "EngramTimeoutError", "Memory", - "MessageContent", - "PreExtractedContent", + "MessageInput", + "PreExtractedInput", "PreExtractedItem", "RetrievalConfig", "Run", "RunStatus", "SearchResults", - "StringContent", + "StringInput", "ToolCallCustomInput", "ToolCallFuncInput", "ToolCallInput", diff --git a/src/engram/_models/__init__.py b/src/engram/_models/__init__.py index 7799517..b9906e2 100644 --- a/src/engram/_models/__init__.py +++ b/src/engram/_models/__init__.py @@ -1,13 +1,13 @@ from .memory import ( - AddContent, - ConversationContent, + AddInput, + ConversationInput, Memory, - MessageContent, - PreExtractedContent, + MessageInput, + PreExtractedInput, PreExtractedItem, RetrievalConfig, SearchResults, - StringContent, + StringInput, ToolCallCustomInput, ToolCallFuncInput, ToolCallInput, @@ -15,19 +15,19 @@ from .run import CommittedOperation, CommittedOperations, Run, RunStatus __all__ = [ - "AddContent", + "AddInput", "CommittedOperation", "CommittedOperations", - "ConversationContent", + "ConversationInput", "Memory", - "MessageContent", - "PreExtractedContent", + "MessageInput", + "PreExtractedInput", "PreExtractedItem", "RetrievalConfig", "Run", "RunStatus", "SearchResults", - "StringContent", + "StringInput", "ToolCallCustomInput", "ToolCallFuncInput", "ToolCallInput", diff --git a/src/engram/_models/memory.py b/src/engram/_models/memory.py index 13f3969..2c824b3 100644 --- a/src/engram/_models/memory.py +++ b/src/engram/_models/memory.py @@ -6,8 +6,8 @@ @dataclass(slots=True) -class PreExtractedContent: - """Pre-extracted content that skips the extraction step continues through the pipeline as-is. +class PreExtractedInput: + """Pre-extracted input that skips the extraction step continues through the pipeline as-is. Each individual item represents a separate memory. """ @@ -23,8 +23,8 @@ class PreExtractedItem: @dataclass(slots=True) -class StringContent: - """String content to extract memories from.""" +class StringInput: + """String input to extract memories from.""" content: str | list[str] @@ -59,7 +59,7 @@ class ToolCallInput: @dataclass(slots=True) -class MessageContent: +class MessageInput: """A message in a conversation using the OpenAI Chat Completions format. - 'tool' role (tool results) is mapped to 'user' by the server. @@ -75,18 +75,18 @@ class MessageContent: @dataclass(slots=True) -class ConversationContent: - """Conversation content that bypasses the extraction pipeline.""" +class ConversationInput: + """Conversation input that bypasses the extraction pipeline.""" - messages: list[MessageContent] + messages: list[MessageInput] metadata: dict[str, Any] | None = None created_at: str | None = None updated_at: str | None = None # Type alias for the content argument to memories.add() -AddContent: TypeAlias = ( - str | list[dict[str, str]] | PreExtractedContent | ConversationContent | StringContent +AddInput: TypeAlias = ( + str | list[dict[str, str]] | PreExtractedInput | ConversationInput | StringInput ) diff --git a/src/engram/_resources/memories.py b/src/engram/_resources/memories.py index fa0bcba..1b22627 100644 --- a/src/engram/_resources/memories.py +++ b/src/engram/_resources/memories.py @@ -3,7 +3,7 @@ from uuid import UUID from .._http import AsyncHttpTransport, HttpTransport -from .._models import AddContent, Memory, RetrievalConfig, Run, SearchResults +from .._models import AddInput, Memory, RetrievalConfig, Run, SearchResults from .._serialization import ( build_add_body, build_memory_params, @@ -29,14 +29,14 @@ def __init__(self, transport: HttpTransport) -> None: def add( self, - content: AddContent, + input_data: AddInput, *, user_id: str | None = None, conversation_id: str | None = None, group: str | None = None, ) -> Run: body = build_add_body( - content, + input_data, user_id=user_id, conversation_id=conversation_id, group=group, @@ -101,14 +101,14 @@ def __init__(self, transport: AsyncHttpTransport) -> None: async def add( self, - content: AddContent, + input_data: AddInput, *, user_id: str | None = None, conversation_id: str | None = None, group: str | None = None, ) -> Run: body = build_add_body( - content, + input_data, user_id=user_id, conversation_id=conversation_id, group=group, diff --git a/src/engram/_serialization/_builders.py b/src/engram/_serialization/_builders.py index ef98c99..7a975db 100644 --- a/src/engram/_serialization/_builders.py +++ b/src/engram/_serialization/_builders.py @@ -3,11 +3,11 @@ from typing import Any from .._models import ( - AddContent, - ConversationContent, - PreExtractedContent, + AddInput, + ConversationInput, + PreExtractedInput, RetrievalConfig, - StringContent, + StringInput, ToolCallInput, ) @@ -21,28 +21,28 @@ def _serialize_tool_call(tc: ToolCallInput) -> dict[str, Any]: return out -def _serialize_input(content: AddContent) -> dict[str, Any]: - """Build the content envelope with the type discriminator.""" - if isinstance(content, str): - return {"string": {"content": [content]}} - if isinstance(content, StringContent): - if isinstance(content.content, list): - return {"string": {"content": content.content}} +def _serialize_input(input_data: AddInput) -> dict[str, Any]: + """Build the input envelope with the type discriminator.""" + if isinstance(input_data, str): + return {"string": {"content": [input_data]}} + if isinstance(input_data, StringInput): + if isinstance(input_data.content, list): + return {"string": {"content": input_data.content}} else: - return {"string": {"content": [content.content]}} - if isinstance(content, PreExtractedContent): - items = [{"content": item.content, "topic": item.topic} for item in content.items] + return {"string": {"content": [input_data.content]}} + if isinstance(input_data, PreExtractedInput): + items = [{"content": item.content, "topic": item.topic} for item in input_data.items] return {"pre_extracted": {"items": items}} - if isinstance(content, list): + if isinstance(input_data, list): return { - "conversation": {"messages": content}, + "conversation": {"messages": input_data}, } - if isinstance(content, ConversationContent): - return _serialize_conversation_content(content) - raise TypeError(f"Unsupported content type: {type(content)}") # pragma: no cover + if isinstance(input_data, ConversationInput): + return _serialize_conversation_content(input_data) + raise TypeError(f"Unsupported input type: {type(input_data)}") # pragma: no cover -def _serialize_conversation_content(content: ConversationContent) -> dict[str, Any]: +def _serialize_conversation_content(content: ConversationInput) -> dict[str, Any]: messages = [] for msg in content.messages: m: dict[str, Any] = {"role": msg.role, "content": msg.content} @@ -66,13 +66,13 @@ def _serialize_conversation_content(content: ConversationContent) -> dict[str, A def build_add_body( - content: AddContent, + input_data: AddInput, *, user_id: str | None, conversation_id: str | None, group: str | None, ) -> dict[str, Any]: - body: dict[str, Any] = {"input": _serialize_input(content)} + body: dict[str, Any] = {"input": _serialize_input(input_data)} if user_id is not None: body["user_id"] = user_id if conversation_id is not None: diff --git a/tests/test_client_async.py b/tests/test_client_async.py index 8834766..64d35f8 100644 --- a/tests/test_client_async.py +++ b/tests/test_client_async.py @@ -6,12 +6,12 @@ from engram._http import AsyncHttpTransport from engram._models import ( - ConversationContent, - MessageContent, - PreExtractedContent, + ConversationInput, + MessageInput, + PreExtractedInput, PreExtractedItem, RetrievalConfig, - StringContent, + StringInput, ToolCallFuncInput, ToolCallInput, ) @@ -118,7 +118,7 @@ async def test_add_str() -> None: async def test_add_pre_extracted() -> None: client = _make_client(body={"run_id": "r2", "status": "pending"}) result = await client.memories.add( - PreExtractedContent(items=[PreExtractedItem(content="fact", topic="topic")]), + PreExtractedInput(items=[PreExtractedItem(content="fact", topic="topic")]), user_id="u1", ) assert result.run_id == "r2" @@ -162,7 +162,7 @@ def handler(request: httpx.Request) -> httpx.Response: return httpx.Response(200, json={"run_id": "r1", "status": "pending"}) client = _make_client_with_handler(handler) - await client.memories.add(StringContent(content=["first", "second"]), user_id="u1") + await client.memories.add(StringInput(content=["first", "second"]), user_id="u1") body = json.loads(captured[0].content) assert body == { "input": {"string": {"content": ["first", "second"]}}, @@ -180,7 +180,7 @@ def handler(request: httpx.Request) -> httpx.Response: client = _make_client_with_handler(handler) await client.memories.add( - PreExtractedContent( + PreExtractedInput( items=[ PreExtractedItem(content="fact one", topic="topic_a"), PreExtractedItem(content="fact two", topic="topic_b"), @@ -205,7 +205,7 @@ def handler(request: httpx.Request) -> httpx.Response: @pytest.mark.asyncio async def test_add_string_content() -> None: client = _make_client(body={"run_id": "r4", "status": "pending"}) - result = await client.memories.add(StringContent(content="hello"), user_id="u1") + result = await client.memories.add(StringInput(content="hello"), user_id="u1") assert result.run_id == "r4" @@ -218,7 +218,7 @@ def handler(request: httpx.Request) -> httpx.Response: return httpx.Response(200, json={"run_id": "r1", "status": "pending"}) client = _make_client_with_handler(handler) - await client.memories.add(StringContent(content="hello"), user_id="u1", group="g1") + await client.memories.add(StringInput(content="hello"), user_id="u1", group="g1") body = json.loads(captured[0].content) assert body == { "input": {"string": {"content": ["hello"]}}, @@ -231,7 +231,7 @@ def handler(request: httpx.Request) -> httpx.Response: async def test_add_conversation_content() -> None: client = _make_client(body={"run_id": "r5", "status": "pending"}) result = await client.memories.add( - ConversationContent(messages=[MessageContent(role="user", content="hi")]), + ConversationInput(messages=[MessageInput(role="user", content="hi")]), user_id="u1", conversation_id="c1", ) @@ -248,10 +248,10 @@ def handler(request: httpx.Request) -> httpx.Response: client = _make_client_with_handler(handler) await client.memories.add( - ConversationContent( + ConversationInput( messages=[ - MessageContent(role="user", content="hi"), - MessageContent( + MessageInput(role="user", content="hi"), + MessageInput( role="assistant", tool_calls=[ ToolCallInput( diff --git a/tests/test_client_sync.py b/tests/test_client_sync.py index c8bc790..5e666c9 100644 --- a/tests/test_client_sync.py +++ b/tests/test_client_sync.py @@ -6,12 +6,12 @@ from engram._http import HttpTransport from engram._models import ( - ConversationContent, - MessageContent, - PreExtractedContent, + ConversationInput, + MessageInput, + PreExtractedInput, PreExtractedItem, RetrievalConfig, - StringContent, + StringInput, ToolCallFuncInput, ToolCallInput, ) @@ -119,7 +119,7 @@ def test_add_str() -> None: def test_add_pre_extracted() -> None: client = _make_client(body={"run_id": "r2", "status": "pending"}) result = client.memories.add( - PreExtractedContent(items=[PreExtractedItem(content="fact", topic="topic")]), + PreExtractedInput(items=[PreExtractedItem(content="fact", topic="topic")]), user_id="u1", ) assert result.run_id == "r2" @@ -177,7 +177,7 @@ def handler(request: httpx.Request) -> httpx.Response: return httpx.Response(200, json={"run_id": "r1", "status": "pending"}) client = _make_client_with_handler(handler) - client.memories.add(StringContent(content=["first", "second"]), user_id="u1") + client.memories.add(StringInput(content=["first", "second"]), user_id="u1") body = json.loads(captured[0].content) assert body == { "input": {"string": {"content": ["first", "second"]}}, @@ -194,7 +194,7 @@ def handler(request: httpx.Request) -> httpx.Response: client = _make_client_with_handler(handler) client.memories.add( - PreExtractedContent( + PreExtractedInput( items=[ PreExtractedItem(content="fact one", topic="topic_a"), PreExtractedItem(content="fact two", topic="topic_b"), @@ -218,7 +218,7 @@ def handler(request: httpx.Request) -> httpx.Response: def test_add_string_content() -> None: client = _make_client(body={"run_id": "r4", "status": "pending"}) - result = client.memories.add(StringContent(content="hello"), user_id="u1") + result = client.memories.add(StringInput(content="hello"), user_id="u1") assert result.run_id == "r4" @@ -230,7 +230,7 @@ def handler(request: httpx.Request) -> httpx.Response: return httpx.Response(200, json={"run_id": "r1", "status": "pending"}) client = _make_client_with_handler(handler) - client.memories.add(StringContent(content="hello"), user_id="u1", group="g1") + client.memories.add(StringInput(content="hello"), user_id="u1", group="g1") body = json.loads(captured[0].content) assert body == { "input": {"string": {"content": ["hello"]}}, @@ -242,7 +242,7 @@ def handler(request: httpx.Request) -> httpx.Response: def test_add_conversation_content() -> None: client = _make_client(body={"run_id": "r5", "status": "pending"}) result = client.memories.add( - ConversationContent(messages=[MessageContent(role="user", content="hi")]), + ConversationInput(messages=[MessageInput(role="user", content="hi")]), user_id="u1", conversation_id="c1", ) @@ -258,10 +258,10 @@ def handler(request: httpx.Request) -> httpx.Response: client = _make_client_with_handler(handler) client.memories.add( - ConversationContent( + ConversationInput( messages=[ - MessageContent(role="user", content="hi"), - MessageContent( + MessageInput(role="user", content="hi"), + MessageInput( role="assistant", tool_calls=[ ToolCallInput( diff --git a/tests/test_imports.py b/tests/test_imports.py index 8bd559a..fb29e1f 100644 --- a/tests/test_imports.py +++ b/tests/test_imports.py @@ -7,19 +7,19 @@ def test_public_imports() -> None: CommittedOperation, CommittedOperations, ConnectionError, - ConversationContent, + ConversationInput, EngramClient, EngramError, EngramTimeoutError, Memory, - MessageContent, - PreExtractedContent, + MessageInput, + PreExtractedInput, PreExtractedItem, RetrievalConfig, Run, RunStatus, SearchResults, - StringContent, + StringInput, ToolCallCustomInput, ToolCallFuncInput, ToolCallInput, @@ -37,14 +37,14 @@ def test_public_imports() -> None: assert isinstance(Run, type) assert isinstance(RunStatus, type) assert isinstance(SearchResults, type) - assert isinstance(PreExtractedContent, type) + assert isinstance(PreExtractedInput, type) assert isinstance(PreExtractedItem, type) assert isinstance(RetrievalConfig, type) assert isinstance(CommittedOperation, type) assert isinstance(CommittedOperations, type) - assert isinstance(ConversationContent, type) - assert isinstance(MessageContent, type) - assert isinstance(StringContent, type) + assert isinstance(ConversationInput, type) + assert isinstance(MessageInput, type) + assert isinstance(StringInput, type) assert isinstance(ToolCallCustomInput, type) assert isinstance(ToolCallFuncInput, type) assert isinstance(ToolCallInput, type) @@ -56,19 +56,19 @@ def test_public_imports() -> None: "CommittedOperation", "CommittedOperations", "ConnectionError", - "ConversationContent", + "ConversationInput", "EngramClient", "EngramError", "EngramTimeoutError", "Memory", - "MessageContent", - "PreExtractedContent", + "MessageInput", + "PreExtractedInput", "PreExtractedItem", "RetrievalConfig", "Run", "RunStatus", "SearchResults", - "StringContent", + "StringInput", "ToolCallCustomInput", "ToolCallFuncInput", "ToolCallInput", diff --git a/tests/test_serialization.py b/tests/test_serialization.py index 19d1fa4..6328d81 100644 --- a/tests/test_serialization.py +++ b/tests/test_serialization.py @@ -1,10 +1,10 @@ from engram._models import ( - ConversationContent, - MessageContent, - PreExtractedContent, + ConversationInput, + MessageInput, + PreExtractedInput, PreExtractedItem, RetrievalConfig, - StringContent, + StringInput, ToolCallCustomInput, ToolCallFuncInput, ToolCallInput, @@ -49,7 +49,7 @@ def test_build_add_body_str_with_options() -> None: def test_build_add_body_pre_extracted() -> None: body = build_add_body( - PreExtractedContent(items=[PreExtractedItem(content="fact", topic="topic")]), + PreExtractedInput(items=[PreExtractedItem(content="fact", topic="topic")]), user_id=None, conversation_id=None, group=None, @@ -79,7 +79,7 @@ def test_build_add_body_conversation() -> None: def test_build_add_body_string_content() -> None: body = build_add_body( - StringContent(content="hello world"), + StringInput(content="hello world"), user_id=None, conversation_id=None, group=None, @@ -89,7 +89,7 @@ def test_build_add_body_string_content() -> None: def test_build_add_body_string_content_with_options() -> None: body = build_add_body( - StringContent(content="hello"), + StringInput(content="hello"), user_id="u1", conversation_id="c1", group="g1", @@ -104,11 +104,11 @@ def test_build_add_body_string_content_with_options() -> None: def test_build_add_body_conversation_content() -> None: messages = [ - MessageContent(role="user", content="hi"), - MessageContent(role="assistant", content="hello"), + MessageInput(role="user", content="hi"), + MessageInput(role="assistant", content="hello"), ] body = build_add_body( - ConversationContent(messages=messages), + ConversationInput(messages=messages), user_id="u1", conversation_id="c1", group=None, @@ -128,9 +128,9 @@ def test_build_add_body_conversation_content() -> None: def test_build_add_body_conversation_content_with_metadata() -> None: - messages = [MessageContent(role="user", content="hi")] + messages = [MessageInput(role="user", content="hi")] body = build_add_body( - ConversationContent( + ConversationInput( messages=messages, metadata={"session_id": "s1"}, created_at="2024-01-01T00:00:00Z", @@ -147,9 +147,9 @@ def test_build_add_body_conversation_content_with_metadata() -> None: def test_build_add_body_conversation_content_with_message_timestamps() -> None: - messages = [MessageContent(role="user", content="hi", created_at="2024-01-01T00:00:00Z")] + messages = [MessageInput(role="user", content="hi", created_at="2024-01-01T00:00:00Z")] body = build_add_body( - ConversationContent(messages=messages), + ConversationInput(messages=messages), user_id=None, conversation_id=None, group=None, @@ -161,7 +161,7 @@ def test_build_add_body_conversation_content_with_message_timestamps() -> None: def test_build_add_body_conversation_content_with_tool_calls() -> None: messages = [ - MessageContent( + MessageInput( role="assistant", tool_calls=[ ToolCallInput( @@ -171,7 +171,7 @@ def test_build_add_body_conversation_content_with_tool_calls() -> None: ) ] body = build_add_body( - ConversationContent(messages=messages), + ConversationInput(messages=messages), user_id=None, conversation_id=None, group=None, @@ -184,7 +184,7 @@ def test_build_add_body_conversation_content_with_tool_calls() -> None: def test_build_add_body_conversation_content_with_custom_tool_calls() -> None: messages = [ - MessageContent( + MessageInput( role="assistant", tool_calls=[ ToolCallInput( @@ -196,7 +196,7 @@ def test_build_add_body_conversation_content_with_custom_tool_calls() -> None: ) ] body = build_add_body( - ConversationContent(messages=messages), + ConversationInput(messages=messages), user_id=None, conversation_id=None, group=None, @@ -208,9 +208,9 @@ def test_build_add_body_conversation_content_with_custom_tool_calls() -> None: def test_build_add_body_conversation_content_with_tool_role() -> None: - messages = [MessageContent(role="tool", content="result", tool_call_id="tc1", name="search")] + messages = [MessageInput(role="tool", content="result", tool_call_id="tc1", name="search")] body = build_add_body( - ConversationContent(messages=messages), + ConversationInput(messages=messages), user_id=None, conversation_id=None, group=None, @@ -223,9 +223,9 @@ def test_build_add_body_conversation_content_with_tool_role() -> None: def test_build_add_body_conversation_content_with_developer_role() -> None: - messages = [MessageContent(role="developer", content="You are a helpful assistant.")] + messages = [MessageInput(role="developer", content="You are a helpful assistant.")] body = build_add_body( - ConversationContent(messages=messages), + ConversationInput(messages=messages), user_id=None, conversation_id=None, group=None, From 943a7b48946ed0abbe159e05fb17cb0b3c6f3cbf Mon Sep 17 00:00:00 2001 From: Dan Jones Date: Thu, 19 Mar 2026 12:35:46 +0000 Subject: [PATCH 4/4] Fix comment --- src/engram/_models/memory.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/engram/_models/memory.py b/src/engram/_models/memory.py index 2c824b3..17ed194 100644 --- a/src/engram/_models/memory.py +++ b/src/engram/_models/memory.py @@ -84,7 +84,7 @@ class ConversationInput: updated_at: str | None = None -# Type alias for the content argument to memories.add() +# Type alias for the input_data argument to memories.add() AddInput: TypeAlias = ( str | list[dict[str, str]] | PreExtractedInput | ConversationInput | StringInput )