291 lines
8.8 KiB
Python
291 lines
8.8 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
RAG Session Ingestor - Load all session transcripts into vector store
|
|
Fixed to handle OpenClaw session event format
|
|
"""
|
|
|
|
import os
|
|
import json
|
|
from pathlib import Path
|
|
from datetime import datetime
|
|
from typing import List, Dict, Any
|
|
|
|
import sys
|
|
sys.path.insert(0, str(Path(__file__).parent))
|
|
|
|
from rag_system import RAGSystem
|
|
|
|
|
|
def parse_jsonl(file_path: Path) -> List[Dict]:
|
|
"""
|
|
Parse OpenClaw session JSONL format
|
|
|
|
Session files contain:
|
|
- Line 1: Session metadata (type: "session")
|
|
- Lines 2+: Events including messages, toolCalls, etc.
|
|
"""
|
|
messages = []
|
|
|
|
try:
|
|
with open(file_path, 'r', encoding='utf-8') as f:
|
|
for line_num, line in enumerate(f, 1):
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
|
|
try:
|
|
event = json.loads(line)
|
|
|
|
# Skip session metadata line
|
|
if line_num == 1 and event.get('type') == 'session':
|
|
continue
|
|
|
|
# Extract message events only
|
|
if event.get('type') == 'message':
|
|
msg_obj = event.get('message', {})
|
|
|
|
messages.append({
|
|
'role': msg_obj.get('role'),
|
|
'content': msg_obj.get('content'),
|
|
'timestamp': event.get('timestamp'),
|
|
'id': event.get('id'),
|
|
'sessionKey': event.get('sessionKey') # Not usually here, but check
|
|
})
|
|
|
|
except json.JSONDecodeError as e:
|
|
continue
|
|
except Exception as e:
|
|
print(f"❌ Error reading {file_path}: {e}")
|
|
|
|
return messages
|
|
|
|
|
|
def extract_session_key(file_name: str) -> str:
|
|
"""Extract session key from filename"""
|
|
return file_name.replace('.jsonl', '')
|
|
|
|
|
|
def extract_session_metadata(session_data: List[Dict], session_key: str) -> Dict:
|
|
"""Extract metadata from session messages"""
|
|
if not session_data:
|
|
return {}
|
|
|
|
first_msg = session_data[0]
|
|
last_msg = session_data[-1]
|
|
|
|
return {
|
|
"start_time": first_msg.get("timestamp"),
|
|
"end_time": last_msg.get("timestamp"),
|
|
"total_messages": len(session_data),
|
|
"has_system": any(msg.get("role") == "system" for msg in session_data),
|
|
"has_user": any(msg.get("role") == "user" for msg in session_data),
|
|
"has_assistant": any(msg.get("role") == "assistant" for msg in session_data),
|
|
}
|
|
|
|
|
|
def format_content(content) -> str:
|
|
"""
|
|
Format message content from OpenClaw format to text
|
|
|
|
Content can be:
|
|
- String
|
|
- List of dicts with 'type' field (text, thinking, toolCall, toolResult)
|
|
"""
|
|
if isinstance(content, str):
|
|
return content
|
|
|
|
if isinstance(content, list):
|
|
texts = []
|
|
|
|
for item in content:
|
|
if not isinstance(item, dict):
|
|
continue
|
|
|
|
item_type = item.get('type', '')
|
|
|
|
if item_type == 'text':
|
|
texts.append(item.get('text', ''))
|
|
elif item_type == 'thinking':
|
|
# Skip reasoning, usually not useful for RAG
|
|
# texts.append(f"[Reasoning: {item.get('thinking', '')[:200]}]")
|
|
pass
|
|
elif item_type == 'toolCall':
|
|
tool_name = item.get('name', 'unknown')
|
|
args = str(item.get('arguments', ''))[:100]
|
|
texts.append(f"[Tool: {tool_name}({args})]")
|
|
elif item_type == 'toolResult':
|
|
result = str(item.get('text', item.get('result', ''))).strip()
|
|
# Truncate large tool results
|
|
if len(result) > 500:
|
|
result = result[:500] + "..."
|
|
texts.append(f"[Tool Result: {result}]")
|
|
|
|
return "\n".join(texts)
|
|
|
|
return str(content)[:500]
|
|
|
|
|
|
def chunk_messages(
|
|
messages: List[Dict],
|
|
session_key: str,
|
|
context_window: int = 20,
|
|
overlap: int = 5
|
|
) -> List[Dict]:
|
|
"""
|
|
Chunk messages for better retrieval
|
|
|
|
Args:
|
|
messages: List of message objects
|
|
session_key: Session identifier for metadata
|
|
context_window: Messages per chunk
|
|
overlap: Message overlap between chunks
|
|
|
|
Returns:
|
|
List of {"text": str, "metadata": dict} chunks
|
|
"""
|
|
chunks = []
|
|
|
|
for i in range(0, len(messages), context_window - overlap):
|
|
chunk_messages = messages[i:i + context_window]
|
|
|
|
# Build text from messages
|
|
text_parts = []
|
|
|
|
for msg in chunk_messages:
|
|
role = msg.get("role", "unknown")
|
|
content = msg.get("content", "")
|
|
|
|
# Format content
|
|
text = format_content(content)
|
|
|
|
if text.strip():
|
|
text_parts.append(f"{role.upper()}: {text}")
|
|
|
|
text = "\n\n".join(text_parts)
|
|
|
|
# Don't add empty chunks
|
|
if not text.strip():
|
|
continue
|
|
|
|
# Metadata
|
|
metadata = {
|
|
"type": "session",
|
|
"source": str(chunk_messages[0].get("sessionKey") or chunk_messages[0].get("id") or session_key),
|
|
"chunk_index": int(i // (context_window - overlap)),
|
|
"chunk_start_time": str(chunk_messages[0].get("timestamp") or ""),
|
|
"chunk_end_time": str(chunk_messages[-1].get("timestamp") or ""),
|
|
"message_count": int(len(chunk_messages)),
|
|
"ingested_at": datetime.now().isoformat(),
|
|
"date": str(chunk_messages[0].get("timestamp") or datetime.now().isoformat())
|
|
}
|
|
|
|
chunks.append({
|
|
"text": text,
|
|
"metadata": metadata
|
|
})
|
|
|
|
return chunks
|
|
|
|
|
|
def ingest_sessions(
|
|
sessions_dir: str = None,
|
|
collection_name: str = "openclaw_knowledge",
|
|
chunk_size: int = 20,
|
|
chunk_overlap: int = 5
|
|
):
|
|
"""
|
|
Ingest all session transcripts into RAG system
|
|
|
|
Args:
|
|
sessions_dir: Directory containing session jsonl files
|
|
collection_name: Name of the ChromaDB collection
|
|
chunk_size: Messages per chunk
|
|
chunk_overlap: Message overlap between chunks
|
|
"""
|
|
if sessions_dir is None:
|
|
sessions_dir = os.path.expanduser("~/.openclaw/agents/main/sessions")
|
|
|
|
sessions_path = Path(sessions_dir)
|
|
|
|
if not sessions_path.exists():
|
|
print(f"❌ Sessions directory not found: {sessions_path}")
|
|
return
|
|
|
|
print(f"🔍 Finding session files in: {sessions_path}")
|
|
|
|
jsonl_files = list(sessions_path.glob("*.jsonl"))
|
|
|
|
if not jsonl_files:
|
|
print(f"⚠️ No jsonl files found in {sessions_path}")
|
|
return
|
|
|
|
print(f"✅ Found {len(jsonl_files)} session files\n")
|
|
|
|
rag = RAGSystem(collection_name=collection_name)
|
|
|
|
total_chunks = 0
|
|
total_messages = 0
|
|
skipped_empty = 0
|
|
|
|
for jsonl_file in sorted(jsonl_files):
|
|
session_key = extract_session_key(jsonl_file.name)
|
|
|
|
print(f"\n📄 Processing: {jsonl_file.name}")
|
|
|
|
messages = parse_jsonl(jsonl_file)
|
|
|
|
if not messages:
|
|
print(f" ⚠️ No messages, skipping")
|
|
skipped_empty += 1
|
|
continue
|
|
|
|
total_messages += len(messages)
|
|
|
|
# Extract session metadata
|
|
session_metadata = extract_session_metadata(messages, session_key)
|
|
print(f" Messages: {len(messages)}")
|
|
|
|
# Chunk messages
|
|
chunks = chunk_messages(messages, session_key, chunk_size, chunk_overlap)
|
|
|
|
if not chunks:
|
|
print(f" ⚠️ No valid chunks, skipping")
|
|
skipped_empty += 1
|
|
continue
|
|
|
|
print(f" Chunks: {len(chunks)}")
|
|
|
|
# Add to RAG
|
|
try:
|
|
ids = rag.add_documents_batch(chunks, batch_size=50)
|
|
total_chunks += len(chunks)
|
|
print(f" ✅ Indexed {len(chunks)} chunks")
|
|
except Exception as e:
|
|
print(f" ❌ Error: {e}")
|
|
|
|
# Summary
|
|
print(f"\n📊 Summary:")
|
|
print(f" Sessions processed: {len(jsonl_files)}")
|
|
print(f" Skipped (empty): {skipped_empty}")
|
|
print(f" Total messages: {total_messages}")
|
|
print(f" Total chunks indexed: {total_chunks}")
|
|
|
|
stats = rag.get_stats()
|
|
print(f" Total documents in collection: {stats['total_documents']}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import argparse
|
|
|
|
parser = argparse.ArgumentParser(description="Ingest OpenClaw session transcripts into RAG")
|
|
parser.add_argument("--sessions-dir", help="Path to sessions directory (default: ~/.openclaw/agents/main/sessions)")
|
|
parser.add_argument("--chunk-size", type=int, default=20, help="Messages per chunk (default: 20)")
|
|
parser.add_argument("--chunk-overlap", type=int, default=5, help="Message overlap (default: 5)")
|
|
|
|
args = parser.parse_args()
|
|
|
|
ingest_sessions(
|
|
sessions_dir=args.sessions_dir,
|
|
chunk_size=args.chunk_size,
|
|
chunk_overlap=args.chunk_overlap
|
|
) |