diff --git a/src/typeagent/storage/sqlite/collections.py b/src/typeagent/storage/sqlite/collections.py index 71a256df..d1ec0121 100644 --- a/src/typeagent/storage/sqlite/collections.py +++ b/src/typeagent/storage/sqlite/collections.py @@ -8,27 +8,36 @@ import typing from ...aitools.embeddings import NormalizedEmbedding -from ...knowpro import interfaces, serialization +from ...knowpro import serialization +from ...knowpro.interfaces import ( + IMessage, + IMessageCollection, + IMessageTextIndex, + ISemanticRefCollection, + SemanticRef, + SemanticRefData, + SemanticRefMetadata, + TextLocation, + TextRange, +) from .schema import ShreddedMessage, ShreddedSemanticRef -class SqliteMessageCollection[TMessage: interfaces.IMessage]( - interfaces.IMessageCollection[TMessage] -): +class SqliteMessageCollection[TMessage: IMessage](IMessageCollection[TMessage]): """SQLite-backed message collection.""" def __init__( self, db: sqlite3.Connection, message_type: type[TMessage] | None = None, - message_text_index: "interfaces.IMessageTextIndex[TMessage] | None" = None, + message_text_index: "IMessageTextIndex[TMessage] | None" = None, ): self.db = db self.message_type = message_type self.message_text_index = message_text_index def set_message_text_index( - self, message_text_index: "interfaces.IMessageTextIndex[TMessage]" + self, message_text_index: "IMessageTextIndex[TMessage]" ) -> None: """Set the message text index for automatic indexing of new messages.""" self.message_text_index = message_text_index @@ -248,7 +257,7 @@ async def extend( ) -class SqliteSemanticRefCollection(interfaces.ISemanticRefCollection): +class SqliteSemanticRefCollection(ISemanticRefCollection): """SQLite-backed semantic reference collection.""" def __init__(self, db: sqlite3.Connection): @@ -256,22 +265,22 @@ def __init__(self, db: sqlite3.Connection): def _deserialize_semantic_ref_from_row( self, row: ShreddedSemanticRef - ) -> interfaces.SemanticRef: + ) -> SemanticRef: """Deserialize a semantic ref from database row columns.""" semref_id, range_json, knowledge_type, knowledge_json = row # Build semantic ref data using camelCase (JSON format) - semantic_ref_data = interfaces.SemanticRefData( + semantic_ref_data = SemanticRefData( semanticRefOrdinal=semref_id, range=json.loads(range_json), knowledgeType=knowledge_type, # type: ignore knowledge=json.loads(knowledge_json), ) - return interfaces.SemanticRef.deserialize(semantic_ref_data) + return SemanticRef.deserialize(semantic_ref_data) def _serialize_semantic_ref_to_row( - self, semantic_ref: interfaces.SemanticRef + self, semantic_ref: SemanticRef ) -> ShreddedSemanticRef: """Serialize a semantic ref object into database columns.""" # Serialize the semantic ref to JSON first (this uses camelCase) @@ -297,7 +306,7 @@ def _size(self) -> int: cursor.execute("SELECT COUNT(*) FROM SemanticRefs") return cursor.fetchone()[0] - async def __aiter__(self) -> typing.AsyncGenerator[interfaces.SemanticRef, None]: + async def __aiter__(self) -> typing.AsyncGenerator[SemanticRef, None]: cursor = self.db.cursor() cursor.execute(""" SELECT semref_id, range_json, knowledge_type, knowledge_json @@ -306,7 +315,7 @@ async def __aiter__(self) -> typing.AsyncGenerator[interfaces.SemanticRef, None] for row in cursor: yield self._deserialize_semantic_ref_from_row(row) - async def get_item(self, arg: int) -> interfaces.SemanticRef: + async def get_item(self, arg: int) -> SemanticRef: if not isinstance(arg, int): raise TypeError(f"Index must be an int, not {type(arg).__name__}") cursor = self.db.cursor() @@ -322,7 +331,7 @@ async def get_item(self, arg: int) -> interfaces.SemanticRef: return self._deserialize_semantic_ref_from_row(row) raise IndexError("SemanticRef not found") - async def get_slice(self, start: int, stop: int) -> list[interfaces.SemanticRef]: + async def get_slice(self, start: int, stop: int) -> list[SemanticRef]: if stop <= start: return [] cursor = self.db.cursor() @@ -337,7 +346,7 @@ async def get_slice(self, start: int, stop: int) -> list[interfaces.SemanticRef] rows = cursor.fetchall() return [self._deserialize_semantic_ref_from_row(row) for row in rows] - async def get_multiple(self, arg: list[int]) -> list[interfaces.SemanticRef]: + async def get_multiple(self, arg: list[int]) -> list[SemanticRef]: size = await self.size() if not all((0 <= i < size) for i in arg): raise IndexError("One or more SemanticRef indices are out of bounds") @@ -355,7 +364,7 @@ async def get_multiple(self, arg: list[int]) -> list[interfaces.SemanticRef]: async def get_metadata_multiple( self, ordinals: list[int] - ) -> list[interfaces.SemanticRefMetadata]: + ) -> list[SemanticRefMetadata]: if not ordinals: return [] cursor = self.db.cursor() @@ -376,15 +385,15 @@ async def get_metadata_multiple( start = range_data["start"] end_data = range_data.get("end") result.append( - interfaces.SemanticRefMetadata( + SemanticRefMetadata( ordinal=row[0], - range=interfaces.TextRange( - start=interfaces.TextLocation( + range=TextRange( + start=TextLocation( start["messageOrdinal"], start.get("chunkOrdinal", 0), ), end=( - interfaces.TextLocation( + TextLocation( end_data["messageOrdinal"], end_data.get("chunkOrdinal", 0), ) @@ -397,7 +406,7 @@ async def get_metadata_multiple( ) return result - async def append(self, item: interfaces.SemanticRef) -> None: + async def append(self, item: SemanticRef) -> None: cursor = self.db.cursor() semref_id, range_json, knowledge_type, knowledge_json = ( self._serialize_semantic_ref_to_row(item) @@ -410,7 +419,7 @@ async def append(self, item: interfaces.SemanticRef) -> None: (semref_id, range_json, knowledge_type, knowledge_json), ) - async def extend(self, items: typing.Iterable[interfaces.SemanticRef]) -> None: + async def extend(self, items: typing.Iterable[SemanticRef]) -> None: items_list = list(items) if not items_list: return diff --git a/src/typeagent/storage/sqlite/messageindex.py b/src/typeagent/storage/sqlite/messageindex.py index 457c1ae5..12e7752c 100644 --- a/src/typeagent/storage/sqlite/messageindex.py +++ b/src/typeagent/storage/sqlite/messageindex.py @@ -10,9 +10,17 @@ from ...aitools.embeddings import NormalizedEmbedding from ...aitools.vectorbase import ScoredInt, VectorBase -from ...knowpro import interfaces from ...knowpro.convsettings import MessageTextIndexSettings -from ...knowpro.interfaces import TextLocationData, TextToTextLocationIndexData +from ...knowpro.interfaces import ( + IMessage, + IMessageCollection, + MessageOrdinal, + MessageTextIndexData, + ScoredMessageOrdinal, + TextLocation, + TextLocationData, + TextToTextLocationIndexData, +) from ...knowpro.textlocindex import ScoredTextLocation from ...storage.memory.messageindex import IMessageTextEmbeddingIndex from .schema import deserialize_embedding, serialize_embedding @@ -25,7 +33,7 @@ def __init__( self, db: sqlite3.Connection, settings: MessageTextIndexSettings, - message_collection: interfaces.IMessageCollection | None = None, + message_collection: IMessageCollection | None = None, ): self.db = db self.settings = settings @@ -55,7 +63,7 @@ def _size(self) -> int: async def add_messages_starting_at( self, start_message_ordinal: int, - messages: list[interfaces.IMessage], + messages: list[IMessage], ) -> None: """Add messages to the text index starting at the given ordinal.""" chunks_to_embed: list[str] = [] @@ -80,7 +88,7 @@ async def add_messages_starting_at( async def add_messages_starting_at_with_embeddings( self, start_message_ordinal: int, - messages: list[interfaces.IMessage], + messages: list[IMessage], chunk_embeddings: list[NormalizedEmbedding], ) -> None: """Add messages to the text index using precomputed chunk embeddings.""" @@ -127,7 +135,7 @@ async def add_messages_starting_at_with_embeddings( async def add_messages( self, - messages: typing.Iterable[interfaces.IMessage], + messages: typing.Iterable[IMessage], ) -> None: """Add messages to the text index (backward compatibility method).""" message_list = list(messages) @@ -182,7 +190,7 @@ async def lookup_text( def _vectorbase_lookup_to_scored_locations( self, fuzzy_results: list[ScoredInt], - predicate: typing.Callable[[interfaces.MessageOrdinal], bool] | None = None, + predicate: typing.Callable[[MessageOrdinal], bool] | None = None, ) -> list[ScoredTextLocation]: """Convert VectorBase fuzzy results to scored text locations using optimized DB query.""" if not fuzzy_results: @@ -217,7 +225,7 @@ def _vectorbase_lookup_to_scored_locations( # Apply predicate filter if provided if predicate is None or predicate(msg_id): - text_location = interfaces.TextLocation( + text_location = TextLocation( message_ordinal=msg_id, chunk_ordinal=chunk_ordinal, ) @@ -231,7 +239,7 @@ def _scored_locations_to_message_ordinals( self, scored_locations: list[ScoredTextLocation], max_matches: int | None = None, - ) -> list[interfaces.ScoredMessageOrdinal]: + ) -> list[ScoredMessageOrdinal]: """Convert scored text locations to scored message ordinals by grouping chunks.""" # Group by message and take the best score per message message_scores: dict[int, float] = {} @@ -245,7 +253,7 @@ def _scored_locations_to_message_ordinals( # Convert to list and sort by score result = [ - interfaces.ScoredMessageOrdinal(msg_ordinal, score) + ScoredMessageOrdinal(msg_ordinal, score) for msg_ordinal, score in message_scores.items() ] result.sort(key=lambda x: x.score, reverse=True) @@ -261,7 +269,7 @@ async def lookup_messages( message_text: str, max_matches: int | None = None, threshold_score: float | None = None, - ) -> list[interfaces.ScoredMessageOrdinal]: + ) -> list[ScoredMessageOrdinal]: """Look up messages by text content.""" scored_locations = await self.lookup_text(message_text, None, threshold_score) return self._scored_locations_to_message_ordinals(scored_locations, max_matches) @@ -269,10 +277,10 @@ async def lookup_messages( async def lookup_messages_in_subset( self, message_text: str, - ordinals_to_search: list[interfaces.MessageOrdinal], + ordinals_to_search: list[MessageOrdinal], max_matches: int | None = None, threshold_score: float | None = None, - ) -> list[interfaces.ScoredMessageOrdinal]: + ) -> list[ScoredMessageOrdinal]: """Look up messages in a subset of ordinals.""" # Get all matches first all_matches = await self.lookup_messages(message_text, None, threshold_score) @@ -298,8 +306,8 @@ async def lookup_by_embedding( text_embedding: NormalizedEmbedding, max_matches: int | None = None, threshold_score: float | None = None, - predicate: typing.Callable[[interfaces.MessageOrdinal], bool] | None = None, - ) -> list[interfaces.ScoredMessageOrdinal]: + predicate: typing.Callable[[MessageOrdinal], bool] | None = None, + ) -> list[ScoredMessageOrdinal]: """Look up messages by embedding using optimized VectorBase similarity search.""" fuzzy_results = self._vectorbase.fuzzy_lookup_embedding( text_embedding, max_hits=max_matches, min_score=threshold_score @@ -312,10 +320,10 @@ async def lookup_by_embedding( async def lookup_in_subset_by_embedding( self, text_embedding: NormalizedEmbedding, - ordinals_to_search: list[interfaces.MessageOrdinal], + ordinals_to_search: list[MessageOrdinal], max_matches: int | None = None, threshold_score: float | None = None, - ) -> list[interfaces.ScoredMessageOrdinal]: + ) -> list[ScoredMessageOrdinal]: """Look up messages in a subset by embedding.""" ordinals_set = set(ordinals_to_search) return await self.lookup_by_embedding( @@ -330,7 +338,7 @@ async def is_empty(self) -> bool: size = await self.size() return size == 0 - async def serialize(self) -> interfaces.MessageTextIndexData: + async def serialize(self) -> MessageTextIndexData: """Serialize the message text index.""" # Get all data from the MessageTextIndex table cursor = self.db.cursor() @@ -371,11 +379,11 @@ async def serialize(self) -> interfaces.MessageTextIndexData: index_data = TextToTextLocationIndexData( textLocations=text_locations, embeddings=embeddings_array ) - return interfaces.MessageTextIndexData(indexData=index_data) + return MessageTextIndexData(indexData=index_data) return {} - async def deserialize(self, data: interfaces.MessageTextIndexData) -> None: + async def deserialize(self, data: MessageTextIndexData) -> None: """Deserialize message text index data.""" cursor = self.db.cursor() diff --git a/src/typeagent/storage/sqlite/propindex.py b/src/typeagent/storage/sqlite/propindex.py index 59a5a111..f05eadfd 100644 --- a/src/typeagent/storage/sqlite/propindex.py +++ b/src/typeagent/storage/sqlite/propindex.py @@ -6,15 +6,18 @@ from collections.abc import Sequence import sqlite3 -from ...knowpro import interfaces -from ...knowpro.interfaces import ScoredSemanticRefOrdinal +from ...knowpro.interfaces import ( + IPropertyToSemanticRefIndex, + ScoredSemanticRefOrdinal, + SemanticRefOrdinal, +) from ...storage.memory.propindex import ( make_property_term_text, split_property_term_text, ) -class SqlitePropertyIndex(interfaces.IPropertyToSemanticRefIndex): +class SqlitePropertyIndex(IPropertyToSemanticRefIndex): """SQLite-backed implementation of property to semantic ref index.""" def __init__(self, db: sqlite3.Connection): @@ -38,12 +41,10 @@ async def add_property( self, property_name: str, value: str, - semantic_ref_ordinal: ( - interfaces.SemanticRefOrdinal | interfaces.ScoredSemanticRefOrdinal - ), + semantic_ref_ordinal: SemanticRefOrdinal | ScoredSemanticRefOrdinal, ) -> None: # Extract semref_id and score from the ordinal - if isinstance(semantic_ref_ordinal, interfaces.ScoredSemanticRefOrdinal): + if isinstance(semantic_ref_ordinal, ScoredSemanticRefOrdinal): semref_id = semantic_ref_ordinal.semantic_ref_ordinal score = semantic_ref_ordinal.score else: @@ -73,7 +74,7 @@ async def add_properties_batch( tuple[ str, str, - interfaces.SemanticRefOrdinal | interfaces.ScoredSemanticRefOrdinal, + SemanticRefOrdinal | ScoredSemanticRefOrdinal, ] ], ) -> None: @@ -81,7 +82,7 @@ async def add_properties_batch( return rows = [] for property_name, value, ordinal in properties: - if isinstance(ordinal, interfaces.ScoredSemanticRefOrdinal): + if isinstance(ordinal, ScoredSemanticRefOrdinal): semref_id = ordinal.semantic_ref_ordinal score = ordinal.score else: @@ -107,7 +108,7 @@ async def lookup_property( self, property_name: str, value: str, - ) -> list[interfaces.ScoredSemanticRefOrdinal] | None: + ) -> list[ScoredSemanticRefOrdinal] | None: # Normalize property name and value (to match in-memory implementation) term_text = make_property_term_text(property_name, value) term_text = term_text.lower() # Matches PropertyIndex._prepare_term_text diff --git a/src/typeagent/storage/sqlite/provider.py b/src/typeagent/storage/sqlite/provider.py index a8ba9c06..cedae7a9 100644 --- a/src/typeagent/storage/sqlite/provider.py +++ b/src/typeagent/storage/sqlite/provider.py @@ -8,9 +8,14 @@ from ...aitools.model_adapters import create_embedding_model from ...aitools.vectorbase import TextEmbeddingIndexSettings -from ...knowpro import interfaces from ...knowpro.convsettings import MessageTextIndexSettings, RelatedTermIndexSettings -from ...knowpro.interfaces import ConversationMetadata, STATUS_INGESTED +from ...knowpro.interfaces import ( + ConversationMetadata, + IMessage, + IStorageProvider, + SemanticRef, + STATUS_INGESTED, +) from ...knowpro.interfaces_storage import ChunkFailure from ..memory.convthreads import ConversationThreads from .collections import SqliteMessageCollection, SqliteSemanticRefCollection @@ -27,9 +32,7 @@ from .timestampindex import SqliteTimestampToTextRangeIndex -class SqliteStorageProvider[TMessage: interfaces.IMessage]( - interfaces.IStorageProvider[TMessage] -): +class SqliteStorageProvider[TMessage: IMessage](IStorageProvider[TMessage]): """SQLite-backed storage provider implementation. This provider performs consistency checks on database initialization to ensure @@ -41,7 +44,7 @@ def __init__( self, db_path: str = ":memory:", message_type: type[TMessage] = None, # type: ignore - semantic_ref_type: type[interfaces.SemanticRef] = None, # type: ignore + semantic_ref_type: type[SemanticRef] = None, # type: ignore message_text_index_settings: MessageTextIndexSettings | None = None, related_term_index_settings: RelatedTermIndexSettings | None = None, metadata: ConversationMetadata | None = None, diff --git a/src/typeagent/storage/sqlite/reltermsindex.py b/src/typeagent/storage/sqlite/reltermsindex.py index 1a27af71..d918fd6c 100644 --- a/src/typeagent/storage/sqlite/reltermsindex.py +++ b/src/typeagent/storage/sqlite/reltermsindex.py @@ -9,27 +9,37 @@ from ...aitools.embeddings import NormalizedEmbedding from ...aitools.vectorbase import TextEmbeddingIndexSettings, VectorBase -from ...knowpro import interfaces +from ...knowpro.interfaces import ( + ITermToRelatedTerms, + ITermToRelatedTermsFuzzy, + ITermToRelatedTermsIndex, + Term, + TermData, + TermsToRelatedTermsDataItem, + TermsToRelatedTermsIndexData, + TermToRelatedTermsData, + TextEmbeddingIndexData, +) from .schema import deserialize_embedding, serialize_embedding -class SqliteRelatedTermsAliases(interfaces.ITermToRelatedTerms): +class SqliteRelatedTermsAliases(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: + async def lookup_term(self, text: str) -> list[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()] + results = [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] + self, text: str, related_terms: Term | list[Term] ) -> None: # Convert single Term to list - if isinstance(related_terms, interfaces.Term): + if isinstance(related_terms, Term): related_terms = [related_terms] cursor = self.db.cursor() @@ -74,7 +84,7 @@ async def is_empty(self) -> bool: cursor.execute("SELECT COUNT(*) FROM RelatedTermsAliases") return cursor.fetchone()[0] == 0 - async def serialize(self) -> interfaces.TermToRelatedTermsData: + async def serialize(self) -> TermToRelatedTermsData: """Serialize the aliases data.""" cursor = self.db.cursor() cursor.execute( @@ -91,16 +101,14 @@ async def serialize(self) -> interfaces.TermToRelatedTermsData: # 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] + term_data_list = [TermData(text=alias) for alias in aliases] items.append( - interfaces.TermsToRelatedTermsDataItem( - termText=term, relatedTerms=term_data_list - ) + TermsToRelatedTermsDataItem(termText=term, relatedTerms=term_data_list) ) - return interfaces.TermToRelatedTermsData(relatedTerms=items) + return TermToRelatedTermsData(relatedTerms=items) - async def deserialize(self, data: interfaces.TermToRelatedTermsData | None) -> None: + async def deserialize(self, data: TermToRelatedTermsData | None) -> None: """Deserialize alias data.""" cursor = self.db.cursor() @@ -130,7 +138,7 @@ async def deserialize(self, data: interfaces.TermToRelatedTermsData | None) -> N ) -class SqliteRelatedTermsFuzzy(interfaces.ITermToRelatedTermsFuzzy): +class SqliteRelatedTermsFuzzy(ITermToRelatedTermsFuzzy): """SQLite-backed implementation of fuzzy term relationships with persistent embeddings.""" def __init__(self, db: sqlite3.Connection, settings: TextEmbeddingIndexSettings): @@ -160,7 +168,7 @@ async def lookup_term( text: str, max_hits: int | None = None, min_score: float | None = None, - ) -> list[interfaces.Term]: + ) -> list[Term]: """Look up similar terms using fuzzy matching.""" # Search for similar terms using VectorBase @@ -174,7 +182,7 @@ async def lookup_term( # 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)) + results.append(Term(term_text, scored_int.score)) return results @@ -261,7 +269,7 @@ async def lookup_terms( texts: list[str], max_hits: int | None = None, min_score: float | None = None, - ) -> list[list[interfaces.Term]]: + ) -> list[list[Term]]: """Look up multiple terms at once.""" # TODO: Some kind of batching? results = [] @@ -270,14 +278,14 @@ async def lookup_terms( results.append(term_results) return results - def serialize(self) -> interfaces.TextEmbeddingIndexData: + def serialize(self) -> TextEmbeddingIndexData: """Serialize the fuzzy index data.""" - return interfaces.TextEmbeddingIndexData( + return TextEmbeddingIndexData( textItems=self._terms_list.copy(), embeddings=self._vector_base.serialize(), ) - async def deserialize(self, data: interfaces.TextEmbeddingIndexData) -> None: + async def deserialize(self, data: TextEmbeddingIndexData) -> None: """Deserialize fuzzy index data from JSON into SQLite database.""" # Clear existing data cursor = self.db.cursor() @@ -325,7 +333,7 @@ async def deserialize(self, data: interfaces.TextEmbeddingIndexData) -> None: ) -class SqliteRelatedTermsIndex(interfaces.ITermToRelatedTermsIndex): +class SqliteRelatedTermsIndex(ITermToRelatedTermsIndex): """SQLite-backed implementation of ITermToRelatedTermsIndex combining aliases and fuzzy index.""" def __init__(self, db: sqlite3.Connection, settings: TextEmbeddingIndexSettings): @@ -335,21 +343,21 @@ def __init__(self, db: sqlite3.Connection, settings: TextEmbeddingIndexSettings) self._fuzzy_index = SqliteRelatedTermsFuzzy(db, settings) @property - def aliases(self) -> interfaces.ITermToRelatedTerms: + def aliases(self) -> ITermToRelatedTerms: return self._aliases @property - def fuzzy_index(self) -> interfaces.ITermToRelatedTermsFuzzy | None: + def fuzzy_index(self) -> ITermToRelatedTermsFuzzy | None: return self._fuzzy_index - async def serialize(self) -> interfaces.TermsToRelatedTermsIndexData: + async def serialize(self) -> TermsToRelatedTermsIndexData: """Serialize the related terms index (both aliases and fuzzy index).""" - return interfaces.TermsToRelatedTermsIndexData( + return TermsToRelatedTermsIndexData( aliasData=await self._aliases.serialize(), textEmbeddingData=self._fuzzy_index.serialize(), ) - async def deserialize(self, data: interfaces.TermsToRelatedTermsIndexData) -> None: + async def deserialize(self, data: TermsToRelatedTermsIndexData) -> None: """Deserialize related terms index data.""" # Deserialize alias data alias_data = data.get("aliasData") diff --git a/src/typeagent/storage/sqlite/semrefindex.py b/src/typeagent/storage/sqlite/semrefindex.py index ac68a1e0..f462955e 100644 --- a/src/typeagent/storage/sqlite/semrefindex.py +++ b/src/typeagent/storage/sqlite/semrefindex.py @@ -8,11 +8,17 @@ import sqlite3 import unicodedata -from ...knowpro import interfaces -from ...knowpro.interfaces import ScoredSemanticRefOrdinal +from ...knowpro.interfaces import ( + ITermToSemanticRefIndex, + ScoredSemanticRefOrdinal, + ScoredSemanticRefOrdinalData, + SemanticRefOrdinal, + TermToSemanticRefIndexData, + TermToSemanticRefIndexItemData, +) -class SqliteTermToSemanticRefIndex(interfaces.ITermToSemanticRefIndex): +class SqliteTermToSemanticRefIndex(ITermToSemanticRefIndex): """SQLite-backed implementation of term to semantic ref index.""" def __init__(self, db: sqlite3.Connection): @@ -31,9 +37,7 @@ async def get_terms(self) -> list[str]: async def add_term( self, term: str, - semantic_ref_ordinal: ( - interfaces.SemanticRefOrdinal | interfaces.ScoredSemanticRefOrdinal - ), + semantic_ref_ordinal: SemanticRefOrdinal | ScoredSemanticRefOrdinal, ) -> str: if not term: return term @@ -41,7 +45,7 @@ async def add_term( term = self._prepare_term(term) # Extract semref_id from the ordinal - if isinstance(semantic_ref_ordinal, interfaces.ScoredSemanticRefOrdinal): + if isinstance(semantic_ref_ordinal, ScoredSemanticRefOrdinal): semref_id = semantic_ref_ordinal.semantic_ref_ordinal else: semref_id = semantic_ref_ordinal @@ -59,11 +63,7 @@ async def add_term( async def add_terms_batch( self, - terms: Sequence[ - tuple[ - str, interfaces.SemanticRefOrdinal | interfaces.ScoredSemanticRefOrdinal - ] - ], + terms: Sequence[tuple[str, SemanticRefOrdinal | ScoredSemanticRefOrdinal]], ) -> None: if not terms: return @@ -72,7 +72,7 @@ async def add_terms_batch( if not term: continue term = self._prepare_term(term) - if isinstance(ordinal, interfaces.ScoredSemanticRefOrdinal): + if isinstance(ordinal, ScoredSemanticRefOrdinal): semref_id = ordinal.semantic_ref_ordinal else: semref_id = ordinal @@ -85,7 +85,7 @@ async def add_terms_batch( ) async def remove_term( - self, term: str, semantic_ref_ordinal: interfaces.SemanticRefOrdinal + self, term: str, semantic_ref_ordinal: SemanticRefOrdinal ) -> None: term = self._prepare_term(term) cursor = self.db.cursor() @@ -94,9 +94,7 @@ async def remove_term( (term, semantic_ref_ordinal), ) - async def lookup_term( - self, term: str - ) -> list[interfaces.ScoredSemanticRefOrdinal] | None: + async def lookup_term(self, term: str) -> list[ScoredSemanticRefOrdinal] | None: term = self._prepare_term(term) cursor = self.db.cursor() cursor.execute( @@ -116,7 +114,7 @@ async def clear(self) -> None: cursor = self.db.cursor() cursor.execute("DELETE FROM SemanticRefIndex") - async def serialize(self) -> interfaces.TermToSemanticRefIndexData: + async def serialize(self) -> TermToSemanticRefIndexData: """Serialize the index data for compatibility with in-memory version.""" cursor = self.db.cursor() cursor.execute( @@ -124,7 +122,7 @@ async def serialize(self) -> interfaces.TermToSemanticRefIndexData: ) # Group by term - term_to_semrefs: dict[str, list[interfaces.ScoredSemanticRefOrdinalData]] = {} + term_to_semrefs: dict[str, list[ScoredSemanticRefOrdinalData]] = {} for term, semref_id in cursor.fetchall(): if term not in term_to_semrefs: term_to_semrefs[term] = [] @@ -135,14 +133,14 @@ async def serialize(self) -> interfaces.TermToSemanticRefIndexData: items = [] for term, semref_ordinals in term_to_semrefs.items(): items.append( - interfaces.TermToSemanticRefIndexItemData( + TermToSemanticRefIndexItemData( term=term, semanticRefOrdinals=semref_ordinals ) ) - return interfaces.TermToSemanticRefIndexData(items=items) + return TermToSemanticRefIndexData(items=items) - async def deserialize(self, data: interfaces.TermToSemanticRefIndexData) -> None: + async def deserialize(self, data: TermToSemanticRefIndexData) -> None: """Deserialize index data by populating the SQLite table.""" cursor = self.db.cursor() diff --git a/src/typeagent/storage/sqlite/timestampindex.py b/src/typeagent/storage/sqlite/timestampindex.py index 8fe017dd..a5394c59 100644 --- a/src/typeagent/storage/sqlite/timestampindex.py +++ b/src/typeagent/storage/sqlite/timestampindex.py @@ -5,11 +5,18 @@ import sqlite3 -from ...knowpro import interfaces +from ...knowpro.interfaces import ( + DateRange, + ITimestampToTextRangeIndex, + MessageOrdinal, + TextLocation, + TextRange, + TimestampedTextRange, +) from ...knowpro.universal_message import format_timestamp_utc -class SqliteTimestampToTextRangeIndex(interfaces.ITimestampToTextRangeIndex): +class SqliteTimestampToTextRangeIndex(ITimestampToTextRangeIndex): """SQL-based timestamp index that queries Messages table directly.""" def __init__(self, db: sqlite3.Connection): @@ -26,13 +33,11 @@ def _size(self) -> int: return cursor.fetchone()[0] async def add_timestamp( - self, message_ordinal: interfaces.MessageOrdinal, timestamp: str + self, message_ordinal: MessageOrdinal, timestamp: str ) -> bool: return self._add_timestamp(message_ordinal, timestamp) - def _add_timestamp( - self, message_ordinal: interfaces.MessageOrdinal, timestamp: str - ) -> bool: + def _add_timestamp(self, message_ordinal: MessageOrdinal, timestamp: str) -> bool: """Add timestamp to Messages table start_timestamp column.""" cursor = self.db.cursor() cursor.execute( @@ -43,7 +48,7 @@ def _add_timestamp( async def get_timestamp_ranges( self, start_timestamp: str, end_timestamp: str | None = None - ) -> list[interfaces.TimestampedTextRange]: + ) -> list[TimestampedTextRange]: """Get timestamp ranges from Messages table.""" cursor = self.db.cursor() @@ -78,14 +83,12 @@ async def get_timestamp_ranges( text_range = TextRange( start=TextLocation(message_ordinal=msg_id, chunk_ordinal=0) ) - results.append( - interfaces.TimestampedTextRange(range=text_range, timestamp=timestamp) - ) + results.append(TimestampedTextRange(range=text_range, timestamp=timestamp)) return results async def add_timestamps( - self, message_timestamps: list[tuple[interfaces.MessageOrdinal, str]] + self, message_timestamps: list[tuple[MessageOrdinal, str]] ) -> None: """Add multiple timestamps.""" if not message_timestamps: @@ -96,9 +99,7 @@ async def add_timestamps( [(ts, ordinal) for ordinal, ts in message_timestamps], ) - async def lookup_range( - self, date_range: interfaces.DateRange - ) -> list[interfaces.TimestampedTextRange]: + async def lookup_range(self, date_range: DateRange) -> list[TimestampedTextRange]: """Lookup messages in a date range.""" cursor = self.db.cursor() @@ -131,14 +132,8 @@ async def lookup_range( results = [] for msg_id, timestamp, _chunks in cursor.fetchall(): - text_location = interfaces.TextLocation( - message_ordinal=msg_id, chunk_ordinal=0 - ) - text_range = interfaces.TextRange( - start=text_location, end=None # Point range - ) - results.append( - interfaces.TimestampedTextRange(timestamp=timestamp, range=text_range) - ) + text_location = TextLocation(message_ordinal=msg_id, chunk_ordinal=0) + text_range = TextRange(start=text_location, end=None) # Point range + results.append(TimestampedTextRange(timestamp=timestamp, range=text_range)) return results