typeagent-py 0.1.0__py3-none-any.whl

typeagent/storage/sqlite/provider.py

@@ -0,0 +1,293 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""SQLite storage provider implementation."""
+
+import json
+import sqlite3
+
+from ...knowpro import interfaces
+from ...knowpro.convsettings import MessageTextIndexSettings, RelatedTermIndexSettings
+from .collections import SqliteMessageCollection, SqliteSemanticRefCollection
+from .messageindex import SqliteMessageTextIndex
+from .propindex import SqlitePropertyIndex
+from .reltermsindex import SqliteRelatedTermsIndex
+from .semrefindex import SqliteTermToSemanticRefIndex
+from .timestampindex import SqliteTimestampToTextRangeIndex
+from .schema import (
+    CONVERSATIONS_SCHEMA,
+    MESSAGE_TEXT_INDEX_SCHEMA,
+    MESSAGES_SCHEMA,
+    PROPERTY_INDEX_SCHEMA,
+    RELATED_TERMS_ALIASES_SCHEMA,
+    RELATED_TERMS_FUZZY_SCHEMA,
+    SEMANTIC_REF_INDEX_SCHEMA,
+    SEMANTIC_REFS_SCHEMA,
+    ConversationMetadata,
+    get_db_schema_version,
+    init_db_schema,
+)
+
+
+class SqliteStorageProvider[TMessage: interfaces.IMessage](
+    interfaces.IStorageProvider[TMessage]
+):
+    """SQLite-backed storage provider implementation."""
+
+    def __init__(
+        self,
+        db_path: str = ":memory:",
+        conversation_id: str = "default",
+        message_type: type[TMessage] = None,  # type: ignore
+        semantic_ref_type: type[interfaces.SemanticRef] = None,  # type: ignore
+        conversation_index_settings=None,
+        message_text_index_settings: MessageTextIndexSettings | None = None,
+        related_term_index_settings: RelatedTermIndexSettings | None = None,
+    ):
+        self.db_path = db_path
+        self.conversation_id = conversation_id
+        self.message_type = message_type
+        self.semantic_ref_type = semantic_ref_type
+
+        # Settings with defaults (require embedding settings)
+        self.conversation_index_settings = conversation_index_settings or {}
+        if message_text_index_settings is None:
+            # Create default embedding settings if not provided
+            from ...aitools.embeddings import AsyncEmbeddingModel
+            from ...aitools.vectorbase import TextEmbeddingIndexSettings
+
+            model = AsyncEmbeddingModel()
+            embedding_settings = TextEmbeddingIndexSettings(model)
+            message_text_index_settings = MessageTextIndexSettings(embedding_settings)
+        self.message_text_index_settings = message_text_index_settings
+
+        if related_term_index_settings is None:
+            # Use the same embedding settings
+            embedding_settings = message_text_index_settings.embedding_index_settings
+            related_term_index_settings = RelatedTermIndexSettings(embedding_settings)
+        self.related_term_index_settings = related_term_index_settings
+
+        # Initialize database connection
+        self.db = sqlite3.connect(db_path)
+
+        # Configure SQLite for optimal bulk insertion performance
+        self.db.execute("PRAGMA foreign_keys = ON")
+        # Improve write performance for bulk operations
+        self.db.execute("PRAGMA synchronous = NORMAL")  # Faster than FULL, still safe
+        self.db.execute(
+            "PRAGMA journal_mode = WAL"
+        )  # Write-Ahead Logging for better concurrency
+        self.db.execute("PRAGMA cache_size = -64000")  # 64MB cache (negative = KB)
+        self.db.execute("PRAGMA temp_store = MEMORY")  # Store temp tables in memory
+        self.db.execute("PRAGMA mmap_size = 268435456")  # 256MB memory-mapped I/O
+
+        # Initialize schema
+        init_db_schema(self.db)
+
+        # Initialize collections
+        # Initialize message collection first
+        self._message_collection = SqliteMessageCollection(self.db, self.message_type)
+        self._semantic_ref_collection = SqliteSemanticRefCollection(self.db)
+
+        # Initialize indexes
+        self._term_to_semantic_ref_index = SqliteTermToSemanticRefIndex(self.db)
+        self._property_index = SqlitePropertyIndex(self.db)
+        self._timestamp_index = SqliteTimestampToTextRangeIndex(self.db)
+        self._message_text_index = SqliteMessageTextIndex(
+            self.db,
+            self.message_text_index_settings,
+            self._message_collection,
+        )
+        # Initialize related terms index
+        self._related_terms_index = SqliteRelatedTermsIndex(
+            self.db, self.related_term_index_settings.embedding_index_settings
+        )
+
+        # Connect message collection to message text index for automatic indexing
+        self._message_collection.set_message_text_index(self._message_text_index)
+
+    async def close(self) -> None:
+        """Close the database connection. COMMITS."""
+        if hasattr(self, "db"):
+            self.db.commit()
+            self.db.close()
+            del self.db
+
+    def __del__(self) -> None:
+        """Ensure database is closed when object is deleted. ROLLS BACK."""
+        # Can't use async in __del__, so close directly
+        if hasattr(self, "db"):
+            self.db.rollback()
+            self.db.close()
+            del self.db
+
+    @property
+    def messages(self) -> SqliteMessageCollection[TMessage]:
+        return self._message_collection
+
+    @property
+    def semantic_refs(self) -> SqliteSemanticRefCollection:
+        return self._semantic_ref_collection
+
+    @property
+    def term_to_semantic_ref_index(self) -> SqliteTermToSemanticRefIndex:
+        return self._term_to_semantic_ref_index
+
+    @property
+    def property_index(self) -> SqlitePropertyIndex:
+        return self._property_index
+
+    @property
+    def timestamp_index(self) -> SqliteTimestampToTextRangeIndex:
+        return self._timestamp_index
+
+    @property
+    def message_text_index(self) -> SqliteMessageTextIndex:
+        return self._message_text_index
+
+    @property
+    def related_terms_index(self) -> SqliteRelatedTermsIndex:
+        return self._related_terms_index
+
+    # Async getters required by base class
+    async def get_message_collection(
+        self, message_type: type[TMessage] | None = None
+    ) -> interfaces.IMessageCollection[TMessage]:
+        """Get the message collection."""
+        return self._message_collection
+
+    async def get_semantic_ref_collection(self) -> interfaces.ISemanticRefCollection:
+        """Get the semantic reference collection."""
+        return self._semantic_ref_collection
+
+    async def get_semantic_ref_index(self) -> interfaces.ITermToSemanticRefIndex:
+        """Get the semantic reference index."""
+        return self._term_to_semantic_ref_index
+
+    async def get_property_index(self) -> interfaces.IPropertyToSemanticRefIndex:
+        """Get the property index."""
+        return self._property_index
+
+    async def get_timestamp_index(self) -> interfaces.ITimestampToTextRangeIndex:
+        """Get the timestamp index."""
+        return self._timestamp_index
+
+    async def get_message_text_index(self) -> interfaces.IMessageTextIndex[TMessage]:
+        """Get the message text index."""
+        return self._message_text_index
+
+    async def get_related_terms_index(self) -> interfaces.ITermToRelatedTermsIndex:
+        """Get the related terms index."""
+        return self._related_terms_index
+
+    async def get_conversation_threads(self) -> interfaces.IConversationThreads:
+        """Get the conversation threads."""
+        # For now, return a simple implementation
+        # In a full implementation, this would be stored/retrieved from SQLite
+        from ...storage.memory.convthreads import ConversationThreads
+
+        return ConversationThreads(
+            self.message_text_index_settings.embedding_index_settings
+        )
+
+    async def clear(self) -> None:
+        """Clear all data from the storage provider."""
+        cursor = self.db.cursor()
+        # Clear in reverse dependency order
+        cursor.execute("DELETE FROM RelatedTermsFuzzy")
+        cursor.execute("DELETE FROM RelatedTermsAliases")
+        cursor.execute("DELETE FROM MessageTextIndex")
+        cursor.execute("DELETE FROM PropertyIndex")
+        cursor.execute("DELETE FROM SemanticRefIndex")
+        cursor.execute("DELETE FROM SemanticRefs")
+        cursor.execute("DELETE FROM Messages")
+        cursor.execute("DELETE FROM ConversationMetadata")
+
+        # Clear in-memory indexes
+        await self._message_text_index.clear()
+
+    def serialize(self) -> dict:
+        """Serialize all storage provider data."""
+        return {
+            "termToSemanticRefIndexData": self._term_to_semantic_ref_index.serialize(),
+            "relatedTermsIndexData": self._related_terms_index.serialize(),
+        }
+
+    async def deserialize(self, data: dict) -> None:
+        """Deserialize storage provider data."""
+        # Deserialize term to semantic ref index
+        if data.get("termToSemanticRefIndexData"):
+            await self._term_to_semantic_ref_index.deserialize(
+                data["termToSemanticRefIndexData"]
+            )
+
+        # Deserialize related terms index
+        if data.get("relatedTermsIndexData"):
+            await self._related_terms_index.deserialize(data["relatedTermsIndexData"])
+
+        # Deserialize message text index
+        if data.get("messageIndexData"):
+            await self._message_text_index.deserialize(data["messageIndexData"])
+
+    def get_conversation_metadata(self) -> ConversationMetadata | None:
+        """Get conversation metadata."""
+        cursor = self.db.cursor()
+        cursor.execute(
+            "SELECT name_tag, schema_version, created_at, updated_at, tags, extra FROM ConversationMetadata LIMIT 1"
+        )
+        row = cursor.fetchone()
+        if row:
+            return ConversationMetadata(
+                name_tag=row[0],
+                schema_version=row[1],
+                created_at=row[2],
+                updated_at=row[3],
+                tags=json.loads(row[4]) if row[4] else [],
+                extra=json.loads(row[5]) if row[5] else {},
+            )
+        return None
+
+    def update_conversation_metadata(
+        self, created_at: str | None = None, updated_at: str | None = None
+    ) -> None:
+        """Update conversation metadata."""
+        cursor = self.db.cursor()
+
+        # Check if conversation metadata exists
+        cursor.execute("SELECT 1 FROM ConversationMetadata LIMIT 1")
+
+        if cursor.fetchone():
+            # Update existing
+            updates = []
+            params = []
+            if created_at is not None:
+                updates.append("created_at = ?")
+                params.append(created_at)
+            if updated_at is not None:
+                updates.append("updated_at = ?")
+                params.append(updated_at)
+
+            if updates:
+                cursor.execute(
+                    f"UPDATE ConversationMetadata SET {', '.join(updates)}",
+                    params,
+                )
+        else:
+            # Insert new with default values
+            name_tag = f"conversation_{self.conversation_id}"
+            schema_version = "1.0"
+            tags = json.dumps([])
+            extra = json.dumps({})
+
+            cursor.execute(
+                "INSERT INTO ConversationMetadata (name_tag, schema_version, created_at, updated_at, tags, extra) VALUES (?, ?, ?, ?, ?, ?)",
+                (name_tag, schema_version, created_at, updated_at, tags, extra),
+            )
+
+    def get_db_version(self) -> int:
+        """Get the database schema version."""
+        version_str = get_db_schema_version(self.db)
+        try:
+            return int(version_str.split(".")[0])  # Get major version as int
+        except (ValueError, AttributeError):
+            return 1  # Default version

typeagent/storage/sqlite/reltermsindex.py

@@ -0,0 +1,328 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""SQLite-based related terms index implementations."""
+
+import sqlite3
+
+from ...aitools.embeddings import AsyncEmbeddingModel, NormalizedEmbeddings
+from ...aitools.vectorbase import TextEmbeddingIndexSettings, VectorBase
+from ...knowpro import interfaces
+
+from .schema import serialize_embedding, deserialize_embedding
+
+
+class SqliteRelatedTermsAliases(interfaces.ITermToRelatedTerms):
+    """SQLite-backed implementation of term to related terms aliases."""
+
+    def __init__(self, db: sqlite3.Connection):
+        self.db = db
+
+    async def lookup_term(self, text: str) -> list[interfaces.Term] | None:
+        cursor = self.db.cursor()
+        cursor.execute("SELECT alias FROM RelatedTermsAliases WHERE term = ?", (text,))
+        results = [interfaces.Term(row[0]) for row in cursor.fetchall()]
+        return results if results else None
+
+    async def add_related_term(
+        self, text: str, related_terms: interfaces.Term | list[interfaces.Term]
+    ) -> None:
+        # Convert single Term to list
+        if isinstance(related_terms, interfaces.Term):
+            related_terms = [related_terms]
+
+        cursor = self.db.cursor()
+        # Add new aliases (use INSERT OR IGNORE to avoid duplicates)
+        for related_term in related_terms:
+            cursor.execute(
+                "INSERT OR IGNORE INTO RelatedTermsAliases (term, alias) VALUES (?, ?)",
+                (text, related_term.text),
+            )
+
+    async def remove_term(self, text: str) -> None:
+        cursor = self.db.cursor()
+        cursor.execute("DELETE FROM RelatedTermsAliases WHERE term = ?", (text,))
+
+    async def clear(self) -> None:
+        cursor = self.db.cursor()
+        cursor.execute("DELETE FROM RelatedTermsAliases")
+
+    async def set_related_terms(self, term: str, related_terms: list[str]) -> None:
+        cursor = self.db.cursor()
+        # Clear existing aliases for this term
+        cursor.execute("DELETE FROM RelatedTermsAliases WHERE term = ?", (term,))
+        # Add new aliases
+        for alias in related_terms:
+            cursor.execute(
+                "INSERT INTO RelatedTermsAliases (term, alias) VALUES (?, ?)",
+                (term, alias),
+            )
+
+    async def size(self) -> int:
+        cursor = self.db.cursor()
+        cursor.execute("SELECT COUNT(DISTINCT term) FROM RelatedTermsAliases")
+        return cursor.fetchone()[0]
+
+    async def get_terms(self) -> list[str]:
+        cursor = self.db.cursor()
+        cursor.execute("SELECT DISTINCT term FROM RelatedTermsAliases ORDER BY term")
+        return [row[0] for row in cursor.fetchall()]
+
+    async def is_empty(self) -> bool:
+        cursor = self.db.cursor()
+        cursor.execute("SELECT COUNT(*) FROM RelatedTermsAliases")
+        return cursor.fetchone()[0] == 0
+
+    async def serialize(self) -> interfaces.TermToRelatedTermsData:
+        """Serialize the aliases data."""
+        cursor = self.db.cursor()
+        cursor.execute(
+            "SELECT term, alias FROM RelatedTermsAliases ORDER BY term, alias"
+        )
+
+        # Group by term
+        term_to_aliases: dict[str, list[str]] = {}
+        for term, alias in cursor.fetchall():
+            if term not in term_to_aliases:
+                term_to_aliases[term] = []
+            term_to_aliases[term].append(alias)
+
+        # Convert to the expected format
+        items = []
+        for term, aliases in term_to_aliases.items():
+            term_data_list = [interfaces.TermData(text=alias) for alias in aliases]
+            items.append(
+                interfaces.TermsToRelatedTermsDataItem(
+                    termText=term, relatedTerms=term_data_list
+                )
+            )
+
+        return interfaces.TermToRelatedTermsData(relatedTerms=items)
+
+    async def deserialize(self, data: interfaces.TermToRelatedTermsData | None) -> None:
+        """Deserialize alias data."""
+        cursor = self.db.cursor()
+
+        # Clear existing data
+        cursor.execute("DELETE FROM RelatedTermsAliases")
+
+        if data is None:
+            return
+
+        related_terms = data.get("relatedTerms", [])
+
+        if related_terms:
+            # Prepare all insertion data for bulk operation
+            insertion_data = []
+            for item in related_terms:
+                if item and item.get("termText") and item.get("relatedTerms"):
+                    term = item["termText"]
+                    for term_data in item["relatedTerms"]:
+                        alias = term_data["text"]
+                        insertion_data.append((term, alias))
+
+            # Bulk insert all the data
+            if insertion_data:
+                cursor.executemany(
+                    "INSERT INTO RelatedTermsAliases (term, alias) VALUES (?, ?)",
+                    insertion_data,
+                )
+
+
+class SqliteRelatedTermsFuzzy(interfaces.ITermToRelatedTermsFuzzy):
+    """SQLite-backed implementation of fuzzy term relationships with persistent embeddings."""
+
+    def __init__(self, db: sqlite3.Connection, settings: TextEmbeddingIndexSettings):
+        self.db = db
+        self._embedding_settings = settings
+        self._vector_base = VectorBase(self._embedding_settings)
+        # Maintain our own list of terms to map ordinals back to keys
+        self._terms_list: list[str] = []  # TODO: Use the database instead?
+        self._added_terms: set[str] = set()  # TODO: Ditto?
+        # If items exist in the db, copy them into the VectorBase, terms list, and added terms
+        if self._size() > 0:
+            cursor = self.db.cursor()
+            cursor.execute(
+                "SELECT term, term_embedding FROM RelatedTermsFuzzy ORDER BY term"
+            )
+            rows = cursor.fetchall()
+            for term, blob in rows:
+                assert blob is not None, term
+                embedding: NormalizedEmbeddings = deserialize_embedding(blob)
+                # Add to VectorBase at the correct ordinal
+                self._vector_base.add_embedding(term, embedding)
+                self._terms_list.append(term)
+                self._added_terms.add(term)
+
+    async def lookup_term(
+        self,
+        text: str,
+        max_hits: int | None = None,
+        min_score: float | None = None,
+    ) -> list[interfaces.Term]:
+        """Look up similar terms using fuzzy matching."""
+
+        # Search for similar terms using VectorBase
+        similar_results = await self._vector_base.fuzzy_lookup(
+            text, max_hits=max_hits, min_score=min_score
+        )
+
+        # Convert VectorBase results to Term objects
+        results = []
+        for scored_int in similar_results:
+            # Get the term text from the list of terms  # TODO: Use the database instead?
+            if scored_int.item < len(self._terms_list):
+                term_text = self._terms_list[scored_int.item]
+                results.append(interfaces.Term(term_text, scored_int.score))
+
+        return results
+
+    async def remove_term(self, term: str) -> None:
+        raise NotImplementedError(
+            "TODO: Removal from VectorBase, _terms_list, _terms_to_ordinal"
+        )
+        # cursor = self.db.cursor()
+        # cursor.execute("DELETE FROM RelatedTermsFuzzy WHERE term = ?", (term,))
+
+        # Clear VectorBase and local mappings - they will be rebuilt on next lookup
+        # NO THEY WON'T
+        # self._vector_base.clear()
+        # self._terms_list.clear()
+        # self._added_terms.clear()
+
+    async def clear(self) -> None:
+        cursor = self.db.cursor()
+        cursor.execute("DELETE FROM RelatedTermsFuzzy")
+
+    async def size(self) -> int:
+        return self._size()
+
+    def _size(self) -> int:
+        cursor = self.db.cursor()
+        cursor.execute("SELECT COUNT(term) FROM RelatedTermsFuzzy")
+        return cursor.fetchone()[0]
+
+    async def get_terms(self) -> list[str]:
+        cursor = self.db.cursor()
+        cursor.execute("SELECT term FROM RelatedTermsFuzzy ORDER BY term")
+        return [row[0] for row in cursor.fetchall()]
+
+    async def add_terms(self, texts: list[str]) -> None:
+        """Add terms."""
+        cursor = self.db.cursor()
+        # TODO: Batch additions to database
+        for text in texts:
+            if text in self._added_terms:
+                continue
+
+            # Add to VectorBase for fuzzy lookup
+            await self._vector_base.add_key(text)
+            self._terms_list.append(text)
+            self._added_terms.add(text)
+
+            # Generate embedding for term and store in database
+            embedding = await self._vector_base.get_embedding(text)  # Cached
+            serialized_embedding = serialize_embedding(embedding)
+            # Insert term and embedding
+            cursor.execute(
+                """
+                INSERT OR REPLACE INTO RelatedTermsFuzzy
+                (term, term_embedding)
+                VALUES (?, ?)
+                """,
+                (text, serialized_embedding),
+            )
+
+    async def lookup_terms(
+        self,
+        texts: list[str],
+        max_hits: int | None = None,
+        min_score: float | None = None,
+    ) -> list[list[interfaces.Term]]:
+        """Look up multiple terms at once."""
+        # TODO: Some kind of batching?
+        results = []
+        for text in texts:
+            term_results = await self.lookup_term(text, max_hits, min_score)
+            results.append(term_results)
+        return results
+
+    async def deserialize(self, data: interfaces.TextEmbeddingIndexData) -> None:
+        """Deserialize fuzzy index data from JSON into SQLite database."""
+        # Clear existing data
+        cursor = self.db.cursor()
+        cursor.execute("DELETE FROM RelatedTermsFuzzy")
+
+        # Clear local mappings
+        self._terms_list.clear()
+        self._added_terms.clear()
+
+        # Get text items and embeddings from the data
+        text_items = data.get("textItems")
+        embeddings_data = data.get("embeddings")
+
+        if not text_items or embeddings_data is None:
+            return
+
+        # Use persistent VectorBase to deserialize embeddings (preserves caching)
+        self._vector_base.deserialize(embeddings_data)
+
+        # Prepare all insertion data for bulk operation
+        from .schema import serialize_embedding
+
+        insertion_data = []
+        for i, text in enumerate(text_items):
+            if i < len(self._vector_base):
+                # Get embedding from persistent VectorBase
+                embedding = self._vector_base.get_embedding_at(i)
+                if embedding is not None:
+                    serialized_embedding = serialize_embedding(embedding)
+                    # Insert as self-referential entry with only term_embedding
+                    insertion_data.append((text, serialized_embedding))
+                    # Update local mappings
+                    self._terms_list.append(text)
+                    self._added_terms.add(text)
+
+        # Bulk insert all the data
+        if insertion_data:
+            cursor.executemany(
+                """
+                INSERT OR REPLACE INTO RelatedTermsFuzzy
+                (term, term_embedding)
+                VALUES (?, ?)
+                """,
+                insertion_data,
+            )
+
+
+class SqliteRelatedTermsIndex(interfaces.ITermToRelatedTermsIndex):
+    """SQLite-backed implementation of ITermToRelatedTermsIndex combining aliases and fuzzy index."""
+
+    def __init__(self, db: sqlite3.Connection, settings: TextEmbeddingIndexSettings):
+        self.db = db
+        # Initialize alias and fuzzy related terms indexes
+        self._aliases = SqliteRelatedTermsAliases(db)
+        self._fuzzy_index = SqliteRelatedTermsFuzzy(db, settings)
+
+    @property
+    def aliases(self) -> interfaces.ITermToRelatedTerms:
+        return self._aliases
+
+    @property
+    def fuzzy_index(self) -> interfaces.ITermToRelatedTermsFuzzy | None:
+        return self._fuzzy_index
+
+    async def serialize(self) -> interfaces.TermsToRelatedTermsIndexData:
+        raise NotImplementedError("TODO")
+
+    async def deserialize(self, data: interfaces.TermsToRelatedTermsIndexData) -> None:
+        """Deserialize related terms index data."""
+        # Deserialize alias data
+        alias_data = data.get("aliasData")
+        if alias_data is not None:
+            await self._aliases.deserialize(alias_data)
+
+        # Deserialize fuzzy index data
+        text_embedding_data = data.get("textEmbeddingData")
+        if text_embedding_data is not None:
+            await self._fuzzy_index.deserialize(text_embedding_data)