From 86139d478d5cdae42e428dedcc4552917d6e4ea7 Mon Sep 17 00:00:00 2001 From: paulpaliychuk Date: Mon, 24 Aug 2026 19:15:16 -0400 Subject: [PATCH 1/2] Reset Fern Replay so the v4 SDK generates clean Unfreezing .fernignore was not enough. Fern Replay records each generation's tree hash in .fern/replay.lock and re-applies everything that has diverged since as a customization, independent of .fernignore. The lock on this branch is v3's history, so every v3 hand-written file still looked like a customization and came back on top of the generated v4 client -- the last generation failed mypy on a client.py that Replay had restored. Removes the lock so v4 starts with no recorded customizations, and stops freezing it: the file is Fern's own state, and pinning it is what made the history outlive the line it described. Anything worth carrying forward from v3 is re-added deliberately, which Replay will then track from a v4 baseline. --- .fern/replay.lock | 28 ---------------------------- .fernignore | 3 --- 2 files changed, 31 deletions(-) delete mode 100644 .fern/replay.lock diff --git a/.fern/replay.lock b/.fern/replay.lock deleted file mode 100644 index afed4a45..00000000 --- a/.fern/replay.lock +++ /dev/null @@ -1,28 +0,0 @@ -# DO NOT EDIT MANUALLY - Managed by Fern Replay -version: "1.0" -generations: - - commit_sha: 36ef7a1c35b8b6f4614d1095dae12b8328832a89 - tree_hash: ff6c7f6dca046cfe71ff8d197ce03fd3f3cefccd - timestamp: 2026-07-28T21:28:05.428Z - cli_version: unknown - generator_versions: {} - - commit_sha: a7d57c567987a6cfead733b64cd20d025c4ea593 - tree_hash: 536ce8a8b7d3e46a6e3b82f36938c3518c6594e9 - timestamp: 2026-07-28T21:28:05.773Z - cli_version: unknown - generator_versions: - fernapi/fern-python-sdk: 4.25.5 - - commit_sha: 06329e0a67f5396af716b0f9065bf54ff9b4e613 - tree_hash: f736bb9cd3ec813048fa6bdc02e54e1f974fb959 - timestamp: 2026-08-04T17:50:26.852Z - cli_version: unknown - generator_versions: - fernapi/fern-python-sdk: 4.25.5 - - commit_sha: 8a3fb853912eadd78157bed5f9811ae936e834a8 - tree_hash: a5aa6b923d390b275f7e12ca55ba9b44d0689ffd - timestamp: 2026-08-11T20:04:50.130Z - cli_version: unknown - generator_versions: - fernapi/fern-python-sdk: 4.25.5 -current_generation: 8a3fb853912eadd78157bed5f9811ae936e834a8 -patches: [] diff --git a/.fernignore b/.fernignore index 0aa54bcd..07850b98 100644 --- a/.fernignore +++ b/.fernignore @@ -4,11 +4,8 @@ # what must not be regenerated stays here: # .github — the SDK Release audit fails closed unless CI is protected # LICENSE — legal text, not a generated artifact -# .fern/ — Fern's own replay state # # Anything ported forward from v3 is re-added deliberately, file by file. .github LICENSE -.fern/replay.lock -.fern/replay.yml .gitattributes From 7e567554ad2e1cef4b6746998f0436d4445ba87e Mon Sep 17 00:00:00 2001 From: paulpaliychuk Date: Mon, 24 Aug 2026 19:19:54 -0400 Subject: [PATCH 2/2] Revert the v4 wrapper port so the branch is consistent again The earlier port rewrote client.py, external_clients, graph/utils.py and their tests against v4 type names, while the checked-in generated code is still v3's. The branch could not typecheck, so nothing on it could be verified, and the port itself is now moot: these files are unfrozen and Replay is reset, so the next generation writes them. Restores them to the v3 state the branch was cut from. Unlike Go, they cannot simply be deleted: __init__.py is generated and imports client.py, so removing it breaks the package before a generation can replace it. The Pydantic ontology DSL is the one piece here Fern will not regenerate. It stays in history and is re-added deliberately if it earns its place on v4. --- src/zep_cloud/client.py | 28 +- src/zep_cloud/external_clients/graph.py | 474 ++++++++++++++++++++---- src/zep_cloud/graph/utils.py | 27 +- tests/graph/test_utils.py | 53 ++- 4 files changed, 451 insertions(+), 131 deletions(-) diff --git a/src/zep_cloud/client.py b/src/zep_cloud/client.py index 3a1a409f..b793251d 100644 --- a/src/zep_cloud/client.py +++ b/src/zep_cloud/client.py @@ -21,7 +21,7 @@ def __init__( ): env_api_url = os.getenv("ZEP_API_URL") if env_api_url: - base_url = f"{env_api_url}/api/v4" + base_url = f"{env_api_url}/api/v2" super().__init__( base_url=base_url, environment=environment, @@ -30,16 +30,8 @@ def __init__( follow_redirects=follow_redirects, httpx_client=httpx_client ) - self._external_user = UserClient(client_wrapper=self._client_wrapper) - self._external_graph = GraphClient(client_wrapper=self._client_wrapper) - - @property - def user(self) -> UserClient: # type: ignore[override] - return self._external_user - - @property - def graph(self) -> GraphClient: # type: ignore[override] - return self._external_graph + self.user = UserClient(client_wrapper=self._client_wrapper) + self.graph = GraphClient(client_wrapper=self._client_wrapper) class AsyncZep(AsyncBaseClient): def __init__( @@ -54,7 +46,7 @@ def __init__( ): env_api_url = os.getenv("ZEP_API_URL") if env_api_url: - base_url = f"{env_api_url}/api/v4" + base_url = f"{env_api_url}/api/v2" super().__init__( base_url=base_url, environment=environment, @@ -63,13 +55,5 @@ def __init__( follow_redirects=follow_redirects, httpx_client=httpx_client ) - self._external_user = AsyncUserClient(client_wrapper=self._client_wrapper) - self._external_graph = AsyncGraphClient(client_wrapper=self._client_wrapper) - - @property - def user(self) -> AsyncUserClient: # type: ignore[override] - return self._external_user - - @property - def graph(self) -> AsyncGraphClient: # type: ignore[override] - return self._external_graph + self.user = AsyncUserClient(client_wrapper=self._client_wrapper) + self.graph = AsyncGraphClient(client_wrapper=self._client_wrapper) diff --git a/src/zep_cloud/external_clients/graph.py b/src/zep_cloud/external_clients/graph.py index 4e620ea3..3c650657 100644 --- a/src/zep_cloud/external_clients/graph.py +++ b/src/zep_cloud/external_clients/graph.py @@ -1,7 +1,7 @@ import typing +from zep_cloud import EdgeType, EntityEdgeSourceTarget from zep_cloud.core.client_wrapper import AsyncClientWrapper, SyncClientWrapper -from zep_cloud.core.request_options import RequestOptions from zep_cloud.external_clients.ontology import ( EdgeModel, edge_model_to_api_schema, @@ -9,107 +9,449 @@ ) from zep_cloud.graph.client import AsyncGraphClient as AsyncBaseGraphClient from zep_cloud.graph.client import GraphClient as BaseGraphClient -from zep_cloud.types import EdgeSourceTarget, EdgeType, EntityType, Ontology +from zep_cloud.types import EntityType if typing.TYPE_CHECKING: from zep_cloud.external_clients.ontology import EntityModel +from zep_cloud.core.request_options import RequestOptions + -EdgeSpec = typing.Union[ - "EdgeModel", - typing.Tuple["EdgeModel", typing.List[EdgeSourceTarget]], -] +class GraphClient(BaseGraphClient): + def __init__(self, *, client_wrapper: SyncClientWrapper): + super().__init__(client_wrapper=client_wrapper) + def set_ontology( + self, + entities: dict[str, "EntityModel"], + edges: typing.Optional[ + dict[ + str, + typing.Union[ + "EdgeModel", + typing.Tuple["EdgeModel", typing.List[EntityEdgeSourceTarget]], + ], + ] + ] = None, + user_ids: typing.Optional[typing.List[str]] = None, + graph_ids: typing.Optional[typing.List[str]] = None, + request_options: typing.Optional[RequestOptions] = None, + ): + """ + Sets the entity and edge types for a project, replacing any existing ones. -def build_ontology( - entities: typing.Dict[str, "EntityModel"], - edges: typing.Optional[typing.Dict[str, EdgeSpec]] = None, -) -> typing.Tuple[typing.List[EntityType], typing.List[EdgeType]]: - """Turn the Pydantic ontology models into the types the v4 API accepts. + Parameters + ---------- + entities : dict[str, "EntityModel"] + Entity type definitions. - This is the whole value the hand-written layer adds: the wire shape is a - list of entity and edge types, and this derives it from Python classes so an - ontology is declared once, in the type system. - """ - api_entity_types: typing.List[EntityType] = [] - for name, entity in entities.items(): - api_entity_types.append(EntityType(**entity_model_to_api_schema(entity, name))) + edges : typing.Optional[dict[str, typing.Union["EdgeModel", typing.Tuple["EdgeModel", typing.List[EntityEdgeSourceTarget]]]]] + Edge type definitions. - api_edge_types: typing.List[EdgeType] = [] - if edges: - for name, edge_data in edges.items(): - if isinstance(edge_data, tuple): - edge_model, source_targets = edge_data - else: - edge_model, source_targets = edge_data, None + user_ids : typing.Optional[typing.List[str]] - edge_dict = edge_model_to_api_schema(edge_model, name) - if source_targets: - edge_dict["source_targets"] = [ - st.dict() if hasattr(st, "dict") else st for st in source_targets - ] - api_edge_types.append(EdgeType(**edge_dict)) + The user identifiers for which to set the ontology. - return api_entity_types, api_edge_types + graph_ids : typing.Optional[typing.List[str]] + The graph identifiers for which to set the ontology. + request_options : typing.Optional[RequestOptions] + Request-specific configuration. -class GraphClient(BaseGraphClient): - def __init__(self, *, client_wrapper: SyncClientWrapper): - super().__init__(client_wrapper=client_wrapper) + Examples + -------- + + class Destination(EntityModel): + + \"""A destination is a place that travelers visit.\""" + destination_name: EntityText = Field( + description="The name of the destination", + default=None + ) + country: EntityText = Field( + description="The country of the destination", + default=None + ) + region: EntityText = Field( + description="The region of the destination", + default=None + ) + description: EntityText = Field( + description="A description of the destination", + default=None + ) + + + class TravelingTo(EdgeModel): + + \"""An edge representing a traveler going to a destination.\""" + travel_date: EntityText = Field( + description="The date of travel to this destination", + default=None + ) + purpose: EntityText = Field( + description="The purpose of travel (Business, Leisure, etc.)", + default=None + ) + + client.graph.set_ontology( + entities={ + "Destination": Destination, + }, + edges={ + "TRAVELING_TO": ( + TravelingTo, + [ + EntityEdgeSourceTarget( + source="User", + target="Destination" + ) + ] + ), + } + ) + """ + return self.set_entity_types( + entities=entities, + edges=edges, + user_ids=user_ids, + graph_ids=graph_ids, + request_options=request_options, + ) def set_entity_types( self, - graph_uuid: str, - entities: typing.Dict[str, "EntityModel"], - edges: typing.Optional[typing.Dict[str, EdgeSpec]] = None, + entities: dict[str, "EntityModel"], + edges: typing.Optional[ + dict[ + str, + typing.Union[ + "EdgeModel", + typing.Tuple["EdgeModel", typing.List[EntityEdgeSourceTarget]], + ], + ] + ] = None, + user_ids: typing.Optional[typing.List[str]] = None, + graph_ids: typing.Optional[typing.List[str]] = None, request_options: typing.Optional[RequestOptions] = None, - ) -> Ontology: + ): """ - Set the entity and edge types for one graph, replacing its existing ontology. + Sets the entity and edge types for a project, replacing any existing ones. + + Parameters + ---------- + entities : dict[str, "EntityModel"] + + edges : typing.Optional[dict[str, typing.Union["EdgeModel", typing.Tuple["EdgeModel", typing.List[EntityEdgeSourceTarget]]]]] + + user_ids : typing.Optional[typing.List[str]] - The graph is an explicit argument. v3 took ``user_ids`` and ``graph_ids`` - and fanned out server-side; v4 has one ontology endpoint per scope - (spec-3 14.4), so a caller targeting several graphs calls this once per - graph, and a caller targeting the project default passes the same - ``build_ontology`` output to ``client.project.set_ontology``. + The user identifiers for which to set the ontology. + + graph_ids : typing.Optional[typing.List[str]] + The graph identifiers for which to set the ontology. + + request_options : typing.Optional[RequestOptions] + Request-specific configuration. Examples -------- + + class Destination(EntityModel): + + \"""A destination is a place that travelers visit.\""" + destination_name: EntityText = Field( + description="The name of the destination", + default=None + ) + country: EntityText = Field( + description="The country of the destination", + default=None + ) + region: EntityText = Field( + description="The region of the destination", + default=None + ) + description: EntityText = Field( + description="A description of the destination", + default=None + ) + + + class TravelingTo(EdgeModel): + + \"""An edge representing a traveler going to a destination.\""" + travel_date: EntityText = Field( + description="The date of travel to this destination", + default=None + ) + purpose: EntityText = Field( + description="The purpose of travel (Business, Leisure, etc.)", + default=None + ) + client.graph.set_entity_types( - graph_uuid="...", - entities={"Traveler": Traveler}, + entities={ + "Destination": Destination, + }, edges={ - "TRAVELED_TO": ( - TraveledTo, - [EdgeSourceTarget(source_entity_type="Traveler", target_entity_type="Destination")], + "TRAVELING_TO": ( + TravelingTo, + [ + EntityEdgeSourceTarget( + source="User", + target="Destination" + ) + ] ), - }, + } ) """ - entity_types, edge_types = build_ontology(entities, edges) - return self.set_ontology( - graph_uuid, - entity_types=entity_types, - edge_types=edge_types, + api_entity_types: list[EntityType] = [] + api_edge_types: list[EdgeType] = [] + + for name, entity in entities.items(): + entity_dict = entity_model_to_api_schema(entity, name) + api_entity_types.append(EntityType(**entity_dict)) + + if edges: + for name, edge_data in edges.items(): + # Handle both EdgeModel directly and tuple of (model, source_targets) + if isinstance(edge_data, tuple): + edge_model, source_targets = edge_data + else: + edge_model = edge_data + source_targets = None + + edge_dict = edge_model_to_api_schema(edge_model, name) + if source_targets: + edge_dict["source_targets"] = [st.dict() for st in source_targets] + api_edge_types.append(EdgeType(**edge_dict)) + res = self.set_entity_types_internal( + entity_types=api_entity_types, + edge_types=api_edge_types, + user_ids=user_ids, + graph_ids=graph_ids, request_options=request_options, ) + return res class AsyncGraphClient(AsyncBaseGraphClient): def __init__(self, *, client_wrapper: AsyncClientWrapper): super().__init__(client_wrapper=client_wrapper) + async def set_ontology( + self, + entities: dict[str, "EntityModel"], + edges: typing.Optional[ + dict[ + str, + typing.Union[ + "EdgeModel", + typing.Tuple["EdgeModel", typing.List[EntityEdgeSourceTarget]], + ], + ] + ] = None, + user_ids: typing.Optional[typing.List[str]] = None, + graph_ids: typing.Optional[typing.List[str]] = None, + request_options: typing.Optional[RequestOptions] = None, + ): + """ + Sets the entity and edge types for a project, replacing any existing ones. + + Parameters + ---------- + entities : dict[str, "EntityModel"] + Entity type definitions. + + edges : typing.Optional[dict[str, typing.Union["EdgeModel", typing.Tuple["EdgeModel", typing.List[EntityEdgeSourceTarget]]]]] + Edge type definitions. + + user_ids : typing.Optional[typing.List[str]] + + The user identifiers for which to set the ontology. + + graph_ids : typing.Optional[typing.List[str]] + The graph identifiers for which to set the ontology. + + request_options : typing.Optional[RequestOptions] + Request-specific configuration. + + Examples + -------- + + class Destination(EntityModel): + + \"""A destination is a place that travelers visit.\""" + destination_name: EntityText = Field( + description="The name of the destination", + default=None + ) + country: EntityText = Field( + description="The country of the destination", + default=None + ) + region: EntityText = Field( + description="The region of the destination", + default=None + ) + description: EntityText = Field( + description="A description of the destination", + default=None + ) + + + class TravelingTo(EdgeModel): + + \"""An edge representing a traveler going to a destination.\""" + travel_date: EntityText = Field( + description="The date of travel to this destination", + default=None + ) + purpose: EntityText = Field( + description="The purpose of travel (Business, Leisure, etc.)", + default=None + ) + + await client.graph.set_ontology( + entities={ + "Destination": Destination, + }, + edges={ + "TRAVELING_TO": ( + TravelingTo, + [ + EntityEdgeSourceTarget( + source="User", + target="Destination" + ) + ] + ), + } + ) + """ + return await self.set_entity_types( + entities=entities, + edges=edges, + request_options=request_options, + user_ids=user_ids, + graph_ids=graph_ids + ) + async def set_entity_types( self, - graph_uuid: str, - entities: typing.Dict[str, "EntityModel"], - edges: typing.Optional[typing.Dict[str, EdgeSpec]] = None, + entities: dict[str, "EntityModel"], + edges: typing.Optional[ + dict[ + str, + typing.Union[ + "EdgeModel", + typing.Tuple["EdgeModel", typing.List[EntityEdgeSourceTarget]], + ], + ] + ] = None, + user_ids: typing.Optional[typing.List[str]] = None, + graph_ids: typing.Optional[typing.List[str]] = None, request_options: typing.Optional[RequestOptions] = None, - ) -> Ontology: - """Asynchronous counterpart of :meth:`GraphClient.set_entity_types`.""" - entity_types, edge_types = build_ontology(entities, edges) - return await self.set_ontology( - graph_uuid, - entity_types=entity_types, - edge_types=edge_types, + ): + """ + Sets the entity and edge types for a project, replacing any existing ones. + + Parameters + ---------- + entities : dict[str, "EntityModel"] + + edges : typing.Optional[dict[str, typing.Union["EdgeModel", typing.Tuple["EdgeModel", typing.List[EntityEdgeSourceTarget]]]]] + + user_ids : typing.Optional[typing.List[str]] + + The user identifiers for which to set the ontology. + + graph_ids : typing.Optional[typing.List[str]] + The graph identifiers for which to set the ontology. + + request_options : typing.Optional[RequestOptions] + Request-specific configuration. + + Examples + -------- + + class Destination(EntityModel): + + \"""A destination is a place that travelers visit.\""" + destination_name: EntityText = Field( + description="The name of the destination", + default=None + ) + country: EntityText = Field( + description="The country of the destination", + default=None + ) + region: EntityText = Field( + description="The region of the destination", + default=None + ) + description: EntityText = Field( + description="A description of the destination", + default=None + ) + + + class TravelingTo(EdgeModel): + + \"""An edge representing a traveler going to a destination.\""" + travel_date: EntityText = Field( + description="The date of travel to this destination", + default=None + ) + purpose: EntityText = Field( + description="The purpose of travel (Business, Leisure, etc.)", + default=None + ) + + await client.graph.set_entity_types( + entities={ + "Destination": Destination, + }, + edges={ + "TRAVELING_TO": ( + TravelingTo, + [ + EntityEdgeSourceTarget( + source="User", + target="Destination" + ) + ] + ), + } + ) + """ + api_entity_types: list[EntityType] = [] + api_edge_types: list[EdgeType] = [] + + for name, entity in entities.items(): + entity_dict = entity_model_to_api_schema(entity, name) + api_entity_types.append(EntityType(**entity_dict)) + + if edges: + for name, edge_data in edges.items(): + # Handle both EdgeModel directly and tuple of (model, source_targets) + if isinstance(edge_data, tuple): + edge_model, source_targets = edge_data + else: + edge_model = edge_data + source_targets = None + + edge_dict = edge_model_to_api_schema(edge_model, name) + if source_targets: + edge_dict["source_targets"] = [st.dict() for st in source_targets] + api_edge_types.append(EdgeType(**edge_dict)) + + res = await self.set_entity_types_internal( + entity_types=api_entity_types, + edge_types=api_edge_types, + user_ids=user_ids, + graph_ids=graph_ids, request_options=request_options, ) + return res diff --git a/src/zep_cloud/graph/utils.py b/src/zep_cloud/graph/utils.py index 35f4d6f2..384a878f 100644 --- a/src/zep_cloud/graph/utils.py +++ b/src/zep_cloud/graph/utils.py @@ -2,12 +2,12 @@ from typing import List, Optional from dateutil import parser as dateutil_parser -from zep_cloud import Edge, Episode, Node +from zep_cloud import EntityEdge, EntityNode, Episode DATE_FORMAT = "%Y-%m-%d %H:%M:%S" -def parse_iso_datetime(iso_string: Optional[str]) -> Optional[datetime]: +def parse_iso_datetime(iso_string: str) -> Optional[datetime]: """Parse ISO datetime string using dateutil parser.""" if not iso_string: return None @@ -39,7 +39,7 @@ def parse_iso_datetime(iso_string: Optional[str]) -> Optional[datetime]: """ -def format_edge_date_range(edge: Edge) -> str: +def format_edge_date_range(edge: EntityEdge) -> str: """ Format the date range of an entity edge. @@ -64,7 +64,7 @@ def format_edge_date_range(edge: Edge) -> str: return f"{valid_at} - {invalid_at}" -def compose_context_string(edges: List[Edge], nodes: List[Node], episodes: List[Episode]) -> str: +def compose_context_string(edges: List[EntityEdge], nodes: List[EntityNode], episodes: List[Episode]) -> str: """ Compose a search context from entity edges, nodes, and episodes. @@ -110,19 +110,14 @@ def compose_context_string(edges: List[Edge], nodes: List[Node], episodes: List[ episodes_list = [] if episodes: for episode in episodes: - # spec-3 8.4 resolves v3's overloaded pair: role carries the enum - # v3 spelled role_type, and role_name carries the sender name v3 - # spelled role. The rendered prefix is unchanged. - role_name = getattr(episode, "role_name", None) - role_type = getattr(episode, "role", None) - role_prefix = "" - if role_name and role_type: - role_prefix = f"{role_name} ({role_type}): " - elif role_name: - role_prefix = f"{role_name}: " - elif role_type: - role_prefix = f"({role_type}): " + if hasattr(episode, 'role') and episode.role: + if hasattr(episode, 'role_type') and episode.role_type: + role_prefix = f"{episode.role} ({episode.role_type}): " + else: + role_prefix = f"{episode.role}: " + elif hasattr(episode, 'role_type') and episode.role_type: + role_prefix = f"({episode.role_type}): " parsed_timestamp = parse_iso_datetime(episode.created_at) timestamp = parsed_timestamp.strftime(DATE_FORMAT) if parsed_timestamp is not None else "date unknown" diff --git a/tests/graph/test_utils.py b/tests/graph/test_utils.py index 909d2a00..7ed919cb 100644 --- a/tests/graph/test_utils.py +++ b/tests/graph/test_utils.py @@ -3,13 +3,13 @@ import pytest -from zep_cloud import Edge, Episode, Node +from zep_cloud import EntityEdge, EntityNode, Episode from zep_cloud.graph.utils import compose_context_string, format_edge_date_range class TestFormatEdgeDateRange: def test_format_edge_date_range_with_valid_dates(self): - edge = Edge( + edge = EntityEdge( fact="Test fact", name="test_edge", uuid_="edge-123", @@ -23,7 +23,7 @@ def test_format_edge_date_range_with_valid_dates(self): assert result == "2024-01-01 10:00:00 - 2024-01-02 10:00:00" def test_format_edge_date_range_with_none_dates(self): - edge = Edge( + edge = EntityEdge( fact="Test fact", name="test_edge", uuid_="edge-123", @@ -37,7 +37,7 @@ def test_format_edge_date_range_with_none_dates(self): assert result == "date unknown - present" def test_format_edge_date_range_with_partial_dates(self): - edge = Edge( + edge = EntityEdge( fact="Test fact", name="test_edge", uuid_="edge-123", @@ -60,7 +60,7 @@ def test_empty_inputs(self): assert "EPISODES" not in result def test_facts_only(self): - edge = Edge( + edge = EntityEdge( fact="User likes pizza", name="likes", uuid_="edge-123", @@ -78,7 +78,7 @@ def test_facts_only(self): assert "EPISODES" not in result def test_entities_basic(self): - node = Node( + node = EntityNode( name="John", summary="A user", uuid_="node-123", @@ -91,7 +91,7 @@ def test_entities_basic(self): assert "" in result def test_entities_with_label_and_attributes(self): - node = Node( + node = EntityNode( name="John", summary="A user", uuid_="node-123", @@ -109,7 +109,7 @@ def test_entities_with_label_and_attributes(self): assert "Summary: A user" in result def test_entities_with_entity_label_removed(self): - node = Node( + node = EntityNode( name="Alice", summary="A customer", uuid_="node-456", @@ -124,7 +124,7 @@ def test_entities_with_entity_label_removed(self): assert "Summary: A customer" in result def test_entities_with_only_entity_label(self): - node = Node( + node = EntityNode( name="Bob", summary="A person", uuid_="node-789", @@ -140,7 +140,7 @@ def test_entities_with_only_entity_label(self): assert "Summary: A person" in result def test_entities_with_labels_attribute_filtered(self): - node = Node( + node = EntityNode( name="stores", summary="Physical locations for shopping", uuid_="node-123", @@ -174,34 +174,33 @@ def test_episodes_with_role(self): content="Hello there!", created_at="2024-01-01T10:00:00Z", uuid_="episode-123", - role_name="user" + role="user" ) result = compose_context_string([], [], [episode]) assert "user: Hello there! (2024-01-01 10:00:00)" in result - def test_episodes_with_sender_name_and_role(self): - # spec-3 8.4: role_name is the sender, role is the enum v3 called - # role_type. + def test_episodes_with_role_and_type(self): + # Create a mock episode with role_type since Episode model uses enum class MockEpisode: def __init__(self): self.content = "Hello there!" self.created_at = "2024-01-01T10:00:00Z" - self.role_name = "assistant" - self.role = "ai" + self.role = "assistant" + self.role_type = "ai" episode = MockEpisode() result = compose_context_string([], [], [episode]) assert "assistant (ai): Hello there! (2024-01-01 10:00:00)" in result - def test_episodes_with_role_only(self): + def test_episodes_with_role_type_only(self): class MockEpisode: def __init__(self): self.content = "Hello there!" self.created_at = "2024-01-01T10:00:00Z" - self.role_name = None - self.role = "system" + self.role = None + self.role_type = "system" episode = MockEpisode() result = compose_context_string([], [], [episode]) @@ -210,7 +209,7 @@ def __init__(self): def test_complete_context_with_all_elements(self): - edge = Edge( + edge = EntityEdge( fact="User prefers coffee", name="prefers", uuid_="edge-123", @@ -221,7 +220,7 @@ def test_complete_context_with_all_elements(self): invalid_at=None ) - node = Node( + node = EntityNode( name="Alice", summary="Regular customer", uuid_="node-123", @@ -234,8 +233,8 @@ class MockEpisode: def __init__(self): self.content = "I'd like my usual coffee" self.created_at = "2024-01-01T09:00:00Z" - self.role_name = "user" - self.role = "customer" + self.role = "user" + self.role_type = "customer" episode = MockEpisode() @@ -256,7 +255,7 @@ def __init__(self): def test_multiple_items(self): edges = [ - Edge( + EntityEdge( fact="Fact 1", name="edge1", uuid_="edge-1", @@ -265,7 +264,7 @@ def test_multiple_items(self): target_node_uuid="target-1", valid_at="2024-01-01T10:00:00Z" ), - Edge( + EntityEdge( fact="Fact 2", name="edge2", uuid_="edge-2", @@ -277,13 +276,13 @@ def test_multiple_items(self): ] nodes = [ - Node( + EntityNode( name="Node1", summary="Summary 1", uuid_="node-1", created_at="2024-01-01T09:00:00Z" ), - Node( + EntityNode( name="Node2", summary="Summary 2", uuid_="node-2",