150 lines
5.3 KiB
Python
150 lines
5.3 KiB
Python
import json
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class GeminiReader:
|
|
"""Reader for Gemini CLI history files."""
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
def load_data(self, history_dir: str, max_count: int = -1) -> list[dict[str, Any]]:
|
|
"""
|
|
Load data from Gemini history directory.
|
|
|
|
Args:
|
|
history_dir: Path to .gemini directory
|
|
max_count: Max number of conversations to load
|
|
|
|
Returns:
|
|
List of dictionaries with 'text' and 'metadata' keys
|
|
"""
|
|
history_path = Path(history_dir).expanduser()
|
|
if not history_path.exists():
|
|
print(f"Gemini history directory not found: {history_path}")
|
|
return []
|
|
|
|
documents = []
|
|
|
|
# 1. Load Memory (GEMINI.md)
|
|
memory_file = history_path / "GEMINI.md"
|
|
if memory_file.exists():
|
|
try:
|
|
text = memory_file.read_text(encoding="utf-8")
|
|
if text.strip():
|
|
documents.append(
|
|
{
|
|
"text": f"Gemini Memory:\n{text}",
|
|
"metadata": {"source": str(memory_file), "type": "memory"},
|
|
}
|
|
)
|
|
except Exception as e:
|
|
print(f"Error reading memory file: {e}")
|
|
|
|
# 2. Find Session Files
|
|
# Legacy JSON sessions
|
|
session_files = list(history_path.glob("session-*.json"))
|
|
# New JSONL sessions
|
|
session_files.extend(list(history_path.glob("session-*.jsonl")))
|
|
# Checkpoints
|
|
session_files.extend(list(history_path.glob("checkpoint-*.json")))
|
|
|
|
# Sort by modification time (newest first)
|
|
session_files.sort(key=lambda x: x.stat().st_mtime, reverse=True)
|
|
|
|
print(f"Found {len(session_files)} session files.")
|
|
|
|
count = 0
|
|
for file_path in session_files:
|
|
if max_count > 0 and count >= max_count:
|
|
break
|
|
|
|
try:
|
|
content = ""
|
|
if file_path.suffix == ".jsonl":
|
|
content = self._parse_jsonl_session(file_path)
|
|
elif file_path.suffix == ".json":
|
|
content = self._parse_json_session(file_path)
|
|
|
|
if content:
|
|
documents.append(
|
|
{
|
|
"text": content,
|
|
"metadata": {
|
|
"source": str(file_path),
|
|
"type": "session",
|
|
"filename": file_path.name,
|
|
},
|
|
}
|
|
)
|
|
count += 1
|
|
except Exception as e:
|
|
print(f"Error reading {file_path.name}: {e}")
|
|
|
|
print(f"Successfully loaded {len(documents)} items from Gemini history.")
|
|
return documents
|
|
|
|
def _parse_json_session(self, file_path: Path) -> str:
|
|
"""Parse legacy JSON session file."""
|
|
data = json.loads(file_path.read_text(encoding="utf-8"))
|
|
|
|
# Handle dict format (standard session)
|
|
messages = []
|
|
if isinstance(data, dict):
|
|
# Check for 'messages' key (standard format)
|
|
if "messages" in data:
|
|
for msg in data["messages"]:
|
|
role = msg.get("role", "unknown")
|
|
content = msg.get("content", "")
|
|
if content:
|
|
messages.append(f"{role.upper()}: {content}")
|
|
# Check for 'parts' key (checkpoint format sometimes)
|
|
elif "parts" in data:
|
|
messages.append(f"Saved Session Content: {data['parts']}")
|
|
|
|
# Handle list format (some older array-based sessions)
|
|
elif isinstance(data, list):
|
|
for item in data:
|
|
if isinstance(item, dict):
|
|
role = item.get("role", "unknown")
|
|
content = item.get("content", "") or item.get("parts", "")
|
|
if content:
|
|
messages.append(f"{role.upper()}: {content}")
|
|
|
|
if not messages:
|
|
return ""
|
|
|
|
return f"File: {file_path.name}\n\n" + "\n\n".join(messages)
|
|
|
|
def _parse_jsonl_session(self, file_path: Path) -> str:
|
|
"""Parse JSONL session file."""
|
|
messages = []
|
|
try:
|
|
with open(file_path, encoding="utf-8") as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
try:
|
|
data = json.loads(line)
|
|
# Skip metadata lines if they don't have content
|
|
if "role" in data and "content" in data:
|
|
messages.append(f"{data['role'].upper()}: {data['content']}")
|
|
elif "parts" in data: # sometimes parts is used
|
|
messages.append(
|
|
f"{data.get('role', 'unknown').upper()}: {data['parts']}"
|
|
)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
except Exception:
|
|
return ""
|
|
|
|
if not messages:
|
|
return ""
|
|
|
|
return f"File: {file_path.name}\n\n" + "\n\n".join(messages)
|