Customizing Graph Structure

1
from zep_cloud.types import Message
2
3
message = {"role": "John Doe", "role_type": "user", "content": "I really like pop music, and I don't like metal"}
4
5
client.memory.add(session_id=session_id, messages=[Message(**message)])

1
search_results = client.graph.search(
2
    user_id=user_id,
3
    query="the user's music preferences",
4
    scope="nodes",
5
    search_filters={
6
        "node_labels": ["Preference"]
7
    }
8
)
9
for i, node in enumerate(search_results.nodes):
10
    preference = node.attributes
11
    print(f"Preference {i+1}:{preference}")

1
from zep_cloud.external_clients.ontology import EntityModel, EntityText, EdgeModel, EntityBoolean
2
from pydantic import Field
3
4
class Restaurant(EntityModel):
5
    """
6
    Represents a specific restaurant.
7
    """
8
    cuisine_type: EntityText = Field(description="The cuisine type of the restaurant, for example: American, Mexican, Indian, etc.", default=None)
9
    dietary_accommodation: EntityText = Field(description="The dietary accommodation of the restaurant, if any, for example: vegetarian, vegan, etc.", default=None)
10
11
class Audiobook(EntityModel):
12
    """
13
    Represents an audiobook entity.
14
    """
15
    genre: EntityText = Field(description="The genre of the audiobook, for example: self-help, fiction, nonfiction, etc.", default=None)
16
17
class RestaurantVisit(EdgeModel):
18
    """
19
    Represents the fact that the user visited a restaurant.
20
    """
21
    restaurant_name: EntityText = Field(description="The name of the restaurant the user visited", default=None)
22
23
class AudiobookListen(EdgeModel):
24
    """
25
    Represents the fact that the user listened to or played an audiobook.
26
    """
27
    audiobook_title: EntityText = Field(description="The title of the audiobook the user listened to or played", default=None)
28
29
class DietaryPreference(EdgeModel):
30
    """
31
    Represents the fact that the user has a dietary preference or dietary restriction.
32
    """
33
    preference_type: EntityText = Field(description="Preference type of the user: anything, vegetarian, vegan, peanut allergy, etc.", default=None)
34
    allergy: EntityBoolean = Field(description="Whether this dietary preference represents a user allergy: True or false", default=None)

1
from zep_cloud import EntityEdgeSourceTarget
2
3
client.graph.set_ontology(
4
    entities={
5
        "Restaurant": Restaurant,
6
        "Audiobook": Audiobook,
7
    },
8
    edges={
9
        "RESTAURANT_VISIT": (
10
            RestaurantVisit,
11
            [EntityEdgeSourceTarget(source="User", target="Restaurant")]
12
        ),
13
        "AUDIOBOOK_LISTEN": (
14
            AudiobookListen,
15
            [EntityEdgeSourceTarget(source="User", target="Audiobook")]
16
        ),
17
        "DIETARY_PREFERENCE": (
18
            DietaryPreference,
19
            [EntityEdgeSourceTarget(source="User")]
20
        ),
21
    }
22
)

1
from zep_cloud import Message
2
import uuid
3
4
messages_session1 = [
5
    Message(content="Take me to a lunch place", role_type="user", role="John Doe"),
6
    Message(content="How about Panera Bread, Chipotle, or Green Leaf Cafe, which are nearby?", role_type="assistant", role="Assistant"),
7
    Message(content="Do any of those have vegetarian options? I’m vegetarian", role_type="user", role="John Doe"),
8
    Message(content="Yes, Green Leaf Cafe has vegetarian options", role_type="assistant", role="Assistant"),
9
    Message(content="Let’s go to Green Leaf Cafe", role_type="user", role="John Doe"),
10
    Message(content="Navigating to Green Leaf Cafe", role_type="assistant", role="Assistant"),
11
]
12
13
messages_session2 = [
14
    Message(content="Play the 7 habits of highly effective people", role_type="user", role="John Doe"),
15
    Message(content="Playing the 7 habits of highly effective people", role_type="assistant", role="Assistant"),
16
]
17
18
user_id = f"user-{uuid.uuid4()}"
19
client.user.add(user_id=user_id, first_name="John", last_name="Doe", email="[email protected]")
20
21
session1_id = f"session-{uuid.uuid4()}"
22
session2_id = f"session-{uuid.uuid4()}"
23
client.memory.add_session(session_id=session1_id, user_id=user_id)
24
client.memory.add_session(session_id=session2_id, user_id=user_id)
25
26
client.memory.add(session_id=session1_id, messages=messages_session1, ignore_roles=["assistant"])
27
client.memory.add(session_id=session2_id, messages=messages_session2, ignore_roles=["assistant"])

1
search_results_restaurants = client.graph.search(
2
    user_id=user_id,
3
    query="Take me to a restaurant",
4
    scope="nodes",
5
    search_filters={
6
        "node_labels": ["Restaurant"]
7
    },
8
    limit=1,
9
)
10
node = search_results_restaurants.nodes[0]
11
print(f"Node name: {node.name}")
12
print(f"Node labels: {node.labels}")
13
print(f"Cuisine type: {node.attributes.get('cuisine_type')}")
14
print(f"Dietary accommodation: {node.attributes.get('dietary_accommodation')}")

1
search_results_audiobook_nodes = client.graph.search(
2
    user_id=user_id,
3
    query="Play an audiobook",
4
    scope="nodes",
5
    search_filters={
6
        "node_labels": ["Audiobook"]
7
    },
8
    limit=1,
9
)
10
node = search_results_audiobook_nodes.nodes[0]
11
print(f"Node name: {node.name}")
12
print(f"Node labels: {node.labels}")
13
print(f"Genre: {node.attributes.get('genre')}")

1
search_results_visits = client.graph.search(
2
    user_id=user_id,
3
    query="Take me to a restaurant",
4
    scope="edges",
5
    search_filters={
6
        "edge_types": ["RESTAURANT_VISIT"]
7
    },
8
    limit=1,
9
)
10
edge = search_results_visits.edges[0]
11
print(f"Edge fact: {edge.fact}")
12
print(f"Edge type: {edge.name}")
13
print(f"Restaurant name: {edge.attributes.get('restaurant_name')}")

1
search_results_audiobook_listens = client.graph.search(
2
    user_id=user_id,
3
    query="Play an audiobook",
4
    scope="edges",
5
    search_filters={
6
        "edge_types": ["AUDIOBOOK_LISTEN"]
7
    },
8
    limit=1,
9
)
10
edge = search_results_audiobook_listens.edges[0]
11
print(f"Edge fact: {edge.fact}")
12
print(f"Edge type: {edge.name}")
13
print(f"Audiobook title: {edge.attributes.get('audiobook_title')}")

1
search_results_dietary_preference = client.graph.search(
2
    user_id=user_id,
3
    query="Find some food around here",
4
    scope="edges",
5
    search_filters={
6
        "edge_types": ["DIETARY_PREFERENCE"]
7
    },
8
    limit=1,
9
)
10
edge = search_results_dietary_preference.edges[0]
11
print(f"Edge fact: {edge.fact}")
12
print(f"Edge type: {edge.name}")
13
print(f"Preference type: {edge.attributes.get('preference_type')}")
14
print(f"Allergy: {edge.attributes.get('allergy')}")

1
search_results_dietary_preference = client.graph.search(
2
    user_id=user_id,
3
    query="Find some food around here",
4
    scope="edges",
5
    search_filters={
6
        "edge_types": ["DIETARY_PREFERENCE", "RESTAURANT_VISIT"]
7
    },
8
    limit=2,
9
)
10
for edge in search_results_dietary_preference.edges:
11
    print(f"Edge fact: {edge.fact}")
12
    print(f"Edge type: {edge.name}")
13
    if edge.name == "DIETARY_PREFERENCE":
14
        print(f"Preference type: {edge.attributes.get('preference_type')}")
15
        print(f"Allergy: {edge.attributes.get('allergy')}")
16
    elif edge.name == "RESTAURANT_VISIT":
17
        print(f"Restaurant name: {edge.attributes.get('restaurant_name')}")
18
    print("\n")