typeagent-py 0.1.0__py3-none-any.whl

typeagent/storage/memory/propindex.py

@@ -0,0 +1,289 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import enum
+from typing import assert_never
+
+from ...knowpro.collections import TextRangesInScope
+from ...knowpro.interfaces import (
+    IConversation,
+    IPropertyToSemanticRefIndex,
+    ISemanticRefCollection,
+    ScoredSemanticRefOrdinal,
+    SemanticRefOrdinal,
+    Tag,
+    Topic,
+)
+from ...knowpro import kplib
+
+
+class PropertyNames(enum.Enum):
+    EntityName = "name"
+    EntityType = "type"
+    FacetName = "facet.name"
+    FacetValue = "facet.value"
+    Verb = "verb"
+    Subject = "subject"
+    Object = "object"
+    IndirectObject = "indirectObject"
+    Tag = "tag"
+    Topic = "topic"
+
+
+async def add_facet(
+    facet: kplib.Facet | None,
+    property_index: IPropertyToSemanticRefIndex,
+    semantic_ref_ordinal: SemanticRefOrdinal,
+) -> None:
+    if facet is not None:
+        await property_index.add_property(
+            PropertyNames.FacetName.value,
+            facet.name,
+            semantic_ref_ordinal,
+        )
+        value = facet.value
+        if value is not None:
+            # If the value is a float, we use .g format store it as a string.
+            if isinstance(value, float) and value:
+                value = f"{value:g}"
+            await property_index.add_property(
+                PropertyNames.FacetValue.value,
+                str(value),
+                semantic_ref_ordinal,
+            )
+
+
+async def add_entity_properties_to_index(
+    entity: kplib.ConcreteEntity,
+    property_index: IPropertyToSemanticRefIndex,
+    semantic_ref_ordinal: SemanticRefOrdinal,
+) -> None:
+    await property_index.add_property(
+        PropertyNames.EntityName.value,
+        entity.name,
+        semantic_ref_ordinal,
+    )
+    for type in entity.type:
+        await property_index.add_property(
+            PropertyNames.EntityType.value,
+            type,
+            semantic_ref_ordinal,
+        )
+    # Add every facet name as a separate term.
+    if entity.facets:
+        for facet in entity.facets:
+            await add_facet(facet, property_index, semantic_ref_ordinal)
+
+
+async def add_action_properties_to_index(
+    action: kplib.Action,
+    property_index: IPropertyToSemanticRefIndex,
+    semantic_ref_ordinal: SemanticRefOrdinal,
+) -> None:
+    await property_index.add_property(
+        PropertyNames.Verb.value,
+        " ".join(action.verbs),
+        semantic_ref_ordinal,
+    )
+    if action.subject_entity_name != "none":
+        await property_index.add_property(
+            PropertyNames.Subject.value,
+            action.subject_entity_name,
+            semantic_ref_ordinal,
+        )
+    if action.object_entity_name != "none":
+        await property_index.add_property(
+            PropertyNames.Object.value,
+            action.object_entity_name,
+            semantic_ref_ordinal,
+        )
+    if action.indirect_object_entity_name != "none":
+        await property_index.add_property(
+            PropertyNames.IndirectObject.value,
+            action.indirect_object_entity_name,
+            semantic_ref_ordinal,
+        )
+
+
+async def build_property_index(conversation: IConversation) -> None:
+    await add_to_property_index(conversation, 0)
+
+
+async def add_to_property_index(
+    conversation: IConversation,
+    start_at_ordinal: SemanticRefOrdinal,
+) -> None:
+    """Add semantic references from a conversation to the property index starting at a specific ordinal."""
+    if (
+        csi := conversation.secondary_indexes
+    ) and conversation.semantic_refs is not None:
+        # Check if semantic_refs collection is not empty
+        if await conversation.semantic_refs.size() == 0:
+            return
+
+        if (property_index := csi.property_to_semantic_ref_index) is None:
+            property_index = csi.property_to_semantic_ref_index = PropertyIndex()
+
+        semantic_refs = conversation.semantic_refs
+        size = await semantic_refs.size()
+
+        for semantic_ref_ordinal, semantic_ref in enumerate(
+            await semantic_refs.get_slice(start_at_ordinal, size),
+            start_at_ordinal,
+        ):
+            assert semantic_ref.semantic_ref_ordinal == semantic_ref_ordinal
+            if isinstance(semantic_ref.knowledge, kplib.Action):
+                await add_action_properties_to_index(
+                    semantic_ref.knowledge, property_index, semantic_ref_ordinal
+                )
+            elif isinstance(semantic_ref.knowledge, kplib.ConcreteEntity):
+                await add_entity_properties_to_index(
+                    semantic_ref.knowledge, property_index, semantic_ref_ordinal
+                )
+            elif isinstance(semantic_ref.knowledge, Tag):
+                tag = semantic_ref.knowledge
+                await property_index.add_property(
+                    PropertyNames.Tag.value, tag.text, semantic_ref_ordinal
+                )
+            elif isinstance(semantic_ref.knowledge, Topic):
+                pass
+            else:
+                assert_never(semantic_ref.knowledge)
+
+
+class PropertyIndex(IPropertyToSemanticRefIndex):
+    def __init__(self):
+        self._map: dict[str, list[ScoredSemanticRefOrdinal]] = {}
+
+    async def size(self) -> int:
+        return len(self._map)
+
+    async def get_values(self) -> list[str]:
+        terms: list[str] = []
+        for key in self._map.keys():
+            nv = split_property_term_text(key)
+            terms.append(nv[1])
+        return terms
+
+    async def add_property(
+        self,
+        property_name: str,
+        value: str,
+        semantic_ref_ordinal: SemanticRefOrdinal | ScoredSemanticRefOrdinal,
+    ) -> None:
+        term_text = make_property_term_text(property_name, value)
+        if isinstance(semantic_ref_ordinal, int):
+            semantic_ref_ordinal = ScoredSemanticRefOrdinal(
+                semantic_ref_ordinal,
+                1.0,
+            )
+        term_text = self._prepare_term_text(term_text)
+        if term_text in self._map:
+            self._map[term_text].append(semantic_ref_ordinal)
+        else:
+            self._map[term_text] = [semantic_ref_ordinal]
+
+    async def clear(self) -> None:
+        self._map = {}
+
+    async def lookup_property(
+        self,
+        property_name: str,
+        value: str,
+    ) -> list[ScoredSemanticRefOrdinal] | None:
+        term_text = make_property_term_text(property_name, value)
+        return self._map.get(self._prepare_term_text(term_text))
+
+    async def remove_property(self, prop_name: str, semref_id: int) -> None:
+        """Remove all properties for a specific property name and semantic ref."""
+        # Find and remove entries matching both property name and semref_id
+        keys_to_remove = []
+        for term_text, scored_refs in self._map.items():
+            prop_name_from_term, _ = split_property_term_text(term_text)
+            # Remove "prop." prefix
+            if prop_name_from_term.startswith("prop."):
+                prop_name_from_term = prop_name_from_term[5:]
+
+            if prop_name_from_term == prop_name:
+                # Filter out entries with matching semref_id
+                filtered_refs = [
+                    ref for ref in scored_refs if ref.semantic_ref_ordinal != semref_id
+                ]
+                if filtered_refs:
+                    self._map[term_text] = filtered_refs
+                else:
+                    keys_to_remove.append(term_text)
+
+        # Remove empty entries
+        for key in keys_to_remove:
+            del self._map[key]
+
+    async def remove_all_for_semref(self, semref_id: int) -> None:
+        """Remove all properties for a specific semantic ref."""
+        keys_to_remove = []
+        for term_text, scored_refs in self._map.items():
+            # Filter out entries with matching semref_id
+            filtered_refs = [
+                ref for ref in scored_refs if ref.semantic_ref_ordinal != semref_id
+            ]
+            if filtered_refs:
+                self._map[term_text] = filtered_refs
+            else:
+                keys_to_remove.append(term_text)
+
+        # Remove empty entries
+        for key in keys_to_remove:
+            del self._map[key]
+
+    def _prepare_term_text(self, term_text: str) -> str:
+        """Do any pre-processing of the term."""
+        return term_text.lower()
+
+
+async def lookup_property_in_property_index(
+    property_index: IPropertyToSemanticRefIndex,
+    property_name: str,
+    property_value: str,
+    semantic_refs: ISemanticRefCollection,
+    ranges_in_scope: TextRangesInScope | None = None,
+) -> list[ScoredSemanticRefOrdinal] | None:
+    scored_refs = await property_index.lookup_property(
+        property_name,
+        property_value,
+    )
+    if ranges_in_scope is not None and scored_refs:
+        filtered_refs = []
+        for sr in scored_refs:
+            semantic_ref = await semantic_refs.get_item(sr.semantic_ref_ordinal)
+            if ranges_in_scope.is_range_in_scope(semantic_ref.range):
+                filtered_refs.append(sr)
+        scored_refs = filtered_refs
+
+    return scored_refs or None  # Return None if no results
+
+
+async def is_known_property(
+    property_index: IPropertyToSemanticRefIndex | None,
+    property_name: PropertyNames,
+    property_value: str,
+) -> bool:
+    if property_index is not None:
+        semantic_refs_with_name = await property_index.lookup_property(
+            property_name.value,
+            property_value,
+        )
+        return semantic_refs_with_name is not None and len(semantic_refs_with_name) > 0
+    else:
+        return False
+
+
+PROPERTY_DELIMITER = "@@"
+
+
+def make_property_term_text(name: str, value: str) -> str:
+    return f"prop.{name}{PROPERTY_DELIMITER}{value}"
+
+
+def split_property_term_text(term_text: str) -> tuple[str, str]:
+    parts = term_text.split(PROPERTY_DELIMITER, 1)
+    return parts[0], parts[1]

typeagent/storage/memory/provider.py

@@ -0,0 +1,84 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Memory storage provider implementation."""
+
+from .collections import MemoryMessageCollection, MemorySemanticRefCollection
+from .semrefindex import TermToSemanticRefIndex
+from .convthreads import ConversationThreads
+from .messageindex import MessageTextIndex
+from .reltermsindex import RelatedTermsIndex
+from .propindex import PropertyIndex
+from .timestampindex import TimestampToTextRangeIndex
+from ...knowpro.convsettings import MessageTextIndexSettings, RelatedTermIndexSettings
+from ...knowpro.interfaces import (
+    IConversationThreads,
+    IMessage,
+    IMessageTextIndex,
+    IPropertyToSemanticRefIndex,
+    IStorageProvider,
+    ITermToRelatedTermsIndex,
+    ITermToSemanticRefIndex,
+    ITimestampToTextRangeIndex,
+)
+
+
+class MemoryStorageProvider[TMessage: IMessage](IStorageProvider[TMessage]):
+    """A storage provider that operates in memory."""
+
+    _message_collection: MemoryMessageCollection[TMessage]
+    _semantic_ref_collection: MemorySemanticRefCollection
+
+    _conversation_index: TermToSemanticRefIndex
+    _property_index: PropertyIndex
+    _timestamp_index: TimestampToTextRangeIndex
+    _message_text_index: MessageTextIndex
+    _related_terms_index: RelatedTermsIndex
+    _conversation_threads: ConversationThreads
+
+    def __init__(
+        self,
+        message_text_settings: MessageTextIndexSettings,
+        related_terms_settings: RelatedTermIndexSettings,
+    ) -> None:
+        """Create and initialize a MemoryStorageProvider with all indexes."""
+        self._message_collection = MemoryMessageCollection[TMessage]()
+        self._semantic_ref_collection = MemorySemanticRefCollection()
+
+        self._conversation_index = TermToSemanticRefIndex()
+        self._property_index = PropertyIndex()
+        self._timestamp_index = TimestampToTextRangeIndex()
+        self._message_text_index = MessageTextIndex(message_text_settings)
+        self._related_terms_index = RelatedTermsIndex(related_terms_settings)
+        thread_settings = message_text_settings.embedding_index_settings
+        self._conversation_threads = ConversationThreads(thread_settings)
+
+    async def get_semantic_ref_index(self) -> ITermToSemanticRefIndex:
+        return self._conversation_index
+
+    async def get_property_index(self) -> IPropertyToSemanticRefIndex:
+        return self._property_index
+
+    async def get_timestamp_index(self) -> ITimestampToTextRangeIndex:
+        return self._timestamp_index
+
+    async def get_message_text_index(self) -> IMessageTextIndex[TMessage]:
+        return self._message_text_index
+
+    async def get_related_terms_index(self) -> ITermToRelatedTermsIndex:
+        return self._related_terms_index
+
+    async def get_conversation_threads(self) -> IConversationThreads:
+        return self._conversation_threads
+
+    async def get_message_collection(
+        self, message_type: type[TMessage] | None = None
+    ) -> MemoryMessageCollection[TMessage]:
+        return self._message_collection
+
+    async def get_semantic_ref_collection(self) -> MemorySemanticRefCollection:
+        return self._semantic_ref_collection
+
+    async def close(self) -> None:
+        """Close the storage provider."""
+        pass

typeagent/storage/memory/reltermsindex.py

@@ -0,0 +1,318 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Protocol, TYPE_CHECKING
+
+from typeagent.aitools.vectorbase import (
+    ScoredInt,
+    TextEmbeddingIndexSettings,
+    VectorBase,
+)
+
+from typeagent.knowpro.collections import TermSet
+from typeagent.knowpro.common import is_search_term_wildcard
+from typeagent.knowpro.convsettings import RelatedTermIndexSettings
+from typeagent.knowpro.interfaces import (
+    IConversation,
+    IMessage,
+    ITermToRelatedTerms,
+    ITermToRelatedTermsFuzzy,
+    ITermToRelatedTermsIndex,
+    ITermToSemanticRefIndex,
+    SearchTerm,
+    Term,
+    TermToRelatedTermsData,
+    TermsToRelatedTermsDataItem,
+    TermsToRelatedTermsIndexData,
+    TextEmbeddingIndexData,
+)
+
+if TYPE_CHECKING:
+    from typeagent.knowpro.query import CompiledSearchTerm, CompiledTermGroup
+
+
+class TermToRelatedTermsMap(ITermToRelatedTerms):
+    def __init__(self):
+        # The inner dict represents a set of terms disregarding their weights.
+        self.map: dict[str, dict[str, Term]] = {}
+
+    async def add_related_term(
+        self, text: str, related_terms: Term | list[Term]
+    ) -> None:
+        if not isinstance(related_terms, list):
+            related_terms = [related_terms]
+        terms: dict[str, Term] = self.map.setdefault(text, {})
+        for related in related_terms:
+            terms.setdefault(related.text, related)
+
+    async def lookup_term(self, text: str) -> list[Term] | None:
+        result = self.map.get(text)
+        if result:
+            return list(result.values())
+        else:
+            return None
+
+    async def remove_term(self, text: str) -> None:
+        self.map.pop(text, None)
+
+    async def clear(self) -> None:
+        self.map.clear()
+
+    async def size(self) -> int:
+        return len(self.map)
+
+    async def is_empty(self) -> bool:
+        return len(self.map) == 0
+
+    async def serialize(self) -> TermToRelatedTermsData:
+        related_terms: list[TermsToRelatedTermsDataItem] = []
+        for key, value in self.map.items():
+            related_terms.append(
+                TermsToRelatedTermsDataItem(
+                    termText=key,
+                    relatedTerms=[term.serialize() for term in value.values()],
+                )
+            )
+        return TermToRelatedTermsData(relatedTerms=related_terms)
+
+    async def deserialize(self, data: TermToRelatedTermsData | None) -> None:
+        self.map.clear()
+        if data is None:
+            return
+        related_terms_data = data.get("relatedTerms")
+        if related_terms_data is None:
+            return
+        for item in related_terms_data:
+            term_text = item["termText"]
+            related_terms_data = item["relatedTerms"]
+            related_terms: list[Term] = [
+                Term(term_data["text"], weight=term_data.get("weight"))
+                for term_data in related_terms_data
+            ]
+            await self.add_related_term(term_text, related_terms)
+
+
+async def build_related_terms_index[
+    TMessage: IMessage,
+    TTermToSemanticRefIndex: ITermToSemanticRefIndex,
+](
+    conversation: IConversation[TMessage, TTermToSemanticRefIndex],
+    settings: RelatedTermIndexSettings,
+) -> None:
+    csr = conversation.semantic_ref_index
+    assert csr is not None
+    csi = conversation.secondary_indexes
+    if csr is not None and csi is not None:
+        if csi.term_to_related_terms_index is None:
+            csi.term_to_related_terms_index = RelatedTermsIndex(settings)
+        fuzzy_index = csi.term_to_related_terms_index.fuzzy_index
+        if fuzzy_index is not None:
+            all_terms = await csr.get_terms()
+            if all_terms:
+                await fuzzy_index.add_terms(all_terms)
+
+
+class RelatedTermsIndex(ITermToRelatedTermsIndex):
+    def __init__(self, settings: RelatedTermIndexSettings):
+        self.settings = settings
+        self._alias_map = TermToRelatedTermsMap()
+        self._term_index = TermEmbeddingIndex(settings.embedding_index_settings)
+
+    @property
+    def aliases(self) -> TermToRelatedTermsMap:
+        return self._alias_map
+
+    @property
+    def fuzzy_index(self) -> ITermToRelatedTermsFuzzy | None:
+        return self._term_index
+
+    async def serialize(self) -> TermsToRelatedTermsIndexData:
+        return TermsToRelatedTermsIndexData(
+            aliasData=await self._alias_map.serialize(),
+            textEmbeddingData=self._term_index.serialize(),
+        )
+
+    async def deserialize(self, data: TermsToRelatedTermsIndexData) -> None:
+        await self._alias_map.clear()
+        self._term_index.clear()
+        await self._alias_map.deserialize(data.get("aliasData"))
+        text_embedding_data = data.get("textEmbeddingData")
+        if text_embedding_data is not None:
+            self._term_index.deserialize(text_embedding_data)
+
+
+async def resolve_related_terms(
+    related_terms_index: ITermToRelatedTermsIndex,
+    compiled_terms: list["CompiledTermGroup"],
+    ensure_single_occurrence: bool = True,
+    should_resolve_fuzzy: Callable[[SearchTerm], bool] | None = None,
+) -> None:
+    """Resolves related terms for those search terms that don't already have them.
+
+    NOTE: This modifies SearchTerm().related_terms in place.
+
+    Optionally ensures that related terms are not duplicated across search terms
+    because this can skew how semantic references are scored during search
+    (over-counting).
+
+    SUBTLE: If a search terms has related_terms == [], don't touch it;
+    only set related_terms if it is None.
+    """
+    all_search_terms = [term for ct in compiled_terms for term in ct.terms]
+    searchable_terms = TermSet()
+    search_terms_needing_related: list[SearchTerm] = []
+
+    for search_term in all_search_terms:
+        if is_search_term_wildcard(search_term):
+            continue
+        searchable_terms.add_or_union(search_term.term)
+        term_text = search_term.term.text
+        # Resolve any specific term to related term mappings
+        if search_term.related_terms is None:
+            search_term.related_terms = await related_terms_index.aliases.lookup_term(
+                term_text
+            )
+        # If no mappings to aliases, add to fuzzy retrieval list
+        if search_term.related_terms is None:
+            if should_resolve_fuzzy is None or should_resolve_fuzzy(search_term):
+                search_terms_needing_related.append(search_term)
+
+    if related_terms_index.fuzzy_index is not None and search_terms_needing_related:
+        related_terms_for_search_terms = (
+            await related_terms_index.fuzzy_index.lookup_terms(
+                [st.term.text for st in search_terms_needing_related]
+            )
+        )
+        for search_term, related_terms in zip(
+            search_terms_needing_related, related_terms_for_search_terms
+        ):
+            search_term.related_terms = related_terms
+
+    # Due to fuzzy matching, a search term may end with related terms that overlap with those of other search terms.
+    # This causes scoring problems... duplicate/redundant scoring that can cause items to seem more relevant than they are
+    # - The same related term can show up for different search terms but with different weights
+    # - related terms may also already be present as search terms
+    for ct in compiled_terms:
+        dedupe_related_terms(
+            ct.terms, ensure_single_occurrence and ct.boolean_op != "and"
+        )
+
+
+def dedupe_related_terms(
+    compiled_terms: list["CompiledSearchTerm"],
+    ensure_single_occurrence: bool,
+) -> None:
+    all_search_terms = TermSet()
+    all_related_terms: TermSet | None = None
+
+    # Collect all unique search and related terms.
+    # We end up with (term, maximum weight for term) pairs.
+    for st in compiled_terms:
+        all_search_terms.add(st.term)
+    if ensure_single_occurrence:
+        all_related_terms = TermSet()
+        for st in compiled_terms:
+            all_related_terms.add_or_union(st.related_terms)
+
+    for search_term in compiled_terms:
+        required = search_term.related_terms_required
+        if required:
+            continue
+        if search_term.related_terms:
+            unique_related_for_search_term: list[Term] = []
+            for candidate_related_term in search_term.related_terms:
+                if candidate_related_term in all_search_terms:
+                    # This related term is already a search term
+                    continue
+                if ensure_single_occurrence and all_related_terms is not None:
+                    # Each unique related term should be searched for only once,
+                    # and (if there were duplicates) assigned the maximum weight assigned to that term
+                    term_with_max_weight = all_related_terms.get(candidate_related_term)
+                    if (
+                        term_with_max_weight is not None
+                        and term_with_max_weight.weight == candidate_related_term.weight
+                    ):
+                        # Associate this related term with the current search term
+                        unique_related_for_search_term.append(term_with_max_weight)
+                        all_related_terms.remove(candidate_related_term)
+                else:
+                    unique_related_for_search_term.append(candidate_related_term)
+            search_term.related_terms = unique_related_for_search_term
+
+
+class ITermEmbeddingIndex(ITermToRelatedTermsFuzzy, Protocol):
+    def serialize(self) -> TextEmbeddingIndexData: ...
+
+    def deserialize(self, data: TextEmbeddingIndexData) -> None: ...
+
+
+# TODO: Inherit from TextEmbeddingCache too.
+class TermEmbeddingIndex(ITermEmbeddingIndex):
+    # The Python version wraps a VectorBase
+
+    settings: TextEmbeddingIndexSettings
+    _vectorbase: VectorBase
+    _texts: list[str]
+
+    def __init__(
+        self,
+        settings: TextEmbeddingIndexSettings,
+        data: TextEmbeddingIndexData | None = None,
+    ):
+        self.settings = settings
+        self._vectorbase = VectorBase(settings)
+        self._texts: list[str] = []
+        if data:
+            self.deserialize(data)
+
+    def clear(self) -> None:
+        self._vectorbase.clear()
+        self._texts.clear()
+
+    def serialize(self) -> TextEmbeddingIndexData:
+        return TextEmbeddingIndexData(
+            textItems=self._texts,
+            embeddings=self._vectorbase.serialize(),
+        )
+
+    def deserialize(self, data: TextEmbeddingIndexData | None) -> None:
+        self.clear()
+        if data is not None:
+            self._texts = data.get("textItems", [])
+            self._vectorbase.deserialize(data.get("embeddings"))
+
+    async def size(self) -> int:
+        return len(self._vectorbase)
+
+    async def add_terms(self, texts: list[str]) -> None:
+        await self._vectorbase.add_keys(texts)
+        self._texts.extend(texts)
+
+    async def lookup_term(
+        self, text: str, max_hits: int | None = None, min_score: float | None = None
+    ) -> list[Term]:
+        matches = await self._vectorbase.fuzzy_lookup(
+            text, max_hits=max_hits, min_score=min_score
+        )
+        return self.matches_to_terms(matches)
+
+    async def lookup_terms(
+        self,
+        texts: list[str],
+        max_hits: int | None = None,
+        min_score: float | None = None,
+    ) -> list[list[Term]]:
+        matches = [
+            await self._vectorbase.fuzzy_lookup(
+                text, max_hits=max_hits, min_score=min_score
+            )
+            for text in texts
+        ]
+        return [self.matches_to_terms(m) for m in matches]
+
+    def matches_to_terms(self, matches: list[ScoredInt]) -> list[Term]:
+        return [
+            Term(text=self._texts[match.item], weight=match.score) for match in matches
+        ]