Skip to content

Commit cd2c937

Browse files
edwinyyyuedwinyyyu
committed
WIP: wiring god commit
Signed-off-by: Edwin Yu <edwinyyyu@gmail.com>
1 parent 14b939f commit cd2c937

22 files changed

Lines changed: 1716 additions & 7 deletions

File tree

‎packages/client/src/memmachine_client/memory.py‎

Lines changed: 106 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
AddSemanticTagResponse,
2727
AddSemanticTagSpec,
2828
ConfigureEpisodicMemorySpec,
29+
ConfigureEventMemorySpec,
2930
ConfigureSemanticSetSpec,
3031
CreateSemanticSetTypeResponse,
3132
CreateSemanticSetTypeSpec,
@@ -36,7 +37,9 @@
3637
DeleteSemanticTagSpec,
3738
DisableSemanticCategorySpec,
3839
EpisodicMemoryConfigEntry,
40+
EventMemoryConfigEntry,
3941
GetEpisodicMemoryConfigSpec,
42+
GetEventMemoryConfigSpec,
4043
GetFeatureSpec,
4144
GetSemanticCategorySetIdsResponse,
4245
GetSemanticCategorySetIdsSpec,
@@ -429,7 +432,7 @@ def search(
429432
agent_mode=agent_mode,
430433
filter=filter_str,
431434
set_metadata=set_metadata,
432-
types=[MemoryType.Episodic, MemoryType.Semantic], # Search both types
435+
types=[MemoryType.Episodic, MemoryType.Semantic, MemoryType.Event],
433436
)
434437
v2_search_data = spec.model_dump(mode="json", exclude_none=True)
435438

@@ -1811,6 +1814,108 @@ def configure_episodic_memory(
18111814
else:
18121815
return True
18131816

1817+
def get_event_memory_config(
1818+
self,
1819+
timeout: int | None = None,
1820+
) -> EventMemoryConfigEntry:
1821+
"""
1822+
Get event memory configuration for this project.
1823+
1824+
Args:
1825+
timeout: Request timeout in seconds (uses client default if not provided).
1826+
1827+
Returns:
1828+
The event memory configuration entry.
1829+
1830+
Raises:
1831+
requests.RequestException: If the request fails.
1832+
RuntimeError: If the client has been closed.
1833+
1834+
"""
1835+
if self._client_closed:
1836+
raise RuntimeError("Cannot get event memory config: client has been closed")
1837+
1838+
spec = GetEventMemoryConfigSpec(
1839+
org_id=self.__org_id,
1840+
project_id=self.__project_id,
1841+
)
1842+
v2_data = spec.model_dump(mode="json", exclude_none=True)
1843+
1844+
try:
1845+
response = self.client.request(
1846+
"POST",
1847+
f"{self.client.base_url}/api/v2/memory/event/config/get",
1848+
json=v2_data,
1849+
timeout=timeout,
1850+
)
1851+
response.raise_for_status()
1852+
response_data = response.json()
1853+
result = EventMemoryConfigEntry(**response_data)
1854+
logger.debug("Successfully retrieved event memory config")
1855+
except Exception:
1856+
logger.exception("Failed to get event memory config")
1857+
raise
1858+
else:
1859+
return result
1860+
1861+
def configure_event_memory(
1862+
self,
1863+
*,
1864+
embedder: str,
1865+
reranker: str | None = None,
1866+
properties_schema: dict[str, str] | None = None,
1867+
derive_sentences: bool = False,
1868+
max_text_chunk_length: int = 2000,
1869+
timeout: int | None = None,
1870+
) -> bool:
1871+
"""
1872+
Configure event memory for this project.
1873+
1874+
Args:
1875+
embedder: Resource ID of the Embedder instance.
1876+
reranker: Resource ID of the Reranker instance, or None.
1877+
properties_schema: User-defined filterable properties and their types.
1878+
derive_sentences: Whether to derive sentence-level derivatives.
1879+
max_text_chunk_length: Max code-point length for text chunking.
1880+
timeout: Request timeout in seconds (uses client default if not provided).
1881+
1882+
Returns:
1883+
True if configuration was successful.
1884+
1885+
Raises:
1886+
requests.RequestException: If the request fails.
1887+
RuntimeError: If the client has been closed.
1888+
1889+
"""
1890+
if self._client_closed:
1891+
raise RuntimeError("Cannot configure event memory: client has been closed")
1892+
1893+
spec = ConfigureEventMemorySpec(
1894+
org_id=self.__org_id,
1895+
project_id=self.__project_id,
1896+
embedder=embedder,
1897+
reranker=reranker,
1898+
properties_schema=properties_schema or {},
1899+
derive_sentences=derive_sentences,
1900+
max_text_chunk_length=max_text_chunk_length,
1901+
)
1902+
v2_data = spec.model_dump(mode="json", exclude_none=True)
1903+
1904+
try:
1905+
response = self.client.request(
1906+
"POST",
1907+
f"{self.client.base_url}/api/v2/memory/event/config",
1908+
json=v2_data,
1909+
timeout=timeout,
1910+
)
1911+
response.raise_for_status()
1912+
logger.debug("Successfully configured event memory")
1913+
except Exception:
1914+
logger.exception("Failed to configure event memory")
1915+
raise
1916+
else:
1917+
return True
1918+
18141919
def mark_client_closed(self) -> None:
18151920
"""Mark this memory instance as closed by its owning client."""
18161921
self._client_closed = True

‎packages/common/common_tests/__init__.py‎

Whitespace-only changes.

‎packages/common/common_tests/test_event_memory_data_types.py‎

Lines changed: 240 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,240 @@
1+
"""Tests for common event memory data types, formatting, and unification."""
2+
3+
from datetime import UTC, datetime, timedelta, timezone
4+
from uuid import uuid4
5+
6+
import pytest
7+
8+
from memmachine_common.api.event_memory.data_types import (
9+
EventMemoryCitationContext,
10+
EventMemoryFormatOptions,
11+
EventMemoryMessageContext,
12+
EventMemoryQueryResult,
13+
EventMemoryScoredSegmentContext,
14+
EventMemorySegment,
15+
EventMemoryText,
16+
format_segment_context,
17+
)
18+
19+
20+
def _make_segment(
21+
*,
22+
event_uuid=None,
23+
index=0,
24+
offset=0,
25+
timestamp=None,
26+
text="hello",
27+
context=None,
28+
properties=None,
29+
):
30+
return EventMemorySegment(
31+
uuid=uuid4(),
32+
event_uuid=event_uuid or uuid4(),
33+
index=index,
34+
offset=offset,
35+
timestamp=timestamp or datetime(2026, 1, 15, 10, 30, tzinfo=UTC),
36+
context=context,
37+
block=EventMemoryText(text=text),
38+
properties=properties or {},
39+
)
40+
41+
42+
class TestSegmentRoundTrip:
43+
def test_serialize_deserialize(self):
44+
seg = _make_segment(
45+
properties={"count": 42, "ts": datetime(2026, 1, 1, tzinfo=UTC)},
46+
)
47+
seg2 = EventMemorySegment.model_validate(seg.model_dump(mode="json"))
48+
assert seg.uuid == seg2.uuid
49+
assert seg.properties == seg2.properties
50+
for key in seg.properties:
51+
assert type(seg.properties[key]) is type(seg2.properties[key])
52+
53+
def test_empty_properties(self):
54+
seg = _make_segment()
55+
seg2 = EventMemorySegment.model_validate(seg.model_dump(mode="json"))
56+
assert seg2.properties == {}
57+
58+
59+
class TestFormatSegmentContext:
60+
def test_message_context(self):
61+
seg = _make_segment(
62+
context=EventMemoryMessageContext(source="alice"),
63+
text="hi there",
64+
)
65+
result = format_segment_context([seg])
66+
assert "alice:" in result
67+
assert "hi there" in result
68+
69+
def test_citation_context(self):
70+
seg = _make_segment(
71+
context=EventMemoryCitationContext(source="doc.pdf"),
72+
text="some quote",
73+
)
74+
result = format_segment_context([seg])
75+
assert "From 'doc.pdf':" in result
76+
assert "some quote" in result
77+
78+
def test_no_context(self):
79+
seg = _make_segment(text="bare text")
80+
result = format_segment_context([seg])
81+
assert "bare text" in result
82+
83+
def test_continuation_same_event_and_index(self):
84+
event_uuid = uuid4()
85+
seg1 = _make_segment(event_uuid=event_uuid, index=0, offset=0, text="part1")
86+
seg2 = _make_segment(event_uuid=event_uuid, index=0, offset=1, text="part2")
87+
result = format_segment_context([seg1, seg2])
88+
assert "part1part2" in result
89+
90+
def test_timezone_formatting(self):
91+
tz = timezone(timedelta(hours=9))
92+
seg = _make_segment(
93+
timestamp=datetime(2026, 1, 15, 10, 30, tzinfo=UTC),
94+
text="test",
95+
)
96+
opts = EventMemoryFormatOptions(timezone=tz, show_timezone_label=True)
97+
result = format_segment_context([seg], format_options=opts)
98+
assert "UTC+09:00" in result
99+
100+
def test_no_timezone_label(self):
101+
seg = _make_segment(text="test")
102+
opts = EventMemoryFormatOptions(show_timezone_label=False)
103+
result = format_segment_context([seg], format_options=opts)
104+
assert "UTC" not in result
105+
106+
107+
class TestScoredSegmentContextToString:
108+
def test_to_string(self):
109+
seg = _make_segment(
110+
context=EventMemoryMessageContext(source="bob"),
111+
text="hey",
112+
)
113+
scored = EventMemoryScoredSegmentContext(
114+
seed_segment_uuid=seg.uuid,
115+
score=0.9,
116+
segments=[seg],
117+
)
118+
result = scored.to_string()
119+
assert "bob:" in result
120+
assert "hey" in result
121+
122+
123+
class TestBuildContext:
124+
def _make_scored_context(self, num_segments=3, score=0.5):
125+
seed_uuid = uuid4()
126+
event_uuid = uuid4()
127+
segments = [
128+
_make_segment(event_uuid=event_uuid, index=i, text=f"seg{i}")
129+
for i in range(num_segments)
130+
]
131+
segments[0] = EventMemorySegment(
132+
uuid=seed_uuid,
133+
event_uuid=event_uuid,
134+
index=0,
135+
offset=0,
136+
timestamp=datetime(2026, 1, 15, 10, 30, tzinfo=UTC),
137+
block=EventMemoryText(text="seed"),
138+
properties={},
139+
)
140+
return EventMemoryScoredSegmentContext(
141+
seed_segment_uuid=seed_uuid,
142+
score=score,
143+
segments=segments,
144+
)
145+
146+
def test_all_fit(self):
147+
ctx = self._make_scored_context(num_segments=3, score=0.9)
148+
qr = EventMemoryQueryResult(scored_segment_contexts=[ctx])
149+
result = qr.build_context(max_num_segments=10)
150+
assert len(result) == 3
151+
152+
def test_budget_respected(self):
153+
ctx = self._make_scored_context(num_segments=5, score=0.9)
154+
qr = EventMemoryQueryResult(scored_segment_contexts=[ctx])
155+
result = qr.build_context(max_num_segments=2)
156+
assert len(result) == 2
157+
158+
def test_higher_score_first(self):
159+
ctx1 = self._make_scored_context(num_segments=2, score=0.3)
160+
ctx2 = self._make_scored_context(num_segments=2, score=0.9)
161+
qr = EventMemoryQueryResult(scored_segment_contexts=[ctx2, ctx1])
162+
result = qr.build_context(max_num_segments=2)
163+
# Should pick from ctx2 (higher score) first
164+
assert all(seg in ctx2.segments for seg in result)
165+
166+
def test_deduplication(self):
167+
shared_seg = _make_segment(text="shared")
168+
ctx1 = EventMemoryScoredSegmentContext(
169+
seed_segment_uuid=shared_seg.uuid,
170+
score=0.9,
171+
segments=[shared_seg],
172+
)
173+
ctx2 = EventMemoryScoredSegmentContext(
174+
seed_segment_uuid=shared_seg.uuid,
175+
score=0.5,
176+
segments=[shared_seg],
177+
)
178+
qr = EventMemoryQueryResult(scored_segment_contexts=[ctx1, ctx2])
179+
result = qr.build_context(max_num_segments=10)
180+
assert len(result) == 1
181+
182+
def test_chronological_order(self):
183+
seg1 = _make_segment(
184+
timestamp=datetime(2026, 1, 15, 10, 0, tzinfo=UTC), text="first"
185+
)
186+
seg2 = _make_segment(
187+
timestamp=datetime(2026, 1, 15, 11, 0, tzinfo=UTC), text="second"
188+
)
189+
ctx = EventMemoryScoredSegmentContext(
190+
seed_segment_uuid=seg2.uuid,
191+
score=0.9,
192+
segments=[seg2, seg1],
193+
)
194+
qr = EventMemoryQueryResult(scored_segment_contexts=[ctx])
195+
result = qr.build_context(max_num_segments=10)
196+
assert result[0].timestamp < result[1].timestamp
197+
198+
def test_empty_result(self):
199+
qr = EventMemoryQueryResult(scored_segment_contexts=[])
200+
result = qr.build_context(max_num_segments=10)
201+
assert result == []
202+
203+
204+
class TestToString:
205+
def test_to_string(self):
206+
seg = _make_segment(
207+
context=EventMemoryMessageContext(source="alice"),
208+
text="hello",
209+
)
210+
ctx = EventMemoryScoredSegmentContext(
211+
seed_segment_uuid=seg.uuid,
212+
score=0.9,
213+
segments=[seg],
214+
)
215+
qr = EventMemoryQueryResult(scored_segment_contexts=[ctx])
216+
result = qr.to_string(max_num_segments=10)
217+
assert "alice:" in result
218+
assert "hello" in result
219+
220+
221+
class TestPropertiesRoundTrip:
222+
@pytest.mark.parametrize(
223+
"props",
224+
[
225+
{"count": 42},
226+
{"name": "test"},
227+
{"ratio": 3.14},
228+
{"flag": True},
229+
{"ts": datetime(2026, 1, 15, tzinfo=UTC)},
230+
{
231+
"count": 42,
232+
"name": "test",
233+
"ts": datetime(2026, 1, 15, tzinfo=UTC),
234+
},
235+
],
236+
)
237+
def test_round_trip(self, props):
238+
seg = _make_segment(properties=props)
239+
seg2 = EventMemorySegment.model_validate(seg.model_dump(mode="json"))
240+
assert seg.properties == seg2.properties

‎packages/common/src/memmachine_common/api/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ class MemoryType(Enum):
88

99
Semantic = "semantic"
1010
Episodic = "episodic"
11+
Event = "event"
1112

1213

1314
class EpisodeType(Enum):

0 commit comments

Comments
 (0)