oghma 0.0.1__py3-none-any.whl → 0.3.0__py3-none-any.whl

oghma/mcp_server.py

@@ -0,0 +1,112 @@
+from collections import Counter
+from contextlib import asynccontextmanager
+from typing import Any
+
+from mcp.server.fastmcp import FastMCP
+
+from oghma.config import load_config
+from oghma.embedder import EmbedConfig, create_embedder
+from oghma.storage import MemoryRecord, Storage
+
+
+@asynccontextmanager
+async def lifespan(_: FastMCP):
+    config = load_config()
+    storage = Storage(config=config, read_only=True)
+    yield {"storage": storage, "config": config}
+
+
+mcp = FastMCP("Oghma Memory", lifespan=lifespan)
+
+
+def _get_storage() -> Storage:
+    return mcp.get_context()["storage"]
+
+
+def _get_config() -> dict[str, Any]:
+    return mcp.get_context().get("config", {})
+
+
+@mcp.tool()
+def oghma_search(
+    query: str,
+    category: str | None = None,
+    source_tool: str | None = None,
+    limit: int = 10,
+    search_mode: str = "keyword",
+) -> list[MemoryRecord]:
+    """Search memories by keyword, vector, or hybrid mode."""
+    if limit < 1:
+        raise ValueError("limit must be >= 1")
+    if search_mode not in {"keyword", "vector", "hybrid"}:
+        raise ValueError("search_mode must be one of: keyword, vector, hybrid")
+
+    storage = _get_storage()
+    if search_mode == "keyword":
+        return storage.search_memories(
+            query=query,
+            category=category,
+            source_tool=source_tool,
+            limit=limit,
+        )
+
+    query_embedding: list[float] | None = None
+    try:
+        embed_config = _get_config().get("embedding", {})
+        embedder = create_embedder(EmbedConfig.from_dict(embed_config))
+        query_embedding = embedder.embed(query)
+    except Exception:
+        if search_mode == "vector":
+            return []
+
+    return storage.search_memories_hybrid(
+        query=query,
+        query_embedding=query_embedding,
+        category=category,
+        source_tool=source_tool,
+        limit=limit,
+        search_mode=search_mode,
+    )
+
+
+@mcp.tool()
+def oghma_get(memory_id: int) -> MemoryRecord | None:
+    """Get a memory by ID."""
+    storage = _get_storage()
+    return storage.get_memory_by_id(memory_id)
+
+
+@mcp.tool()
+def oghma_stats() -> dict[str, Any]:
+    """Get memory database statistics."""
+    storage = _get_storage()
+    memories = storage.get_all_memories(status="active")
+    extraction_logs = storage.get_recent_extraction_logs(limit=1)
+
+    return {
+        "total_memories": storage.get_memory_count(),
+        "memories_by_category": dict(Counter(memory["category"] for memory in memories)),
+        "memories_by_source": dict(Counter(memory["source_tool"] for memory in memories)),
+        "last_extraction_time": extraction_logs[0]["created_at"] if extraction_logs else None,
+    }
+
+
+@mcp.tool()
+def oghma_categories() -> list[dict[str, Any]]:
+    """List categories with memory counts."""
+    storage = _get_storage()
+    memories = storage.get_all_memories(status="active")
+    category_counts = Counter(memory["category"] for memory in memories)
+
+    return [
+        {"category": category, "count": count}
+        for category, count in sorted(category_counts.items(), key=lambda item: (-item[1], item[0]))
+    ]
+
+
+def main() -> None:
+    mcp.run()
+
+
+if __name__ == "__main__":
+    main()

oghma/migration.py

@@ -0,0 +1,63 @@
+from dataclasses import dataclass
+
+from openai import APIError
+
+from oghma.embedder import Embedder
+from oghma.storage import Storage
+
+
+@dataclass
+class MigrationResult:
+    processed: int
+    migrated: int
+    skipped: int
+    failed: int
+
+
+class EmbeddingMigration:
+    def __init__(
+        self,
+        storage: Storage,
+        embedder: Embedder,
+        batch_size: int = 100,
+    ):
+        self.storage = storage
+        self.embedder = embedder
+        self.batch_size = batch_size
+
+    def run(self, dry_run: bool = False) -> MigrationResult:
+        processed = 0
+        migrated = 0
+        failed = 0
+
+        while True:
+            batch = self.storage.get_memories_without_embeddings(limit=self.batch_size)
+            if not batch:
+                break
+
+            contents = [memory["content"] for memory in batch]
+            processed += len(batch)
+
+            if dry_run:
+                continue
+
+            try:
+                vectors = self.embedder.embed_batch(contents)
+            except (APIError, RuntimeError, ValueError):
+                failed += len(batch)
+                continue
+
+            for memory, vector in zip(batch, vectors, strict=False):
+                success = self.storage.upsert_memory_embedding(memory["id"], vector)
+                if success:
+                    migrated += 1
+                else:
+                    failed += 1
+
+        skipped = max(processed - migrated - failed, 0)
+        return MigrationResult(
+            processed=processed,
+            migrated=migrated,
+            skipped=skipped,
+            failed=failed,
+        )

oghma/parsers/__init__.py

@@ -0,0 +1,26 @@
+from pathlib import Path
+
+from oghma.parsers.base import BaseParser
+from oghma.parsers.base import Message as _Message
+from oghma.parsers.claude_code import ClaudeCodeParser
+from oghma.parsers.codex import CodexParser
+from oghma.parsers.openclaw import OpenClawParser
+from oghma.parsers.opencode import OpenCodeParser
+
+__all__ = ["BaseParser", "Message", "get_parser_for_file", "PARSERS"]
+
+Message = _Message
+
+PARSERS: list[BaseParser] = [
+    ClaudeCodeParser(),
+    CodexParser(),
+    OpenClawParser(),
+    OpenCodeParser(),
+]
+
+
+def get_parser_for_file(file_path: Path) -> BaseParser | None:
+    for parser in PARSERS:
+        if parser.can_parse(file_path):
+            return parser
+    return None

oghma/parsers/base.py

@@ -0,0 +1,24 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    pass
+
+
+@dataclass
+class Message:
+    role: str
+    content: str
+    timestamp: str | None = None
+
+
+class BaseParser(ABC):
+    @abstractmethod
+    def parse(self, file_path: Path) -> list[Message]:
+        pass
+
+    @abstractmethod
+    def can_parse(self, file_path: Path) -> bool:
+        pass

oghma/parsers/claude_code.py

@@ -0,0 +1,62 @@
+import json
+from pathlib import Path
+
+from oghma.parsers.base import BaseParser, Message
+
+
+class ClaudeCodeParser(BaseParser):
+    def can_parse(self, file_path: Path) -> bool:
+        if not file_path.name.endswith(".jsonl"):
+            return False
+        path_str = str(file_path)
+        return ".claude/projects/-Users-" in path_str
+
+    def parse(self, file_path: Path) -> list[Message]:
+        messages: list[Message] = []
+
+        try:
+            with open(file_path, encoding="utf-8") as f:
+                for _line_num, line in enumerate(f, 1):
+                    line = line.strip()
+                    if not line:
+                        continue
+
+                    try:
+                        data = json.loads(line)
+                        role = self._extract_role(data)
+                        content = self._extract_content(data)
+
+                        if role and content:
+                            messages.append(Message(role=role, content=content[:3000]))
+                    except (json.JSONDecodeError, KeyError, TypeError):
+                        continue
+        except (OSError, UnicodeDecodeError):
+            return []
+
+        return messages
+
+    def _extract_role(self, data: dict) -> str | None:
+        msg_type = data.get("type")
+        if msg_type == "user":
+            return "user"
+        elif msg_type == "assistant":
+            return "assistant"
+        return None
+
+    def _extract_content(self, data: dict) -> str:
+        message = data.get("message", {})
+        content = message.get("content", "")
+
+        if isinstance(content, str):
+            return content
+
+        if isinstance(content, list):
+            parts: list[str] = []
+            for block in content:
+                if isinstance(block, dict) and block.get("type") == "text":
+                    text = block.get("text", "")
+                    if text:
+                        parts.append(text)
+            return "\n".join(parts)
+
+        return str(content)

oghma/parsers/codex.py

@@ -0,0 +1,84 @@
+import json
+from pathlib import Path
+
+from oghma.parsers.base import BaseParser, Message
+
+
+class CodexParser(BaseParser):
+    def can_parse(self, file_path: Path) -> bool:
+        if not file_path.name.endswith(".jsonl"):
+            return False
+        path_str = str(file_path)
+        return ".codex/sessions/" in path_str and "rollout-" in file_path.name
+
+    def parse(self, file_path: Path) -> list[Message]:
+        messages: list[Message] = []
+
+        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)
+                        role = self._extract_role(data)
+                        content = self._extract_content(data)
+
+                        if role and content:
+                            messages.append(Message(role=role, content=content[:3000]))
+                    except (json.JSONDecodeError, KeyError, TypeError):
+                        continue
+        except (OSError, UnicodeDecodeError):
+            return []
+
+        return messages
+
+    def _extract_role(self, data: dict) -> str | None:
+        msg_type = data.get("type")
+        # Support both old format (item) and new format (response_item, event_msg)
+        if msg_type not in ("item", "response_item", "event_msg"):
+            return None
+
+        payload = data.get("payload", {})
+
+        # New format: role directly in payload
+        if "role" in payload:
+            role = payload.get("role")
+            # Map developer/assistant to assistant, user to user
+            if role in ("developer", "assistant"):
+                return "assistant"
+            elif role == "user":
+                return "user"
+            return None
+
+        # Old format: nested in payload.item
+        item = payload.get("item", {})
+        return item.get("role")
+
+    def _extract_content(self, data: dict) -> str:
+        payload = data.get("payload", {})
+
+        # New format: content directly in payload
+        if "content" in payload:
+            content = payload.get("content", "")
+        else:
+            # Old format: nested in payload.item
+            item = payload.get("item", {})
+            content = item.get("content", "")
+
+        if isinstance(content, str):
+            return content
+
+        if isinstance(content, list):
+            parts: list[str] = []
+            for block in content:
+                if isinstance(block, dict):
+                    block_type = block.get("type")
+                    text = block.get("text", "")
+                    if text and block_type in ("input_text", "output_text", "text"):
+                        parts.append(text)
+            return "\n".join(parts)
+
+        return str(content)

oghma/parsers/openclaw.py

@@ -0,0 +1,64 @@
+import json
+from pathlib import Path
+
+from oghma.parsers.base import BaseParser, Message
+
+
+class OpenClawParser(BaseParser):
+    def can_parse(self, file_path: Path) -> bool:
+        if not file_path.name.endswith(".jsonl"):
+            return False
+        path_str = str(file_path)
+        return ".openclaw/agents/" in path_str and "/sessions/" in path_str
+
+    def parse(self, file_path: Path) -> list[Message]:
+        messages: list[Message] = []
+
+        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)
+                        role = self._extract_role(data)
+                        content = self._extract_content(data)
+
+                        if role and content:
+                            messages.append(Message(role=role, content=content[:3000]))
+                    except (json.JSONDecodeError, KeyError, TypeError):
+                        continue
+        except (OSError, UnicodeDecodeError):
+            return []
+
+        return messages
+
+    def _extract_role(self, data: dict) -> str | None:
+        if data.get("type") != "message":
+            return None
+
+        message = data.get("message", {})
+        role = message.get("role")
+        if role in ("user", "assistant"):
+            return role
+        return None
+
+    def _extract_content(self, data: dict) -> str:
+        message = data.get("message", {})
+        content = message.get("content", "")
+
+        if isinstance(content, str):
+            return content
+
+        if isinstance(content, list):
+            parts: list[str] = []
+            for item in content:
+                if isinstance(item, dict) and "text" in item:
+                    text = item.get("text", "")
+                    if text:
+                        parts.append(text)
+            return "\n".join(parts)
+
+        return str(content)

oghma/parsers/opencode.py

@@ -0,0 +1,90 @@
+import json
+from pathlib import Path
+
+from oghma.parsers.base import BaseParser, Message
+
+
+class OpenCodeParser(BaseParser):
+    def can_parse(self, file_path: Path) -> bool:
+        if not file_path.is_dir():
+            return False
+        path_str = str(file_path)
+        return ".local/share/opencode/storage/message/ses_" in path_str
+
+    def parse(self, file_path: Path) -> list[Message]:
+        messages: list[Message] = []
+
+        message_files = sorted(file_path.glob("msg_*.json"))
+        part_files = list(file_path.glob("part/msg_*/prt_*.json"))
+
+        parts_map = self._build_parts_map(part_files)
+
+        for msg_file in message_files:
+            try:
+                with open(msg_file, encoding="utf-8") as f:
+                    data = json.load(f)
+            except (OSError, json.JSONDecodeError):
+                continue
+
+            role = self._extract_role(data)
+            if not role:
+                continue
+
+            content_parts = self._get_message_content(data, parts_map)
+            if content_parts:
+                content = "\n".join(content_parts)
+                messages.append(Message(role=role, content=content[:3000]))
+
+        return messages
+
+    def _build_parts_map(self, part_files: list[Path]) -> dict[str, list[str]]:
+        parts_map: dict[str, list[str]] = {}
+
+        for part_file in part_files:
+            try:
+                with open(part_file, encoding="utf-8") as f:
+                    data = json.load(f)
+            except (OSError, json.JSONDecodeError):
+                continue
+
+            msg_id = data.get("message_id")
+            if not msg_id:
+                continue
+
+            if msg_id not in parts_map:
+                parts_map[msg_id] = []
+
+            text = data.get("text", "")
+            if text:
+                parts_map[msg_id].append(text)
+
+        return parts_map
+
+    def _extract_role(self, data: dict) -> str | None:
+        role = data.get("role")
+        if role in ("user", "assistant"):
+            return role
+        return None
+
+    def _get_message_content(self, data: dict, parts_map: dict[str, list[str]]) -> list[str]:
+        msg_id = data.get("id")
+        if msg_id in parts_map:
+            return parts_map[msg_id]
+
+        content = data.get("content", "")
+        if isinstance(content, str):
+            return [content]
+
+        if isinstance(content, list):
+            parts: list[str] = []
+            for item in content:
+                if isinstance(item, dict) and "text" in item:
+                    text = item.get("text", "")
+                    if text:
+                        parts.append(text)
+            return parts
+
+        if content:
+            return [str(content)]
+
+        return []