Using LangGraph and Graphiti

1
import asyncio
2
import json
3
import logging
4
import os
5
import sys
6
import uuid
7
from contextlib import suppress
8
from datetime import datetime
9
from pathlib import Path
10
from typing import Annotated
11
12
import ipywidgets as widgets
13
from dotenv import load_dotenv
14
from IPython.display import Image, display
15
from typing_extensions import TypedDict
16
17
load_dotenv()

1
def setup_logging():
2
    logger = logging.getLogger()
3
    logger.setLevel(logging.ERROR)
4
    console_handler = logging.StreamHandler(sys.stdout)
5
    console_handler.setLevel(logging.INFO)
6
    formatter = logging.Formatter('%(name)s - %(levelname)s - %(message)s')
7
    console_handler.setFormatter(formatter)
8
    logger.addHandler(console_handler)
9
    return logger
10
11
12
logger = setup_logging()

1
# Configure Graphiti
2
3
from graphiti_core import Graphiti
4
from graphiti_core.edges import EntityEdge
5
from graphiti_core.nodes import EpisodeType
6
from graphiti_core.utils.bulk_utils import RawEpisode
7
from graphiti_core.utils.maintenance.graph_data_operations import clear_data
8
9
neo4j_uri = os.environ.get('NEO4J_URI', 'bolt://localhost:7687')
10
neo4j_user = os.environ.get('NEO4J_USER', 'neo4j')
11
neo4j_password = os.environ.get('NEO4J_PASSWORD', 'password')
12
13
client = Graphiti(
14
    neo4j_uri,
15
    neo4j_user,
16
    neo4j_password,
17
)

1
# Note: This will clear the database
2
await clear_data(client.driver)
3
await client.build_indices_and_constraints()

1
async def ingest_products_data(client: Graphiti):
2
    script_dir = Path.cwd().parent
3
    json_file_path = script_dir / 'data' / 'manybirds_products.json'
4
5
    with open(json_file_path) as file:
6
        products = json.load(file)['products']
7
8
    episodes: list[RawEpisode] = [
9
        RawEpisode(
10
            name=product.get('title', f'Product {i}'),
11
            content=str({k: v for k, v in product.items() if k != 'images'}),
12
            source_description='ManyBirds products',
13
            source=EpisodeType.json,
14
            reference_time=datetime.now(),
15
        )
16
        for i, product in enumerate(products)
17
    ]
18
19
    await client.add_episode_bulk(episodes)
20
21
22
await ingest_products_data(client)

1
user_name = 'jess'
2
3
await client.add_episode(
4
    name='User Creation',
5
    episode_body=(f'{user_name} is interested in buying a pair of shoes'),
6
    source=EpisodeType.text,
7
    reference_time=datetime.now(),
8
    source_description='SalesBot',
9
)
10
11
# let's get Jess's node uuid
12
nl = await client.get_nodes_by_query(user_name)
13
14
user_node_uuid = nl[0].uuid
15
16
# and the ManyBirds node uuid
17
nl = await client.get_nodes_by_query('ManyBirds')
18
manybirds_node_uuid = nl[0].uuid

1
def edges_to_facts_string(entities: list[EntityEdge]):
2
    return '-' + '\n- '.join([edge.fact for edge in entities])

1
from langchain_core.messages import AIMessage, SystemMessage
2
from langchain_core.tools import tool
3
from langchain_openai import ChatOpenAI
4
from langgraph.checkpoint.memory import MemorySaver
5
from langgraph.graph import END, START, StateGraph, add_messages
6
from langgraph.prebuilt import ToolNode

1
@tool
2
async def get_shoe_data(query: str) -> str:
3
    """Search the graphiti graph for information about shoes"""
4
    edge_results = await client.search(
5
        query,
6
        center_node_uuid=manybirds_node_uuid,
7
        num_results=10,
8
    )
9
    return edges_to_facts_string(edge_results)
10
11
12
tools = [get_shoe_data]
13
tool_node = ToolNode(tools)

1
llm = ChatOpenAI(model='gpt-4o-mini', temperature=0).bind_tools(tools)

1
await tool_node.ainvoke({'messages': [await llm.ainvoke('wool shoes')]})

1
class State(TypedDict):
2
    messages: Annotated[list, add_messages]
3
    user_name: str
4
    user_node_uuid: str
5
6
7
async def chatbot(state: State):
8
    facts_string = None
9
    if len(state['messages']) > 0:
10
        last_message = state['messages'][-1]
11
        graphiti_query = f'{"SalesBot" if isinstance(last_message, AIMessage) else state["user_name"]}: {last_message.content}'
12
        # search graphiti using Jess's node uuid as the center node
13
        # graph edges (facts) further from the Jess node will be ranked lower
14
        edge_results = await client.search(
15
            graphiti_query, center_node_uuid=state['user_node_uuid'], num_results=5
16
        )
17
        facts_string = edges_to_facts_string(edge_results)
18
19
    system_message = SystemMessage(
20
        content=f"""You are a skillfull shoe salesperson working for ManyBirds. Review information about the user and their prior conversation below and respond accordingly.
21
        Keep responses short and concise. And remember, always be selling (and helpful!)
22
23
        Things you'll need to know about the user in order to close a sale:
24
        - the user's shoe size
25
        - any other shoe needs? maybe for wide feet?
26
        - the user's preferred colors and styles
27
        - their budget
28
29
        Ensure that you ask the user for the above if you don't already know.
30
31
        Facts about the user and their conversation:
32
        {facts_string or 'No facts about the user and their conversation'}"""
33
    )
34
35
    messages = [system_message] + state['messages']
36
37
    response = await llm.ainvoke(messages)
38
39
    # add the response to the graphiti graph.
40
    # this will allow us to use the graphiti search later in the conversation
41
    # we're doing async here to avoid blocking the graph execution
42
    asyncio.create_task(
43
        client.add_episode(
44
            name='Chatbot Response',
45
            episode_body=f"{state['user_name']}: {state['messages'][-1]}\nSalesBot: {response.content}",
46
            source=EpisodeType.message,
47
            reference_time=datetime.now(),
48
            source_description='Chatbot',
49
        )
50
    )
51
52
    return {'messages': [response]}

1
graph_builder = StateGraph(State)
2
3
memory = MemorySaver()
4
5
6
# Define the function that determines whether to continue or not
7
async def should_continue(state, config):
8
    messages = state['messages']
9
    last_message = messages[-1]
10
    # If there is no function call, then we finish
11
    if not last_message.tool_calls:
12
        return 'end'
13
    # Otherwise if there is, we continue
14
    else:
15
        return 'continue'
16
17
18
graph_builder.add_node('agent', chatbot)
19
graph_builder.add_node('tools', tool_node)
20
21
graph_builder.add_edge(START, 'agent')
22
graph_builder.add_conditional_edges('agent', should_continue, {'continue': 'tools', 'end': END})
23
graph_builder.add_edge('tools', 'agent')
24
25
26
graph = graph_builder.compile(checkpointer=memory)

1
with suppress(Exception):
2
    display(Image(graph.get_graph().draw_mermaid_png()))

1
await graph.ainvoke(
2
    {
3
        'messages': [
4
            {
5
                'role': 'user',
6
                'content': 'What sizes do the TinyBirds Wool Runners in Natural Black come in?',
7
            }
8
        ],
9
        'user_name': user_name,
10
        'user_node_uuid': user_node_uuid,
11
    },
12
    config={'configurable': {'thread_id': uuid.uuid4().hex}},
13
)

1
conversation_output = widgets.Output()
2
config = {'configurable': {'thread_id': uuid.uuid4().hex}}
3
user_state = {'user_name': user_name, 'user_node_uuid': user_node_uuid}
4
5
6
async def process_input(user_state: State, user_input: str):
7
    conversation_output.append_stdout(f'\nUser: {user_input}\n')
8
    conversation_output.append_stdout('\nAssistant: ')
9
10
    graph_state = {
11
        'messages': [{'role': 'user', 'content': user_input}],
12
        'user_name': user_state['user_name'],
13
        'user_node_uuid': user_state['user_node_uuid'],
14
    }
15
16
    try:
17
        async for event in graph.astream(
18
            graph_state,
19
            config=config,
20
        ):
21
            for value in event.values():
22
                if 'messages' in value:
23
                    last_message = value['messages'][-1]
24
                    if isinstance(last_message, AIMessage) and isinstance(
25
                        last_message.content, str
26
                    ):
27
                        conversation_output.append_stdout(last_message.content)
28
    except Exception as e:
29
        conversation_output.append_stdout(f'Error: {e}')
30
31
32
def on_submit(b):
33
    user_input = input_box.value
34
    input_box.value = ''
35
    asyncio.create_task(process_input(user_state, user_input))
36
37
38
input_box = widgets.Text(placeholder='Type your message here...')
39
submit_button = widgets.Button(description='Send')
40
submit_button.on_click(on_submit)
41
42
conversation_output.append_stdout('Asssistant: Hello, how can I help you find shoes today?')
43
44
display(widgets.VBox([input_box, submit_button, conversation_output]))