136 lines
4.2 KiB
Python
136 lines
4.2 KiB
Python
import json
|
|
import sys
|
|
from pathlib import Path
|
|
from datetime import datetime
|
|
from sentence_transformers import SentenceTransformer
|
|
import chromadb
|
|
|
|
# Paths
|
|
db_path = str(Path.home() / "aaronai" / "db")
|
|
EXPORT_DIR = "/home/aaron/nextcloud/data/data/aaron/files/Archive/Misc/Claude Export"
|
|
|
|
print("Loading embedding model...")
|
|
embedder = SentenceTransformer("all-MiniLM-L6-v2")
|
|
client = chromadb.PersistentClient(path=db_path)
|
|
collection = client.get_or_create_collection(
|
|
name="aaronai",
|
|
metadata={"hnsw:space": "cosine"}
|
|
)
|
|
|
|
def extract_messages(convo):
|
|
"""Extract messages from a Claude conversation object."""
|
|
messages = []
|
|
for msg in convo.get("chat_messages", []):
|
|
role = msg.get("sender", "")
|
|
if role not in ["human", "assistant"]:
|
|
continue
|
|
# Claude export stores content as a list of content blocks
|
|
content = msg.get("content", [])
|
|
text = ""
|
|
if isinstance(content, str):
|
|
text = content
|
|
elif isinstance(content, list):
|
|
for block in content:
|
|
if isinstance(block, dict) and block.get("type") == "text":
|
|
text += block.get("text", "")
|
|
elif isinstance(block, str):
|
|
text += block
|
|
text = text.strip()
|
|
if not text:
|
|
continue
|
|
created_at = msg.get("created_at", "")
|
|
messages.append((created_at, role, text))
|
|
return messages
|
|
|
|
def chunk_conversation(convo):
|
|
"""Turn a conversation into indexable chunks."""
|
|
chunks = []
|
|
title = convo.get("name", "Untitled conversation")
|
|
uuid = convo.get("uuid", "")
|
|
created_at = convo.get("created_at", "")
|
|
messages = extract_messages(convo)
|
|
if not messages:
|
|
return chunks
|
|
|
|
# Chunk into sliding windows of 3 messages
|
|
window = []
|
|
for i, (ts, role, text) in enumerate(messages):
|
|
label = "You" if role == "human" else "Claude"
|
|
window.append(f"{label}: {text}")
|
|
if len(window) >= 3 or i == len(messages) - 1:
|
|
chunk_text = f"[Claude conversation: {title}]\n\n" + "\n\n".join(window)
|
|
chunk_id = f"claude_{uuid}_{i}"
|
|
chunks.append((chunk_id, chunk_text, {
|
|
"source": f"Claude: {title}",
|
|
"type": "claude_conversation",
|
|
"created_at": created_at,
|
|
}))
|
|
window = window[-1:] # overlap by 1
|
|
|
|
return chunks
|
|
|
|
def ingest_file(jsonl_path):
|
|
print(f"Processing {jsonl_path.name}...")
|
|
conversations = []
|
|
with open(jsonl_path, encoding="utf-8") as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
try:
|
|
conversations.append(json.loads(line))
|
|
except json.JSONDecodeError:
|
|
continue
|
|
|
|
print(f"Found {len(conversations)} conversations")
|
|
total_chunks = 0
|
|
skipped = 0
|
|
|
|
for convo in conversations:
|
|
chunks = chunk_conversation(convo)
|
|
if not chunks:
|
|
skipped += 1
|
|
continue
|
|
|
|
ids = [c[0] for c in chunks]
|
|
texts = [c[1] for c in chunks]
|
|
metas = [c[2] for c in chunks]
|
|
|
|
# Check existing
|
|
existing = collection.get(ids=ids)
|
|
existing_ids = set(existing["ids"])
|
|
new = [(id, txt, meta) for id, txt, meta in zip(ids, texts, metas) if id not in existing_ids]
|
|
|
|
if not new:
|
|
continue
|
|
|
|
embeddings = embedder.encode([n[1] for n in new]).tolist()
|
|
collection.add(
|
|
ids=[n[0] for n in new],
|
|
documents=[n[1] for n in new],
|
|
metadatas=[n[2] for n in new],
|
|
embeddings=embeddings,
|
|
)
|
|
total_chunks += len(new)
|
|
|
|
print(f"Done. {total_chunks} chunks added, {skipped} conversations skipped.")
|
|
return total_chunks
|
|
|
|
# Find the export file
|
|
export_dir = Path(EXPORT_DIR)
|
|
export_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
jsonl_files = list(export_dir.glob("*.jsonl")) + list(export_dir.glob("**/*.jsonl"))
|
|
|
|
if not jsonl_files:
|
|
print(f"No .jsonl files found in {EXPORT_DIR}")
|
|
print("Place your Claude export conversations.jsonl file there and run again.")
|
|
sys.exit(0)
|
|
|
|
total = 0
|
|
for f in jsonl_files:
|
|
total += ingest_file(f)
|
|
|
|
print(f"\nTotal chunks added to corpus: {total}")
|
|
print(f"Database at: {db_path}")
|