From 679f4b306e47c6e610f097a2fdd9214ca4fef163 Mon Sep 17 00:00:00 2001 From: Charles Pierse Date: Mon, 9 Mar 2026 15:45:56 +0000 Subject: [PATCH 1/2] Standardise conversation input with OpenAI format --- src/engram/__init__.py | 6 ++-- src/engram/_models/__init__.py | 6 ++-- src/engram/_models/memory.py | 27 ++++++++++++---- src/engram/_serialization/_builders.py | 18 ++++++++--- tests/test_client_async.py | 14 +++++--- tests/test_client_sync.py | 14 +++++--- tests/test_imports.py | 9 ++++-- tests/test_serialization.py | 44 +++++++++++++++++++++++--- 8 files changed, 107 insertions(+), 31 deletions(-) diff --git a/src/engram/__init__.py b/src/engram/__init__.py index d05fdf0..7e4ed6a 100644 --- a/src/engram/__init__.py +++ b/src/engram/__init__.py @@ -10,7 +10,8 @@ RunStatus, SearchResults, StringContent, - ToolCallMetadata, + ToolCallFuncInput, + ToolCallInput, ) from .async_client import AsyncEngramClient from .client import EngramClient @@ -43,7 +44,8 @@ "RunStatus", "SearchResults", "StringContent", - "ToolCallMetadata", + "ToolCallFuncInput", + "ToolCallInput", "ValidationError", "__version__", ] diff --git a/src/engram/_models/__init__.py b/src/engram/_models/__init__.py index 594e196..b239181 100644 --- a/src/engram/_models/__init__.py +++ b/src/engram/_models/__init__.py @@ -7,7 +7,8 @@ RetrievalConfig, SearchResults, StringContent, - ToolCallMetadata, + ToolCallFuncInput, + ToolCallInput, ) from .run import CommittedOperation, CommittedOperations, Run, RunStatus @@ -24,5 +25,6 @@ "RunStatus", "SearchResults", "StringContent", - "ToolCallMetadata", + "ToolCallFuncInput", + "ToolCallInput", ] diff --git a/src/engram/_models/memory.py b/src/engram/_models/memory.py index 775865a..8378211 100644 --- a/src/engram/_models/memory.py +++ b/src/engram/_models/memory.py @@ -21,21 +21,36 @@ class StringContent: @dataclass(slots=True) -class ToolCallMetadata: - """Tool call metadata.""" +class ToolCallFuncInput: + """The function details of an OpenAI-format tool call.""" name: str + arguments: str + + +@dataclass(slots=True) +class ToolCallInput: + """A single tool call in OpenAI Chat Completions format.""" + id: str + function: ToolCallFuncInput + type: str = "function" @dataclass(slots=True) class MessageContent: - """A message in a conversation.""" + """A message in a conversation using the OpenAI Chat Completions format. - role: Literal["user", "assistant", "system"] - content: str + - 'tool' role (tool results) is mapped to 'user' by the server. + - 'developer' role is mapped to 'system' by the server. + """ + + role: Literal["user", "assistant", "system", "tool", "developer"] + content: str = "" created_at: str | None = None - tool_call_metadata: ToolCallMetadata | None = None + tool_call_id: str | None = None + name: str | None = None + tool_calls: list[ToolCallInput] | None = None @dataclass(slots=True) diff --git a/src/engram/_serialization/_builders.py b/src/engram/_serialization/_builders.py index ce80559..ac6cae5 100644 --- a/src/engram/_serialization/_builders.py +++ b/src/engram/_serialization/_builders.py @@ -39,11 +39,19 @@ def _serialize_conversation_content(content: ConversationContent) -> dict[str, A m: dict[str, Any] = {"role": msg.role, "content": msg.content} if msg.created_at is not None: m["created_at"] = msg.created_at - if msg.tool_call_metadata is not None: - m["tool_call_metadata"] = { - "name": msg.tool_call_metadata.name, - "id": msg.tool_call_metadata.id, - } + if msg.tool_call_id is not None: + m["tool_call_id"] = msg.tool_call_id + if msg.name is not None: + m["name"] = msg.name + if msg.tool_calls is not None: + m["tool_calls"] = [ + { + "id": tc.id, + "type": tc.type, + "function": {"name": tc.function.name, "arguments": tc.function.arguments}, + } + for tc in msg.tool_calls + ] messages.append(m) conversation: dict[str, Any] = {"messages": messages} if content.metadata is not None: diff --git a/tests/test_client_async.py b/tests/test_client_async.py index 2228431..bbc036b 100644 --- a/tests/test_client_async.py +++ b/tests/test_client_async.py @@ -11,7 +11,8 @@ PreExtractedContent, RetrievalConfig, StringContent, - ToolCallMetadata, + ToolCallFuncInput, + ToolCallInput, ) from engram.async_client import DEFAULT_BASE_URL, AsyncEngramClient from engram.errors import APIError, AuthenticationError, ValidationError @@ -202,8 +203,11 @@ def handler(request: httpx.Request) -> httpx.Response: MessageContent(role="user", content="hi"), MessageContent( role="assistant", - content="using tool", - tool_call_metadata=ToolCallMetadata(name="search", id="tc1"), + tool_calls=[ + ToolCallInput( + id="tc1", function=ToolCallFuncInput(name="search", arguments="{}") + ) + ], ), ], metadata={"session_id": "s1"}, @@ -214,7 +218,9 @@ def handler(request: httpx.Request) -> httpx.Response: assert body["content"]["type"] == "conversation" conv = body["content"]["conversation"] assert conv["metadata"] == {"session_id": "s1"} - assert conv["messages"][1]["tool_call_metadata"] == {"name": "search", "id": "tc1"} + assert conv["messages"][1]["tool_calls"] == [ + {"id": "tc1", "type": "function", "function": {"name": "search", "arguments": "{}"}} + ] assert body["conversation_id"] == "c1" diff --git a/tests/test_client_sync.py b/tests/test_client_sync.py index 98b4c97..fea7666 100644 --- a/tests/test_client_sync.py +++ b/tests/test_client_sync.py @@ -11,7 +11,8 @@ PreExtractedContent, RetrievalConfig, StringContent, - ToolCallMetadata, + ToolCallFuncInput, + ToolCallInput, ) from engram.client import DEFAULT_BASE_URL, EngramClient from engram.errors import APIError, AuthenticationError, ValidationError @@ -217,8 +218,11 @@ def handler(request: httpx.Request) -> httpx.Response: MessageContent(role="user", content="hi"), MessageContent( role="assistant", - content="using tool", - tool_call_metadata=ToolCallMetadata(name="search", id="tc1"), + tool_calls=[ + ToolCallInput( + id="tc1", function=ToolCallFuncInput(name="search", arguments="{}") + ) + ], ), ], metadata={"session_id": "s1"}, @@ -229,7 +233,9 @@ def handler(request: httpx.Request) -> httpx.Response: assert body["content"]["type"] == "conversation" conv = body["content"]["conversation"] assert conv["metadata"] == {"session_id": "s1"} - assert conv["messages"][1]["tool_call_metadata"] == {"name": "search", "id": "tc1"} + assert conv["messages"][1]["tool_calls"] == [ + {"id": "tc1", "type": "function", "function": {"name": "search", "arguments": "{}"}} + ] assert body["conversation_id"] == "c1" diff --git a/tests/test_imports.py b/tests/test_imports.py index b802205..ed79fd0 100644 --- a/tests/test_imports.py +++ b/tests/test_imports.py @@ -19,7 +19,8 @@ def test_public_imports() -> None: RunStatus, SearchResults, StringContent, - ToolCallMetadata, + ToolCallFuncInput, + ToolCallInput, ValidationError, ) @@ -41,7 +42,8 @@ def test_public_imports() -> None: assert isinstance(ConversationContent, type) assert isinstance(MessageContent, type) assert isinstance(StringContent, type) - assert isinstance(ToolCallMetadata, type) + assert isinstance(ToolCallFuncInput, type) + assert isinstance(ToolCallInput, type) expected_exports = { "APIError", @@ -62,7 +64,8 @@ def test_public_imports() -> None: "RunStatus", "SearchResults", "StringContent", - "ToolCallMetadata", + "ToolCallFuncInput", + "ToolCallInput", "ValidationError", "__version__", } diff --git a/tests/test_serialization.py b/tests/test_serialization.py index 8b4fa2f..97d19d2 100644 --- a/tests/test_serialization.py +++ b/tests/test_serialization.py @@ -4,7 +4,8 @@ PreExtractedContent, RetrievalConfig, StringContent, - ToolCallMetadata, + ToolCallFuncInput, + ToolCallInput, ) from engram._serialization import ( build_add_body, @@ -160,12 +161,15 @@ def test_build_add_body_conversation_content_with_message_timestamps() -> None: assert "tool_call_metadata" not in msg -def test_build_add_body_conversation_content_with_tool_call_metadata() -> None: +def test_build_add_body_conversation_content_with_tool_calls() -> None: messages = [ MessageContent( role="assistant", - content="using tool", - tool_call_metadata=ToolCallMetadata(name="search", id="tc1"), + tool_calls=[ + ToolCallInput( + id="tc1", function=ToolCallFuncInput(name="search", arguments='{"q":"x"}') + ) + ], ) ] body = build_add_body( @@ -175,7 +179,37 @@ def test_build_add_body_conversation_content_with_tool_call_metadata() -> None: group=None, ) msg = body["content"]["conversation"]["messages"][0] - assert msg["tool_call_metadata"] == {"name": "search", "id": "tc1"} + assert msg["tool_calls"] == [ + {"id": "tc1", "type": "function", "function": {"name": "search", "arguments": '{"q":"x"}'}} + ] + + +def test_build_add_body_conversation_content_with_tool_role() -> None: + messages = [MessageContent(role="tool", content="result", tool_call_id="tc1", name="search")] + body = build_add_body( + ConversationContent(messages=messages), + user_id=None, + conversation_id=None, + group=None, + ) + msg = body["content"]["conversation"]["messages"][0] + assert msg["role"] == "tool" + assert msg["tool_call_id"] == "tc1" + assert msg["name"] == "search" + assert msg["content"] == "result" + + +def test_build_add_body_conversation_content_with_developer_role() -> None: + messages = [MessageContent(role="developer", content="You are a helpful assistant.")] + body = build_add_body( + ConversationContent(messages=messages), + user_id=None, + conversation_id=None, + group=None, + ) + msg = body["content"]["conversation"]["messages"][0] + assert msg["role"] == "developer" + assert msg["content"] == "You are a helpful assistant." # ── build_memory_params ───────────────────────────────────────────────── From 215e643c743e910af66b5f2c9ad89130c05083a8 Mon Sep 17 00:00:00 2001 From: Charles Pierse Date: Mon, 9 Mar 2026 16:48:26 +0000 Subject: [PATCH 2/2] align with backend changes --- src/engram/__init__.py | 2 ++ src/engram/_models/__init__.py | 2 ++ src/engram/_models/memory.py | 18 +++++++++++++++--- src/engram/_serialization/_builders.py | 19 +++++++++++-------- tests/test_imports.py | 3 +++ tests/test_serialization.py | 26 ++++++++++++++++++++++++++ 6 files changed, 59 insertions(+), 11 deletions(-) diff --git a/src/engram/__init__.py b/src/engram/__init__.py index 7e4ed6a..5e6eee2 100644 --- a/src/engram/__init__.py +++ b/src/engram/__init__.py @@ -10,6 +10,7 @@ RunStatus, SearchResults, StringContent, + ToolCallCustomInput, ToolCallFuncInput, ToolCallInput, ) @@ -44,6 +45,7 @@ "RunStatus", "SearchResults", "StringContent", + "ToolCallCustomInput", "ToolCallFuncInput", "ToolCallInput", "ValidationError", diff --git a/src/engram/_models/__init__.py b/src/engram/_models/__init__.py index b239181..06656aa 100644 --- a/src/engram/_models/__init__.py +++ b/src/engram/_models/__init__.py @@ -7,6 +7,7 @@ RetrievalConfig, SearchResults, StringContent, + ToolCallCustomInput, ToolCallFuncInput, ToolCallInput, ) @@ -25,6 +26,7 @@ "RunStatus", "SearchResults", "StringContent", + "ToolCallCustomInput", "ToolCallFuncInput", "ToolCallInput", ] diff --git a/src/engram/_models/memory.py b/src/engram/_models/memory.py index 8378211..6153f1a 100644 --- a/src/engram/_models/memory.py +++ b/src/engram/_models/memory.py @@ -22,19 +22,31 @@ class StringContent: @dataclass(slots=True) class ToolCallFuncInput: - """The function details of an OpenAI-format tool call.""" + """The function details of an OpenAI-format function tool call.""" name: str arguments: str +@dataclass(slots=True) +class ToolCallCustomInput: + """The details of an OpenAI-format custom tool call.""" + + name: str + input: str + + @dataclass(slots=True) class ToolCallInput: - """A single tool call in OpenAI Chat Completions format.""" + """A single tool call in OpenAI Chat Completions format. + + Set either `function` or `custom` depending on the tool type. + """ id: str - function: ToolCallFuncInput type: str = "function" + function: ToolCallFuncInput | None = None + custom: ToolCallCustomInput | None = None @dataclass(slots=True) diff --git a/src/engram/_serialization/_builders.py b/src/engram/_serialization/_builders.py index ac6cae5..d0fa161 100644 --- a/src/engram/_serialization/_builders.py +++ b/src/engram/_serialization/_builders.py @@ -8,9 +8,19 @@ PreExtractedContent, RetrievalConfig, StringContent, + ToolCallInput, ) +def _serialize_tool_call(tc: ToolCallInput) -> dict[str, Any]: + out: dict[str, Any] = {"id": tc.id, "type": tc.type} + if tc.function is not None: + out["function"] = {"name": tc.function.name, "arguments": tc.function.arguments} + if tc.custom is not None: + out["custom"] = {"name": tc.custom.name, "input": tc.custom.input} + return out + + def _serialize_content(content: AddContent) -> dict[str, Any]: """Build the content envelope with the type discriminator.""" if isinstance(content, str): @@ -44,14 +54,7 @@ def _serialize_conversation_content(content: ConversationContent) -> dict[str, A if msg.name is not None: m["name"] = msg.name if msg.tool_calls is not None: - m["tool_calls"] = [ - { - "id": tc.id, - "type": tc.type, - "function": {"name": tc.function.name, "arguments": tc.function.arguments}, - } - for tc in msg.tool_calls - ] + m["tool_calls"] = [_serialize_tool_call(tc) for tc in msg.tool_calls] messages.append(m) conversation: dict[str, Any] = {"messages": messages} if content.metadata is not None: diff --git a/tests/test_imports.py b/tests/test_imports.py index ed79fd0..0a8e8e4 100644 --- a/tests/test_imports.py +++ b/tests/test_imports.py @@ -19,6 +19,7 @@ def test_public_imports() -> None: RunStatus, SearchResults, StringContent, + ToolCallCustomInput, ToolCallFuncInput, ToolCallInput, ValidationError, @@ -42,6 +43,7 @@ def test_public_imports() -> None: assert isinstance(ConversationContent, type) assert isinstance(MessageContent, type) assert isinstance(StringContent, type) + assert isinstance(ToolCallCustomInput, type) assert isinstance(ToolCallFuncInput, type) assert isinstance(ToolCallInput, type) @@ -64,6 +66,7 @@ def test_public_imports() -> None: "RunStatus", "SearchResults", "StringContent", + "ToolCallCustomInput", "ToolCallFuncInput", "ToolCallInput", "ValidationError", diff --git a/tests/test_serialization.py b/tests/test_serialization.py index 97d19d2..1b36ddb 100644 --- a/tests/test_serialization.py +++ b/tests/test_serialization.py @@ -4,6 +4,7 @@ PreExtractedContent, RetrievalConfig, StringContent, + ToolCallCustomInput, ToolCallFuncInput, ToolCallInput, ) @@ -184,6 +185,31 @@ 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( + role="assistant", + tool_calls=[ + ToolCallInput( + id="tc2", + type="custom", + custom=ToolCallCustomInput(name="my_tool", input="some input"), + ) + ], + ) + ] + body = build_add_body( + ConversationContent(messages=messages), + user_id=None, + conversation_id=None, + group=None, + ) + msg = body["content"]["conversation"]["messages"][0] + assert msg["tool_calls"] == [ + {"id": "tc2", "type": "custom", "custom": {"name": "my_tool", "input": "some input"}} + ] + + def test_build_add_body_conversation_content_with_tool_role() -> None: messages = [MessageContent(role="tool", content="result", tool_call_id="tc1", name="search")] body = build_add_body(