graphiti/runner.py
Pavlo Paliychuk 40e74a2e97
fix: Address graph disconnect (#7)
* fix: Address graph disconnect

* chore: Remove valid_to and valid_from setting in extract edges step (will be handled during invalidation step)
2024-08-19 09:37:56 -04:00

68 lines
1.8 KiB
Python

from core import Graphiti
from core.utils.maintenance.graph_data_operations import clear_data
from dotenv import load_dotenv
import os
import asyncio
import logging
import sys
load_dotenv()
neo4j_uri = os.environ.get("NEO4J_URI") or "bolt://localhost:7687"
neo4j_user = os.environ.get("NEO4J_USER") or "neo4j"
neo4j_password = os.environ.get("NEO4J_PASSWORD") or "password"
def setup_logging():
# Create a logger
logger = logging.getLogger()
logger.setLevel(logging.INFO) # Set the logging level to INFO
# Create console handler and set level to INFO
console_handler = logging.StreamHandler(sys.stdout)
console_handler.setLevel(logging.INFO)
# Create formatter
formatter = logging.Formatter(
"%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
# Add formatter to console handler
console_handler.setFormatter(formatter)
# Add console handler to logger
logger.addHandler(console_handler)
return logger
async def main():
setup_logging()
client = Graphiti(neo4j_uri, neo4j_user, neo4j_password)
await clear_data(client.driver)
# await client.build_indices()
await client.add_episode(
name="Message 1",
episode_body="Paul: I love apples",
source_description="WhatsApp Message",
)
await client.add_episode(
name="Message 2",
episode_body="Paul: I own many bananas",
source_description="WhatsApp Message",
)
await client.add_episode(
name="Message 3",
episode_body="Assistant: The best type of apples available are Fuji apples",
source_description="WhatsApp Message",
)
await client.add_episode(
name="Message 4",
episode_body="Paul: Oh, I actually hate those",
source_description="WhatsApp Message",
)
asyncio.run(main())