From 295a75db5e057006681888222052d30fb1beb34b Mon Sep 17 00:00:00 2001 From: Augustas Date: Mon, 4 May 2026 15:35:21 +0300 Subject: [PATCH 1/3] add properties --- src/engram/__init__.py | 2 + src/engram/_models/__init__.py | 4 ++ src/engram/_models/memory.py | 16 ++++++ src/engram/_resources/memories.py | 24 +++++++- src/engram/_serialization/_builders.py | 32 ++++++++++- src/engram/_serialization/_parsers.py | 1 + tests/test_client_async.py | 40 +++++++++++++ tests/test_client_sync.py | 57 +++++++++++++++++++ tests/test_imports.py | 3 + tests/test_serialization.py | 79 ++++++++++++++++++++++++++ 10 files changed, 252 insertions(+), 6 deletions(-) diff --git a/src/engram/__init__.py b/src/engram/__init__.py index 9fbb4b5..6e06efb 100644 --- a/src/engram/__init__.py +++ b/src/engram/__init__.py @@ -14,6 +14,7 @@ ToolCallCustomInput, ToolCallFuncInput, ToolCallInput, + Topic, ) from .async_client import AsyncEngramClient from .client import EngramClient @@ -50,6 +51,7 @@ "ToolCallCustomInput", "ToolCallFuncInput", "ToolCallInput", + "Topic", "ValidationError", "__version__", ] diff --git a/src/engram/_models/__init__.py b/src/engram/_models/__init__.py index b9906e2..04cbfc1 100644 --- a/src/engram/_models/__init__.py +++ b/src/engram/_models/__init__.py @@ -11,6 +11,8 @@ ToolCallCustomInput, ToolCallFuncInput, ToolCallInput, + Topic, + TopicSelector, ) from .run import CommittedOperation, CommittedOperations, Run, RunStatus @@ -31,4 +33,6 @@ "ToolCallCustomInput", "ToolCallFuncInput", "ToolCallInput", + "Topic", + "TopicSelector", ] diff --git a/src/engram/_models/memory.py b/src/engram/_models/memory.py index 17ed194..43357d0 100644 --- a/src/engram/_models/memory.py +++ b/src/engram/_models/memory.py @@ -96,6 +96,21 @@ class RetrievalConfig: limit: int | None = None +@dataclass(slots=True) +class Topic: + """A topic with an optional per-topic property filter. + + Use ``None`` as a property value to clear an inherited global filter + for this topic only. + """ + + name: str + properties: dict[str, str | None] | None = None + + +TopicSelector: TypeAlias = str | Topic + + @dataclass(slots=True) class Memory: id: str @@ -109,6 +124,7 @@ class Memory: conversation_id: str | None = None tags: list[str] | None = None score: float | None = None + properties: dict[str, str] | None = None class SearchResults(Sequence[Memory]): diff --git a/src/engram/_resources/memories.py b/src/engram/_resources/memories.py index 1b22627..0d93f36 100644 --- a/src/engram/_resources/memories.py +++ b/src/engram/_resources/memories.py @@ -1,9 +1,17 @@ from __future__ import annotations +from typing import TypeAlias from uuid import UUID from .._http import AsyncHttpTransport, HttpTransport -from .._models import AddInput, Memory, RetrievalConfig, Run, SearchResults +from .._models import ( + AddInput, + Memory, + RetrievalConfig, + Run, + SearchResults, + TopicSelector, +) from .._serialization import ( build_add_body, build_memory_params, @@ -16,6 +24,8 @@ _MEMORIES_PATH = "/v1/memories" _MEMORIES_SEARCH_PATH = "/v1/memories/search" +_Topics: TypeAlias = list[TopicSelector] | None + def _memory_path(memory_id: str | UUID) -> str: return f"{_MEMORIES_PATH}/{memory_id}" @@ -34,12 +44,14 @@ def add( user_id: str | None = None, conversation_id: str | None = None, group: str | None = None, + properties: dict[str, str] | None = None, ) -> Run: body = build_add_body( input_data, user_id=user_id, conversation_id=conversation_id, group=group, + properties=properties, ) data = self._transport.request("POST", _MEMORIES_PATH, json=body) return parse_run(data) @@ -75,11 +87,12 @@ def search( self, *, query: str, - topics: list[str] | None = None, + topics: _Topics = None, user_id: str | None = None, conversation_id: str | None = None, group: str | None = None, retrieval_config: RetrievalConfig | None = None, + properties: dict[str, str] | None = None, ) -> SearchResults: body = build_search_body( query=query, @@ -88,6 +101,7 @@ def search( conversation_id=conversation_id, group=group, retrieval_config=retrieval_config, + properties=properties, ) data = self._transport.request("POST", _MEMORIES_SEARCH_PATH, json=body) return parse_search_results(data) @@ -106,12 +120,14 @@ async def add( user_id: str | None = None, conversation_id: str | None = None, group: str | None = None, + properties: dict[str, str] | None = None, ) -> Run: body = build_add_body( input_data, user_id=user_id, conversation_id=conversation_id, group=group, + properties=properties, ) data = await self._transport.request("POST", _MEMORIES_PATH, json=body) return parse_run(data) @@ -147,11 +163,12 @@ async def search( self, *, query: str, - topics: list[str] | None = None, + topics: _Topics = None, user_id: str | None = None, conversation_id: str | None = None, group: str | None = None, retrieval_config: RetrievalConfig | None = None, + properties: dict[str, str] | None = None, ) -> SearchResults: body = build_search_body( query=query, @@ -160,6 +177,7 @@ async def search( conversation_id=conversation_id, group=group, retrieval_config=retrieval_config, + properties=properties, ) data = await self._transport.request("POST", _MEMORIES_SEARCH_PATH, json=body) return parse_search_results(data) diff --git a/src/engram/_serialization/_builders.py b/src/engram/_serialization/_builders.py index 7a975db..c6f5a77 100644 --- a/src/engram/_serialization/_builders.py +++ b/src/engram/_serialization/_builders.py @@ -9,6 +9,8 @@ RetrievalConfig, StringInput, ToolCallInput, + Topic, + TopicSelector, ) @@ -65,12 +67,30 @@ def _serialize_conversation_content(content: ConversationInput) -> dict[str, Any return {"conversation": conversation} +def _serialize_topic(topic: TopicSelector) -> str | dict[str, Any]: + if isinstance(topic, str): + return topic + if isinstance(topic, Topic): + out: dict[str, Any] = {"name": topic.name} + if topic.properties is not None: + out["properties"] = dict(topic.properties) + return out + raise TypeError(f"Unsupported topic type: {type(topic)}") # pragma: no cover + + +def _serialize_topics(topics: list[TopicSelector] | None) -> list[str | dict[str, Any]] | None: + if topics is None: + return None + return [_serialize_topic(t) for t in topics] + + def build_add_body( input_data: AddInput, *, user_id: str | None, conversation_id: str | None, group: str | None, + properties: dict[str, str] | None = None, ) -> dict[str, Any]: body: dict[str, Any] = {"input": _serialize_input(input_data)} if user_id is not None: @@ -79,6 +99,8 @@ def build_add_body( body["conversation_id"] = conversation_id if group is not None: body["group"] = group + if properties is not None: + body["properties"] = dict(properties) return body @@ -98,11 +120,12 @@ def build_memory_params( def build_search_body( *, query: str, - topics: list[str] | None, + topics: list[TopicSelector] | None, user_id: str | None, conversation_id: str | None, group: str | None, retrieval_config: RetrievalConfig | None, + properties: dict[str, str] | None = None, ) -> dict[str, Any]: body: dict[str, Any] = {"query": query} if retrieval_config is not None: @@ -110,12 +133,15 @@ def build_search_body( "retrieval_type": retrieval_config.retrieval_type, "limit": retrieval_config.limit, } - if topics is not None: - body["topics"] = topics + serialized_topics = _serialize_topics(topics) + if serialized_topics is not None: + body["topics"] = serialized_topics if user_id is not None: body["user_id"] = user_id if conversation_id is not None: body["conversation_id"] = conversation_id if group is not None: body["group"] = group + if properties is not None: + body["properties"] = dict(properties) return body diff --git a/src/engram/_serialization/_parsers.py b/src/engram/_serialization/_parsers.py index 92a01fc..f998847 100644 --- a/src/engram/_serialization/_parsers.py +++ b/src/engram/_serialization/_parsers.py @@ -33,6 +33,7 @@ def parse_memory(data: dict[str, Any]) -> Memory: conversation_id=data.get("conversation_id"), tags=data.get("tags"), score=data.get("score"), + properties=data.get("properties"), ) diff --git a/tests/test_client_async.py b/tests/test_client_async.py index 64d35f8..3975240 100644 --- a/tests/test_client_async.py +++ b/tests/test_client_async.py @@ -14,6 +14,7 @@ StringInput, ToolCallFuncInput, ToolCallInput, + Topic, ) from engram.async_client import DEFAULT_BASE_URL, AsyncEngramClient from engram.errors import APIError, AuthenticationError, ValidationError @@ -371,6 +372,45 @@ def handler(request: httpx.Request) -> httpx.Response: assert body["retrieval_config"]["limit"] == 5 +# ── properties / list ─────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_add_sends_properties() -> 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( + "hello", + properties={"region": "eu"}, + ) + body = json.loads(captured[0].content) + assert body["properties"] == {"region": "eu"} + + +@pytest.mark.asyncio +async def test_search_sends_topic_filters() -> None: + captured: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured.append(request) + return httpx.Response(200, json={"memories": [], "total": 0}) + + client = _make_client_with_handler(handler) + await client.memories.search( + query="q", + topics=[Topic(name="t1", properties={"region": "eu"})], + properties={"tier": "pro"}, + ) + body = json.loads(captured[0].content) + assert body["topics"] == [{"name": "t1", "properties": {"region": "eu"}}] + assert body["properties"] == {"tier": "pro"} + + # ── runs.get ──────────────────────────────────────────────────────────── diff --git a/tests/test_client_sync.py b/tests/test_client_sync.py index 5e666c9..8365435 100644 --- a/tests/test_client_sync.py +++ b/tests/test_client_sync.py @@ -14,6 +14,7 @@ StringInput, ToolCallFuncInput, ToolCallInput, + Topic, ) from engram.client import DEFAULT_BASE_URL, EngramClient from engram.errors import APIError, AuthenticationError, ValidationError @@ -388,6 +389,62 @@ def handler(request: httpx.Request) -> httpx.Response: assert "retrieval_config" not in body +# ── properties support ────────────────────────────────────────────────── + + +def test_add_sends_properties() -> 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( + "hello", + user_id="u1", + properties={"region": "eu", "tier": "pro"}, + ) + body = json.loads(captured[0].content) + assert body["properties"] == {"region": "eu", "tier": "pro"} + + +def test_search_sends_properties_and_topic_filters() -> None: + captured: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured.append(request) + return httpx.Response(200, json={"memories": [], "total": 0}) + + client = _make_client_with_handler(handler) + client.memories.search( + query="q", + topics=[ + "plain", + Topic(name="scoped", properties={"region": "eu"}), + Topic(name="cleared", properties={"region": None}), + ], + properties={"tier": "pro"}, + ) + body = json.loads(captured[0].content) + assert body["properties"] == {"tier": "pro"} + assert body["topics"] == [ + "plain", + {"name": "scoped", "properties": {"region": "eu"}}, + {"name": "cleared", "properties": {"region": None}}, + ] + + +def test_get_memory_returns_properties() -> None: + response_body = { + **SAMPLE_MEMORY_RESPONSE, + "properties": {"region": "eu", "tier": "pro"}, + } + client = _make_client(body=response_body) + mem = client.memories.get("m1") + assert mem.properties == {"region": "eu", "tier": "pro"} + + # ── runs.get ──────────────────────────────────────────────────────────── diff --git a/tests/test_imports.py b/tests/test_imports.py index fb29e1f..3453f89 100644 --- a/tests/test_imports.py +++ b/tests/test_imports.py @@ -23,6 +23,7 @@ def test_public_imports() -> None: ToolCallCustomInput, ToolCallFuncInput, ToolCallInput, + Topic, ValidationError, ) @@ -48,6 +49,7 @@ def test_public_imports() -> None: assert isinstance(ToolCallCustomInput, type) assert isinstance(ToolCallFuncInput, type) assert isinstance(ToolCallInput, type) + assert isinstance(Topic, type) expected_exports = { "APIError", @@ -72,6 +74,7 @@ def test_public_imports() -> None: "ToolCallCustomInput", "ToolCallFuncInput", "ToolCallInput", + "Topic", "ValidationError", "__version__", } diff --git a/tests/test_serialization.py b/tests/test_serialization.py index 6328d81..d1fe991 100644 --- a/tests/test_serialization.py +++ b/tests/test_serialization.py @@ -8,6 +8,7 @@ ToolCallCustomInput, ToolCallFuncInput, ToolCallInput, + Topic, ) from engram._serialization import ( build_add_body, @@ -281,6 +282,82 @@ def test_build_search_body_full() -> None: assert body["retrieval_config"]["limit"] == 5 +# ── properties on add ─────────────────────────────────────────────────── + + +def test_build_add_body_with_properties() -> None: + body = build_add_body( + "hello", + user_id=None, + conversation_id=None, + group=None, + properties={"region": "eu", "tier": "pro"}, + ) + assert body == { + "input": {"string": {"content": ["hello"]}}, + "properties": {"region": "eu", "tier": "pro"}, + } + + +def test_build_add_body_properties_none_omitted() -> None: + body = build_add_body( + "hello", + user_id=None, + conversation_id=None, + group=None, + properties=None, + ) + assert "properties" not in body + + +# ── properties + topic filters on search ──────────────────────────────── + + +def test_build_search_body_with_properties() -> None: + body = build_search_body( + query="q", + topics=None, + user_id=None, + conversation_id=None, + group=None, + retrieval_config=None, + properties={"region": "eu"}, + ) + assert body == {"query": "q", "properties": {"region": "eu"}} + + +def test_build_search_body_with_topic_filter() -> None: + body = build_search_body( + query="q", + topics=[ + "plain", + Topic(name="scoped", properties={"region": "eu"}), + Topic(name="cleared", properties={"region": None}), + ], + user_id=None, + conversation_id=None, + group=None, + retrieval_config=None, + ) + assert body["topics"] == [ + "plain", + {"name": "scoped", "properties": {"region": "eu"}}, + {"name": "cleared", "properties": {"region": None}}, + ] + + +def test_build_search_body_topic_filter_without_properties() -> None: + body = build_search_body( + query="q", + topics=[Topic(name="t1")], + user_id=None, + conversation_id=None, + group=None, + retrieval_config=None, + ) + assert body["topics"] == [{"name": "t1"}] + + # ── parse_run ─────────────────────────────────────────────────────────── @@ -325,11 +402,13 @@ def test_parse_memory_with_optional_fields() -> None: "conversation_id": "c1", "tags": ["x"], "score": 0.95, + "properties": {"region": "eu", "tier": "pro"}, } mem = parse_memory(data) assert mem.user_id == "u1" assert mem.tags == ["x"] assert mem.score == 0.95 + assert mem.properties == {"region": "eu", "tier": "pro"} # ── parse_search_results ──────────────────────────────────────────────── From 40b6905fac0de79db5153c58d0c986efd9754f93 Mon Sep 17 00:00:00 2001 From: Augustas Date: Mon, 4 May 2026 15:46:04 +0300 Subject: [PATCH 2/3] remove conversation_id --- src/engram/_models/memory.py | 1 - src/engram/_resources/memories.py | 8 ------- src/engram/_serialization/_builders.py | 6 ------ src/engram/_serialization/_parsers.py | 1 - tests/test_client_async.py | 8 +++---- tests/test_client_sync.py | 12 +++++------ tests/test_serialization.py | 29 ++------------------------ 7 files changed, 12 insertions(+), 53 deletions(-) diff --git a/src/engram/_models/memory.py b/src/engram/_models/memory.py index 43357d0..6cd79d8 100644 --- a/src/engram/_models/memory.py +++ b/src/engram/_models/memory.py @@ -121,7 +121,6 @@ class Memory: created_at: str updated_at: str user_id: str | None = None - conversation_id: str | None = None tags: list[str] | None = None score: float | None = None properties: dict[str, str] | None = None diff --git a/src/engram/_resources/memories.py b/src/engram/_resources/memories.py index 0d93f36..4ab55e7 100644 --- a/src/engram/_resources/memories.py +++ b/src/engram/_resources/memories.py @@ -42,14 +42,12 @@ def add( input_data: AddInput, *, user_id: str | None = None, - conversation_id: str | None = None, group: str | None = None, properties: dict[str, str] | None = None, ) -> Run: body = build_add_body( input_data, user_id=user_id, - conversation_id=conversation_id, group=group, properties=properties, ) @@ -89,7 +87,6 @@ def search( query: str, topics: _Topics = None, user_id: str | None = None, - conversation_id: str | None = None, group: str | None = None, retrieval_config: RetrievalConfig | None = None, properties: dict[str, str] | None = None, @@ -98,7 +95,6 @@ def search( query=query, topics=topics, user_id=user_id, - conversation_id=conversation_id, group=group, retrieval_config=retrieval_config, properties=properties, @@ -118,14 +114,12 @@ async def add( input_data: AddInput, *, user_id: str | None = None, - conversation_id: str | None = None, group: str | None = None, properties: dict[str, str] | None = None, ) -> Run: body = build_add_body( input_data, user_id=user_id, - conversation_id=conversation_id, group=group, properties=properties, ) @@ -165,7 +159,6 @@ async def search( query: str, topics: _Topics = None, user_id: str | None = None, - conversation_id: str | None = None, group: str | None = None, retrieval_config: RetrievalConfig | None = None, properties: dict[str, str] | None = None, @@ -174,7 +167,6 @@ async def search( query=query, topics=topics, user_id=user_id, - conversation_id=conversation_id, group=group, retrieval_config=retrieval_config, properties=properties, diff --git a/src/engram/_serialization/_builders.py b/src/engram/_serialization/_builders.py index c6f5a77..f3dcf56 100644 --- a/src/engram/_serialization/_builders.py +++ b/src/engram/_serialization/_builders.py @@ -88,15 +88,12 @@ def build_add_body( input_data: AddInput, *, user_id: str | None, - conversation_id: str | None, group: str | None, properties: dict[str, str] | None = None, ) -> dict[str, Any]: 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: - body["conversation_id"] = conversation_id if group is not None: body["group"] = group if properties is not None: @@ -122,7 +119,6 @@ def build_search_body( query: str, topics: list[TopicSelector] | None, user_id: str | None, - conversation_id: str | None, group: str | None, retrieval_config: RetrievalConfig | None, properties: dict[str, str] | None = None, @@ -138,8 +134,6 @@ def build_search_body( body["topics"] = serialized_topics if user_id is not None: body["user_id"] = user_id - if conversation_id is not None: - body["conversation_id"] = conversation_id if group is not None: body["group"] = group if properties is not None: diff --git a/src/engram/_serialization/_parsers.py b/src/engram/_serialization/_parsers.py index f998847..662493d 100644 --- a/src/engram/_serialization/_parsers.py +++ b/src/engram/_serialization/_parsers.py @@ -30,7 +30,6 @@ def parse_memory(data: dict[str, Any]) -> Memory: created_at=data["created_at"], updated_at=data["updated_at"], user_id=data.get("user_id"), - conversation_id=data.get("conversation_id"), tags=data.get("tags"), score=data.get("score"), properties=data.get("properties"), diff --git a/tests/test_client_async.py b/tests/test_client_async.py index 3975240..e77e1b3 100644 --- a/tests/test_client_async.py +++ b/tests/test_client_async.py @@ -131,7 +131,7 @@ async def test_add_conversation() -> None: result = await client.memories.add( [{"role": "user", "content": "hi"}], user_id="u1", - conversation_id="c1", + properties={"conversation_id": "c1"}, ) assert result.run_id == "r3" @@ -234,7 +234,7 @@ async def test_add_conversation_content() -> None: result = await client.memories.add( ConversationInput(messages=[MessageInput(role="user", content="hi")]), user_id="u1", - conversation_id="c1", + properties={"conversation_id": "c1"}, ) assert result.run_id == "r5" @@ -263,7 +263,7 @@ def handler(request: httpx.Request) -> httpx.Response: ], metadata={"session_id": "s1"}, ), - conversation_id="c1", + properties={"conversation_id": "c1"}, ) body = json.loads(captured[0].content) conv = body["input"]["conversation"] @@ -271,7 +271,7 @@ def handler(request: httpx.Request) -> httpx.Response: assert conv["messages"][1]["tool_calls"] == [ {"id": "tc1", "type": "function", "function": {"name": "search", "arguments": "{}"}} ] - assert body["conversation_id"] == "c1" + assert body["properties"] == {"conversation_id": "c1"} # ── memories.get ──────────────────────────────────────────────────────── diff --git a/tests/test_client_sync.py b/tests/test_client_sync.py index 8365435..682823d 100644 --- a/tests/test_client_sync.py +++ b/tests/test_client_sync.py @@ -131,7 +131,7 @@ def test_add_conversation() -> None: result = client.memories.add( [{"role": "user", "content": "hi"}], user_id="u1", - conversation_id="c1", + properties={"conversation_id": "c1"}, ) assert result.run_id == "r3" @@ -162,11 +162,11 @@ def handler(request: httpx.Request) -> httpx.Response: client = _make_client_with_handler(handler) messages = [{"role": "user", "content": "hi"}] - client.memories.add(messages, conversation_id="c1") + client.memories.add(messages, properties={"conversation_id": "c1"}) body = json.loads(captured[0].content) assert body == { "input": {"conversation": {"messages": messages}}, - "conversation_id": "c1", + "properties": {"conversation_id": "c1"}, } @@ -245,7 +245,7 @@ def test_add_conversation_content() -> None: result = client.memories.add( ConversationInput(messages=[MessageInput(role="user", content="hi")]), user_id="u1", - conversation_id="c1", + properties={"conversation_id": "c1"}, ) assert result.run_id == "r5" @@ -273,7 +273,7 @@ def handler(request: httpx.Request) -> httpx.Response: ], metadata={"session_id": "s1"}, ), - conversation_id="c1", + properties={"conversation_id": "c1"}, ) body = json.loads(captured[0].content) conv = body["input"]["conversation"] @@ -281,7 +281,7 @@ def handler(request: httpx.Request) -> httpx.Response: assert conv["messages"][1]["tool_calls"] == [ {"id": "tc1", "type": "function", "function": {"name": "search", "arguments": "{}"}} ] - assert body["conversation_id"] == "c1" + assert body["properties"] == {"conversation_id": "c1"} # ── memories.get ──────────────────────────────────────────────────────── diff --git a/tests/test_serialization.py b/tests/test_serialization.py index d1fe991..98283cb 100644 --- a/tests/test_serialization.py +++ b/tests/test_serialization.py @@ -27,7 +27,6 @@ def test_build_add_body_str() -> None: body = build_add_body( "hello world", user_id=None, - conversation_id=None, group=None, ) assert body == {"input": {"string": {"content": ["hello world"]}}} @@ -37,13 +36,11 @@ def test_build_add_body_str_with_options() -> None: body = build_add_body( "hello", user_id="u1", - conversation_id="c1", group="g1", ) assert body == { "input": {"string": {"content": ["hello"]}}, "user_id": "u1", - "conversation_id": "c1", "group": "g1", } @@ -52,7 +49,6 @@ def test_build_add_body_pre_extracted() -> None: body = build_add_body( PreExtractedInput(items=[PreExtractedItem(content="fact", topic="topic")]), user_id=None, - conversation_id=None, group=None, ) assert body == { @@ -68,13 +64,11 @@ def test_build_add_body_conversation() -> None: body = build_add_body( messages, user_id="u1", - conversation_id="c1", group=None, ) assert body == { "input": {"conversation": {"messages": messages}}, "user_id": "u1", - "conversation_id": "c1", } @@ -82,7 +76,6 @@ def test_build_add_body_string_content() -> None: body = build_add_body( StringInput(content="hello world"), user_id=None, - conversation_id=None, group=None, ) assert body == {"input": {"string": {"content": ["hello world"]}}} @@ -92,13 +85,11 @@ def test_build_add_body_string_content_with_options() -> None: body = build_add_body( StringInput(content="hello"), user_id="u1", - conversation_id="c1", group="g1", ) assert body == { "input": {"string": {"content": ["hello"]}}, "user_id": "u1", - "conversation_id": "c1", "group": "g1", } @@ -111,7 +102,6 @@ def test_build_add_body_conversation_content() -> None: body = build_add_body( ConversationInput(messages=messages), user_id="u1", - conversation_id="c1", group=None, ) assert body == { @@ -124,7 +114,6 @@ def test_build_add_body_conversation_content() -> None: }, }, "user_id": "u1", - "conversation_id": "c1", } @@ -138,7 +127,6 @@ def test_build_add_body_conversation_content_with_metadata() -> None: updated_at="2024-01-02T00:00:00Z", ), user_id=None, - conversation_id=None, group=None, ) conv = body["input"]["conversation"] @@ -152,7 +140,6 @@ def test_build_add_body_conversation_content_with_message_timestamps() -> None: body = build_add_body( ConversationInput(messages=messages), user_id=None, - conversation_id=None, group=None, ) msg = body["input"]["conversation"]["messages"][0] @@ -174,7 +161,6 @@ def test_build_add_body_conversation_content_with_tool_calls() -> None: body = build_add_body( ConversationInput(messages=messages), user_id=None, - conversation_id=None, group=None, ) msg = body["input"]["conversation"]["messages"][0] @@ -199,7 +185,6 @@ def test_build_add_body_conversation_content_with_custom_tool_calls() -> None: body = build_add_body( ConversationInput(messages=messages), user_id=None, - conversation_id=None, group=None, ) msg = body["input"]["conversation"]["messages"][0] @@ -213,7 +198,6 @@ def test_build_add_body_conversation_content_with_tool_role() -> None: body = build_add_body( ConversationInput(messages=messages), user_id=None, - conversation_id=None, group=None, ) msg = body["input"]["conversation"]["messages"][0] @@ -228,7 +212,6 @@ def test_build_add_body_conversation_content_with_developer_role() -> None: body = build_add_body( ConversationInput(messages=messages), user_id=None, - conversation_id=None, group=None, ) msg = body["input"]["conversation"]["messages"][0] @@ -260,7 +243,6 @@ def test_build_search_body_defaults() -> None: query="test", topics=None, user_id=None, - conversation_id=None, group=None, retrieval_config=None, ) @@ -272,7 +254,6 @@ def test_build_search_body_full() -> None: query="test", topics=["a", "b"], user_id="u1", - conversation_id="c1", group="g1", retrieval_config=RetrievalConfig(retrieval_type="vector", limit=5), ) @@ -289,7 +270,6 @@ def test_build_add_body_with_properties() -> None: body = build_add_body( "hello", user_id=None, - conversation_id=None, group=None, properties={"region": "eu", "tier": "pro"}, ) @@ -303,7 +283,6 @@ def test_build_add_body_properties_none_omitted() -> None: body = build_add_body( "hello", user_id=None, - conversation_id=None, group=None, properties=None, ) @@ -318,7 +297,6 @@ def test_build_search_body_with_properties() -> None: query="q", topics=None, user_id=None, - conversation_id=None, group=None, retrieval_config=None, properties={"region": "eu"}, @@ -335,7 +313,6 @@ def test_build_search_body_with_topic_filter() -> None: Topic(name="cleared", properties={"region": None}), ], user_id=None, - conversation_id=None, group=None, retrieval_config=None, ) @@ -351,7 +328,6 @@ def test_build_search_body_topic_filter_without_properties() -> None: query="q", topics=[Topic(name="t1")], user_id=None, - conversation_id=None, group=None, retrieval_config=None, ) @@ -399,16 +375,15 @@ def test_parse_memory_with_optional_fields() -> None: data = { **SAMPLE_MEMORY, "user_id": "u1", - "conversation_id": "c1", "tags": ["x"], "score": 0.95, - "properties": {"region": "eu", "tier": "pro"}, + "properties": {"region": "eu", "conversation_id": "c1"}, } mem = parse_memory(data) assert mem.user_id == "u1" assert mem.tags == ["x"] assert mem.score == 0.95 - assert mem.properties == {"region": "eu", "tier": "pro"} + assert mem.properties == {"region": "eu", "conversation_id": "c1"} # ── parse_search_results ──────────────────────────────────────────────── From e278f47ad83308d284f44fca46c7a35fddaf4d70 Mon Sep 17 00:00:00 2001 From: Augustas Date: Tue, 5 May 2026 15:16:57 +0300 Subject: [PATCH 3/3] improve types --- src/engram/_resources/memories.py | 7 ++----- src/engram/_serialization/_builders.py | 11 +++++++---- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/src/engram/_resources/memories.py b/src/engram/_resources/memories.py index 4ab55e7..2db72d5 100644 --- a/src/engram/_resources/memories.py +++ b/src/engram/_resources/memories.py @@ -1,6 +1,5 @@ from __future__ import annotations -from typing import TypeAlias from uuid import UUID from .._http import AsyncHttpTransport, HttpTransport @@ -24,8 +23,6 @@ _MEMORIES_PATH = "/v1/memories" _MEMORIES_SEARCH_PATH = "/v1/memories/search" -_Topics: TypeAlias = list[TopicSelector] | None - def _memory_path(memory_id: str | UUID) -> str: return f"{_MEMORIES_PATH}/{memory_id}" @@ -85,7 +82,7 @@ def search( self, *, query: str, - topics: _Topics = None, + topics: list[TopicSelector] | None = None, user_id: str | None = None, group: str | None = None, retrieval_config: RetrievalConfig | None = None, @@ -157,7 +154,7 @@ async def search( self, *, query: str, - topics: _Topics = None, + topics: list[TopicSelector] | None = None, user_id: str | None = None, group: str | None = None, retrieval_config: RetrievalConfig | None = None, diff --git a/src/engram/_serialization/_builders.py b/src/engram/_serialization/_builders.py index f3dcf56..ef28334 100644 --- a/src/engram/_serialization/_builders.py +++ b/src/engram/_serialization/_builders.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Any +from typing import Any, TypeAlias from .._models import ( AddInput, @@ -67,18 +67,21 @@ def _serialize_conversation_content(content: ConversationInput) -> dict[str, Any return {"conversation": conversation} -def _serialize_topic(topic: TopicSelector) -> str | dict[str, Any]: +_SerializedTopic: TypeAlias = str | dict[str, str | dict[str, str | None]] + + +def _serialize_topic(topic: TopicSelector) -> _SerializedTopic: if isinstance(topic, str): return topic if isinstance(topic, Topic): - out: dict[str, Any] = {"name": topic.name} + out: dict[str, str | dict[str, str | None]] = {"name": topic.name} if topic.properties is not None: out["properties"] = dict(topic.properties) return out raise TypeError(f"Unsupported topic type: {type(topic)}") # pragma: no cover -def _serialize_topics(topics: list[TopicSelector] | None) -> list[str | dict[str, Any]] | None: +def _serialize_topics(topics: list[TopicSelector] | None) -> list[_SerializedTopic] | None: if topics is None: return None return [_serialize_topic(t) for t in topics]