topos-node 0.1.0__py3-none-any.whl

topos/storage/canonical/ai_chat/mapper.py

@@ -0,0 +1,168 @@
+"""Canonical mapper - maps staging data to canonical models.
+
+Migrated from engine/canonical/mapper.py (commit 7b709af).
+"""
+
+from __future__ import annotations
+
+import logging
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List
+
+from .model import CanonicalAIChatMessage, CanonicalAIChatConversation
+
+logger = logging.getLogger("topos.storage.canonical.ai_chat.mapper")
+
+
+class CanonicalMapper(ABC):
+    """Base class for mapping staging data to canonical models."""
+
+    @abstractmethod
+    def map_to_canonical(
+        self,
+        staging_record: Dict[str, Any],
+        source: str
+    ) -> List[CanonicalAIChatMessage]:
+        """Map a staging record to canonical format.
+
+        Args:
+            staging_record: Record from staging table
+            source: Source identifier (e.g., "chatgpt")
+
+        Returns:
+            List of canonical messages (usually one, but could be multiple)
+        """
+        pass
+
+    @abstractmethod
+    def extract_conversation_id(self, staging_record: Dict[str, Any]) -> str:
+        """Extract conversation ID from staging record.
+
+        Args:
+            staging_record: Record from staging table
+
+        Returns:
+            Conversation ID (unified across sources)
+        """
+        pass
+
+
+class ChatGPTToAIChatMapper(CanonicalMapper):
+    """Maps ChatGPT staging records to canonical AI chat model."""
+
+    def map_to_canonical(
+        self,
+        staging_record: Dict[str, Any],
+        source: str = "chatgpt"
+    ) -> List[CanonicalAIChatMessage]:
+        """Map ChatGPT staging record to canonical AI chat message.
+
+        ChatGPT staging format:
+        {
+            "message_id": "...",
+            "dataset_id": "...",
+            "thread_id": "...",
+            "ts": "...",
+            "sender_type": "human" | "assistant",
+            "content": "..."
+        }
+
+        Returns:
+            List with single canonical message
+        """
+        conversation_id = self.extract_conversation_id(staging_record)
+
+        # Use source_id from staging_record if available (actual source_id like "chatgpt_file_ingestion"),
+        # otherwise fall back to source parameter (mapper ID like "chatgpt")
+        actual_source_id = staging_record.get("source_id") or source
+
+        # Preserve _metadata if present (for conversation tree reconstruction)
+        metadata = {
+            "thread_id": staging_record.get("thread_id"),
+            "original_source": source,
+        }
+        if "_metadata" in staging_record:
+            metadata.update(staging_record["_metadata"])
+
+        message = CanonicalAIChatMessage(
+            message_id=staging_record.get("message_id", ""),
+            conversation_id=conversation_id,
+            sender_type=staging_record.get("sender_type", ""),
+            sender_id=None,
+            ts=staging_record.get("ts", ""),
+            content=staging_record.get("content") or "",
+            content_rendered=None,
+            metadata_json=metadata,
+            seq=0,
+            source_id=actual_source_id,
+        )
+        return [message]
+
+    def extract_conversation_id(self, staging_record: Dict[str, Any]) -> str:
+        """Extract conversation ID from ChatGPT staging record."""
+        thread_id = staging_record.get("thread_id")
+        if thread_id:
+            return f"chatgpt:{thread_id}"
+        dataset_id = staging_record.get("dataset_id", "")
+        return f"chatgpt:{dataset_id}"
+
+
+class StoreMessageToAIChatMapper(CanonicalMapper):
+    """Maps messages from the messages table (store_message) to canonical AI chat model."""
+
+    def map_to_canonical(
+        self,
+        staging_record: Dict[str, Any],
+        source: str = "store_message"
+    ) -> List[CanonicalAIChatMessage]:
+        """Map store_message record to canonical AI chat message."""
+        conversation_id = self.extract_conversation_id(staging_record)
+        message = CanonicalAIChatMessage(
+            message_id=staging_record.get("message_id", ""),
+            conversation_id=conversation_id,
+            sender_type=staging_record.get("sender_type", ""),
+            sender_id=None,
+            ts=staging_record.get("ts", ""),
+            content=staging_record.get("content") or "",
+            content_rendered=None,
+            metadata_json={
+                "thread_id": staging_record.get("thread_id"),
+                "original_source": source,
+            },
+            seq=0,
+            source_id=source,
+        )
+        return [message]
+
+    def extract_conversation_id(self, staging_record: Dict[str, Any]) -> str:
+        """Extract conversation ID from store_message record."""
+        thread_id = staging_record.get("thread_id")
+        if thread_id:
+            return f"store_message:{thread_id}"
+        dataset_id = staging_record.get("dataset_id", "")
+        return f"store_message:{dataset_id}"
+
+
+_MAPPER_REGISTRY: Dict[str, type[CanonicalMapper]] = {
+    "chatgpt": ChatGPTToAIChatMapper,
+    "store_message": StoreMessageToAIChatMapper,
+}
+
+
+def register_mapper(source: str, mapper_class: type[CanonicalMapper]) -> None:
+    """Register a mapper class for a source."""
+    if not issubclass(mapper_class, CanonicalMapper):
+        raise TypeError(f"Mapper class must be a subclass of CanonicalMapper, got {mapper_class}")
+    _MAPPER_REGISTRY[source] = mapper_class
+    logger.info("Registered canonical mapper for source: %s -> %s", source, mapper_class.__name__)
+
+
+def get_mapper(source: str) -> CanonicalMapper:
+    """Get a mapper instance for a source."""
+    mapper_class = _MAPPER_REGISTRY.get(source)
+    if not mapper_class:
+        available = ", ".join(_MAPPER_REGISTRY.keys())
+        raise ValueError(
+            f"Unknown source for canonical mapping: {source}. Available sources: {available}"
+        )
+    return mapper_class()

topos/storage/canonical/ai_chat/model.py

@@ -0,0 +1,87 @@
+"""Canonical data models - unified schemas for compatible sources.
+
+Migrated from engine/canonical/model.py (commit 7b709af).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+
+@dataclass
+class CanonicalAIChatMessage:
+    """Canonical AI chat message model.
+
+    This is the unified format for all AI chat conversations regardless of source
+    (ChatGPT, Claude, Gemini, etc.).
+    """
+    message_id: str
+    conversation_id: str  # Unified conversation identifier
+    sender_type: str  # "human" | "assistant" | "system"
+    ts: str  # ISO timestamp
+    content: str  # Message content
+    source_id: str = ""  # Original source identifier (e.g., "chatgpt")
+    sender_id: Optional[str] = None  # Optional sender identifier
+    content_rendered: Optional[str] = None  # Optional rendered content
+    metadata_json: Optional[Dict[str, Any]] = None  # Source-specific metadata
+    seq: int = 0  # Sequence number within conversation
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for database storage."""
+        import json
+        return {
+            "message_id": self.message_id,
+            "conversation_id": self.conversation_id,
+            "sender_type": self.sender_type,
+            "sender_id": self.sender_id,
+            "ts": self.ts,
+            "content": self.content,
+            "content_rendered": self.content_rendered,
+            "metadata_json": json.dumps(self.metadata_json) if self.metadata_json else None,
+            "seq": self.seq,
+            "source_id": self.source_id,
+        }
+
+
+@dataclass
+class CanonicalAIChatConversation:
+    """Canonical AI chat conversation model.
+
+    Represents a unified conversation across all AI chat sources.
+    """
+    conversation_id: str
+    owner_user_id: str
+    title: Optional[str] = None
+    source: str = ""  # Original source (e.g., "chatgpt")
+    created_at: str = ""  # ISO timestamp
+    updated_at: str = ""  # ISO timestamp
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for database storage."""
+        return {
+            "conversation_id": self.conversation_id,
+            "owner_user_id": self.owner_user_id,
+            "title": self.title,
+            "source": self.source,
+            "created_at": self.created_at,
+            "updated_at": self.updated_at,
+        }
+
+
+class CanonicalAIChatModel:
+    """Canonical AI chat conversation model.
+
+    This model unifies all AI chat sources (ChatGPT, Claude, Gemini, etc.)
+    into a single canonical format.
+    """
+
+    @staticmethod
+    def get_conversation_table() -> str:
+        """Return the canonical conversations table name."""
+        return "ai_chat_conversations"
+
+    @staticmethod
+    def get_messages_table() -> str:
+        """Return the canonical messages table name."""
+        return "ai_chat_messages"

topos/storage/canonical/ai_chat/tables.py

@@ -0,0 +1,179 @@
+"""Canonical tables manager - manages canonical database tables.
+
+Migrated from engine/canonical/tables.py (commit 7b709af).
+"""
+
+from __future__ import annotations
+
+import logging
+import sqlite3
+from typing import Any, Dict, List
+
+from .model import CanonicalAIChatMessage, CanonicalAIChatConversation
+
+logger = logging.getLogger("topos.storage.canonical.ai_chat.tables")
+
+
+class CanonicalTablesManager:
+    """Manages canonical tables for unified data models."""
+
+    def __init__(self, conn: sqlite3.Connection):
+        """Initialize with database connection."""
+        self.conn = conn
+        self._ensure_tables()
+
+    def _ensure_tables(self) -> None:
+        """Ensure canonical tables exist. Creates them if they don't exist."""
+        try:
+            self.conn.execute("""
+                CREATE TABLE IF NOT EXISTS ai_chat_conversations (
+                    conversation_id TEXT PRIMARY KEY,
+                    owner_user_id TEXT NOT NULL,
+                    title TEXT,
+                    source_id TEXT NOT NULL,
+                    created_at TEXT NOT NULL,
+                    updated_at TEXT NOT NULL
+                )
+            """)
+            self.conn.execute("""
+                CREATE INDEX IF NOT EXISTS idx_ai_chat_conversations_owner
+                ON ai_chat_conversations(owner_user_id)
+            """)
+            self.conn.execute("""
+                CREATE INDEX IF NOT EXISTS idx_ai_chat_conversations_source_id
+                ON ai_chat_conversations(source_id)
+            """)
+            self.conn.execute("""
+                CREATE TABLE IF NOT EXISTS ai_chat_messages (
+                    message_id TEXT PRIMARY KEY,
+                    conversation_id TEXT NOT NULL,
+                    sender_type TEXT NOT NULL,
+                    sender_id TEXT,
+                    event_at TEXT NOT NULL,
+                    content TEXT NOT NULL,
+                    content_rendered TEXT,
+                    metadata_json TEXT,
+                    sequence INTEGER NOT NULL DEFAULT 0,
+                    source_id TEXT NOT NULL
+                )
+            """)
+            self.conn.execute("""
+                CREATE INDEX IF NOT EXISTS idx_ai_chat_messages_conversation
+                ON ai_chat_messages(conversation_id, sequence)
+            """)
+            self.conn.execute("""
+                CREATE INDEX IF NOT EXISTS idx_ai_chat_messages_event_at
+                ON ai_chat_messages(event_at)
+            """)
+            self.conn.execute("""
+                CREATE INDEX IF NOT EXISTS idx_ai_chat_messages_source
+                ON ai_chat_messages(source_id)
+            """)
+            self.conn.commit()
+            logger.debug("Ensured canonical tables exist")
+        except Exception as e:
+            self.conn.rollback()
+            logger.error("Failed to ensure canonical tables: %s", e)
+
+    def write_conversations_batch(
+        self,
+        conversations: List[CanonicalAIChatConversation],
+        batch_size: int = 1000,
+    ) -> int:
+        """Write multiple conversations to canonical table in batches."""
+        if not conversations:
+            return 0
+        written = 0
+        try:
+            for i in range(0, len(conversations), batch_size):
+                batch = conversations[i:i + batch_size]
+                values = [
+                    (
+                        conv.conversation_id,
+                        conv.owner_user_id,
+                        conv.title,
+                        conv.source,
+                        conv.created_at,
+                        conv.updated_at,
+                    )
+                    for conv in batch
+                ]
+                self.conn.executemany("""
+                    INSERT OR REPLACE INTO ai_chat_conversations (
+                        conversation_id, owner_user_id, title, source_id, created_at, updated_at
+                    ) VALUES (?, ?, ?, ?, ?, ?)
+                """, values)
+                self.conn.commit()
+                written += len(batch)
+                logger.debug("Wrote batch of %d conversations (total: %d)", len(batch), written)
+        except Exception as e:
+            self.conn.rollback()
+            logger.error("Failed to write conversations batch: %s", e)
+            raise
+        return written
+
+    def write_messages_batch(
+        self,
+        messages: List[CanonicalAIChatMessage],
+        batch_size: int = 1000,
+    ) -> int:
+        """Write multiple messages to canonical table in batches."""
+        if not messages:
+            return 0
+        import json
+        written = 0
+        try:
+            for i in range(0, len(messages), batch_size):
+                batch = messages[i:i + batch_size]
+                values = [
+                    (
+                        msg.message_id,
+                        msg.conversation_id,
+                        msg.sender_type,
+                        msg.sender_id,
+                        msg.ts,
+                        msg.content,
+                        msg.content_rendered,
+                        json.dumps(msg.metadata_json) if msg.metadata_json else None,
+                        msg.seq,
+                        msg.source_id,
+                    )
+                    for msg in batch
+                ]
+                self.conn.executemany("""
+                    INSERT OR REPLACE INTO ai_chat_messages (
+                        message_id, conversation_id, sender_type, sender_id, event_at,
+                        content, content_rendered, metadata_json, sequence, source_id
+                    ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+                """, values)
+                self.conn.commit()
+                written += len(batch)
+                logger.debug("Wrote batch of %d messages (total: %d)", len(batch), written)
+        except Exception as e:
+            self.conn.rollback()
+            logger.error("Failed to write messages batch: %s", e)
+            raise
+        return written
+
+    def update_message_sequences(self, conversation_id: str) -> None:
+        """Update sequence numbers for messages in a conversation."""
+        try:
+            cursor = self.conn.execute("""
+                SELECT message_id, event_at
+                FROM ai_chat_messages
+                WHERE conversation_id = ?
+                ORDER BY event_at ASC
+            """, (conversation_id,))
+            messages = cursor.fetchall()
+            for seq, (message_id, _) in enumerate(messages):
+                self.conn.execute("""
+                    UPDATE ai_chat_messages
+                    SET sequence = ?
+                    WHERE message_id = ?
+                """, (seq, message_id))
+            self.conn.commit()
+            logger.debug("Updated sequences for conversation %s (%d messages)", conversation_id, len(messages))
+        except Exception as e:
+            self.conn.rollback()
+            logger.error("Failed to update message sequences: %s", e)
+            raise

topos/storage/canonical/canonical_store.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict
+
+
+@dataclass(frozen=True)
+class CanonicalRef:
+    record_id: str
+
+
+class CanonicalStore:
+    def upsert(self, record: Dict[str, str]) -> CanonicalRef:
+        raise NotImplementedError
+
+
+class InMemoryCanonicalStore(CanonicalStore):
+    def __init__(self):
+        self._records: Dict[str, Dict[str, str]] = {}
+
+    def upsert(self, record: Dict[str, str]) -> CanonicalRef:
+        record_id = record.get("record_id") or record.get("message_id") or ""
+        self._records[record_id] = record
+        return CanonicalRef(record_id=record_id)