graphiti/runner.py

import asyncio
import logging
import os
import sys
from datetime import datetime

from dotenv import load_dotenv

from core import Graphiti
from core.utils.maintenance.graph_data_operations import clear_data

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 3',
		episode_body='Jane: I am married to Paul',
		source_description='WhatsApp Message',
		reference_time=datetime.now(),
	)
	await client.add_episode(
		name='Message 4',
		episode_body='Paul: I have divorced Jane',
		source_description='WhatsApp Message',
		reference_time=datetime.now(),
	)
	await client.add_episode(
		name='Message 5',
		episode_body='Jane: I  miss Paul',
		source_description='WhatsApp Message',
		reference_time=datetime.now(),
	)
	await client.add_episode(
		name='Message 6',
		episode_body='Jane: I dont miss Paul anymore, I hate him',
		source_description='WhatsApp Message',
		reference_time=datetime.now(),
	)

	# 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())