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 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",