wehub-resource-sync
108 lines
3.2 KiB
Python
108 lines
3.2 KiB
Python
"""Ingest chunked data into Zep knowledge graph."""
|
|
import sys
|
|
import os
|
|
import json
|
|
import time
|
|
from pathlib import Path
|
|
from zep_cloud import Zep
|
|
from dotenv import load_dotenv
|
|
|
|
# Load environment variables
|
|
load_dotenv()
|
|
ZEP_API_KEY = os.getenv("ZEP_API_KEY")
|
|
if not ZEP_API_KEY:
|
|
raise ValueError("ZEP_API_KEY environment variable is required")
|
|
|
|
USER_ID = os.getenv("ZEP_DOCS_USER_ID")
|
|
|
|
|
|
def main():
|
|
"""Main function to ingest data to user graph."""
|
|
print("Initializing Zep client...")
|
|
zep = Zep(api_key=ZEP_API_KEY)
|
|
|
|
# Create or verify user exists
|
|
print(f"Setting up user: \"{USER_ID}\"...")
|
|
try:
|
|
zep.user.add(
|
|
user_id=USER_ID,
|
|
email="docs@system.local",
|
|
first_name="Zep",
|
|
last_name="KG"
|
|
)
|
|
print(f"User \"{USER_ID}\" created successfully")
|
|
except Exception as e:
|
|
if "already exists" in str(e).lower() or "409" in str(e):
|
|
print(f"User \"{USER_ID}\" already exists")
|
|
else:
|
|
print(f"Warning: Error creating user: {e}")
|
|
print("Continuing with ingestion...")
|
|
|
|
# Load chunked docs
|
|
print("Loading chunked docs...")
|
|
chunks_path = Path(__file__).parent.parent / "data" / "chunked-docs.json"
|
|
|
|
if not chunks_path.exists():
|
|
print(f"Error: {chunks_path} not found!")
|
|
print("Please create chunked docs first.")
|
|
return
|
|
|
|
with open(chunks_path, "r", encoding="utf-8") as f:
|
|
chunks = json.load(f)
|
|
|
|
print(f"Loaded {len(chunks)} chunks")
|
|
|
|
episode_uuids = []
|
|
|
|
for i, chunk in enumerate(chunks):
|
|
print(f"Ingesting chunk {i + 1}/{len(chunks)}: \"{chunk['heading']}\"")
|
|
|
|
MAX_CONTENT_LENGTH = 10000
|
|
header = f"{chunk['context']}\n\n# {chunk['heading']}\n\n"
|
|
available_space = MAX_CONTENT_LENGTH - len(header)
|
|
|
|
content = chunk["content"]
|
|
if len(content) > available_space:
|
|
print(
|
|
f" Warning: Chunk content exceeds available space ({len(content)} chars), truncating to {available_space} chars...")
|
|
content = content[:available_space - 20] + "\n\n[truncated...]"
|
|
|
|
episode_content = header + content
|
|
|
|
try:
|
|
episode = zep.graph.add(
|
|
user_id=USER_ID,
|
|
type="text",
|
|
data=episode_content,
|
|
)
|
|
|
|
episode_uuids.append(episode.uuid_)
|
|
print(
|
|
f" Created episode: {episode.uuid_} ({len(episode_content)} chars)")
|
|
except Exception as e:
|
|
print(f" Error ingesting chunk: {e}")
|
|
|
|
# Rate limiting
|
|
time.sleep(0.5)
|
|
|
|
# Save episode mapping
|
|
mapping_path = Path(__file__).parent.parent / \
|
|
"data" / "episode-mapping.json"
|
|
mapping_data = {
|
|
"user_id": USER_ID,
|
|
"total_episodes": len(episode_uuids),
|
|
"episode_ids": episode_uuids,
|
|
"created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
|
|
}
|
|
|
|
with open(mapping_path, "w", encoding="utf-8") as f:
|
|
json.dump(mapping_data, f, indent=2)
|
|
|
|
print(
|
|
f"\nDone! Ingested {len(episode_uuids)} episodes to user graph \"{USER_ID}\"")
|
|
print(f"Episode mapping saved to {mapping_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|