typeagent-py 0.1.0__py3-none-any.whl

typeagent/knowpro/search.py

@@ -0,0 +1,628 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from collections.abc import Callable
+from pydantic.dataclasses import dataclass
+from pydantic import Field, AliasChoices
+from typing import TypeGuard, cast, Annotated
+
+from .collections import MessageAccumulator, SemanticRefAccumulator
+from .field_helpers import CamelCaseField
+from .interfaces import (
+    IConversation,
+    IConversationSecondaryIndexes,
+    KnowledgeType,
+    PropertySearchTerm,
+    ScoredMessageOrdinal,
+    ScoredSemanticRefOrdinal,
+    SearchSelectExpr,
+    SearchTerm,
+    SearchTermGroup,
+    SemanticRef,
+    SemanticRefSearchResult,
+    Term,
+    WhenFilter,
+)
+from .kplib import ConcreteEntity
+from ..storage.memory.messageindex import IMessageTextEmbeddingIndex
+from .searchlib import create_tag_search_term_group
+from .query import (
+    BooleanOp,
+    CompiledSearchTerm,
+    CompiledTermGroup,
+    GetScopeExpr,
+    GetScoredMessagesExpr,
+    GroupByKnowledgeTypeExpr,
+    GroupSearchResultsExpr,
+    IQueryOpExpr,
+    IQuerySemanticRefPredicate,
+    IQueryTextRangeSelector,
+    KnowledgeTypePredicate,
+    MatchMessagesAndExpr,
+    MatchMessagesBooleanExpr,
+    MatchMessagesOrExpr,
+    MatchMessagesOrMaxExpr,
+    MatchPropertySearchTermExpr,
+    MatchSearchTermExpr,
+    MatchTagExpr,
+    MatchTermsAndExpr,
+    MatchTermsBooleanExpr,
+    MatchTermsOrExpr,
+    MatchTermsOrMaxExpr,
+    MatchTopicExpr,
+    MessagesFromKnowledgeExpr,
+    NoOpExpr,
+    QueryEvalContext,
+    RankMessagesBySimilarityExpr,
+    SelectMessagesInCharBudget,
+    SelectTopNExpr,
+    SelectTopNKnowledgeGroupExpr,
+    TextRangeSelector,
+    TextRangesFromMessagesSelector,
+    TextRangesInDateRangeSelector,
+    ThreadSelector,
+    WhereSemanticRefExpr,
+    is_conversation_searchable,
+    match_entity_name_or_type,
+    to_non_required_search_term,
+    to_required_search_term,
+)
+from ..storage.memory.reltermsindex import resolve_related_terms
+
+
+@dataclass
+class SearchQueryExpr:
+    select_expressions: list[SearchSelectExpr] = CamelCaseField(
+        "List of selection expressions for the search"
+    )
+    raw_query: str | None = None
+
+
+@dataclass
+class SearchOptions:
+    max_knowledge_matches: int | None = None
+    exact_match: bool = False
+    max_message_matches: int | None = None
+    # The maximum # of total message characters to select
+    # The query processor will ensure that the cumulative character count of message matches
+    # is less than this number
+    max_chars_in_budget: int | None = None
+    threshold_score: float | None = None
+
+    def __repr__(self):
+        parts = []
+        for key in dir(self):
+            if not key.startswith("_"):
+                value = getattr(self, key)
+                if value is not None:
+                    parts.append(f"{key}={value!r}")
+        return f"{self.__class__.__name__}({', '.join(parts)})"
+
+
+@dataclass
+class ConversationSearchResult:
+    message_matches: list[ScoredMessageOrdinal]
+    knowledge_matches: dict[KnowledgeType, SemanticRefSearchResult]
+    raw_query_text: str | None = None
+
+
+async def search_conversation(
+    conversation: IConversation,
+    search_term_group: SearchTermGroup,
+    when_filter: WhenFilter | None = None,
+    options: SearchOptions | None = None,
+    raw_search_query: str | None = None,
+) -> ConversationSearchResult | None:
+    options = options or SearchOptions()
+    knowledge_matches = await search_conversation_knowledge(
+        conversation, search_term_group, when_filter, options
+    )
+    if knowledge_matches is None:
+        return None
+    # return ConversationSearchResult([], knowledge_matches, raw_search_query)
+    # Future: Combine knowledge and message queries into single query tree.
+    compiler = QueryCompiler(conversation, conversation.secondary_indexes)
+    message_query = await compiler.compile_message_query(
+        knowledge_matches, options, raw_search_query
+    )
+    message_matches: list[ScoredMessageOrdinal] = await run_query(
+        conversation, options, message_query
+    )
+    return ConversationSearchResult(
+        message_matches, knowledge_matches, raw_search_query
+    )
+
+
+async def search_conversation_knowledge(
+    conversation: IConversation,
+    search_term_group: SearchTermGroup,
+    when_filter: WhenFilter | None = None,
+    options: SearchOptions | None = None,
+) -> dict[KnowledgeType, SemanticRefSearchResult] | None:
+    """Search a conversation for knowledge that matches the given search terms and filter."""
+    options = options or SearchOptions()
+    if not is_conversation_searchable(conversation):
+        return None
+    assert (
+        conversation.secondary_indexes is not None
+    ), "Conversation secondary indexes must be initialized before searching"
+    compiler = QueryCompiler(conversation, conversation.secondary_indexes)
+    knowledge_query = await compiler.compile_knowledge_query(
+        search_term_group, when_filter, options
+    )
+    return await run_query(conversation, options, knowledge_query)
+
+
+# TODO: search_conversation_by_text_similarity
+
+
+async def run_search_query(
+    conversation: IConversation,
+    query: SearchQueryExpr,
+    options: SearchOptions | None = None,
+    original_query_text: str | None = None,
+) -> list[ConversationSearchResult]:
+    options = options or SearchOptions()
+    results: list[ConversationSearchResult] = []
+    for expr in query.select_expressions:
+        search_results = await search_conversation(
+            conversation,
+            expr.search_term_group,
+            expr.when,
+            options,
+            original_query_text or query.raw_query,
+        )
+        if search_results is not None:
+            results.append(search_results)
+    return results
+
+
+# TODO: run_search_query_by_text_similarity
+
+
+async def run_query[T](
+    conversation: IConversation,
+    options: SearchOptions | None,
+    query: IQueryOpExpr[T],
+) -> T:
+    secondary_indexes = conversation.secondary_indexes
+    assert (
+        secondary_indexes is not None
+    ), "Conversation secondary indexes must be initialized before running queries"
+    return await query.eval(
+        QueryEvalContext(
+            conversation,
+            secondary_indexes.property_to_semantic_ref_index,
+            secondary_indexes.timestamp_index,
+        )
+    )
+
+
+# NOTE: QueryCompiler instances are stateful, and not thread-safe.
+#       Create a new one for each query.
+class QueryCompiler:
+    def __init__(
+        self,
+        conversation: IConversation,
+        secondary_indexes: IConversationSecondaryIndexes | None,
+        entity_term_match_weight: float = 100.0,
+        default_term_match_weight: float = 10.0,
+        related_is_exact_threshold: float = 0.95,
+    ):
+        self.conversation = conversation
+        self.secondary_indexes = secondary_indexes
+        self.entity_term_match_weight = entity_term_match_weight
+        self.default_term_match_weight = default_term_match_weight
+        self.related_is_exact_threshold = related_is_exact_threshold
+        # All SearchTerms used which compiling the 'select' portion of the query.
+        self.all_search_terms: list[CompiledTermGroup] = []
+        # All search terms used while compiling predicates in the query.
+        self.all_predicate_search_terms: list[CompiledTermGroup] = []
+        self.all_scope_search_terms: list[CompiledTermGroup] = []
+
+    # NOTE: Everything is async because we sometimes use embeddings.
+
+    async def compile_knowledge_query(
+        self,
+        terms: SearchTermGroup,
+        filter: WhenFilter | None = None,
+        options: SearchOptions | None = None,
+    ) -> GroupSearchResultsExpr:
+        query = await self.compile_query(terms, filter, options)
+
+        exact_match = (
+            options.exact_match
+            if options is not None and options.exact_match is not None
+            else False
+        )
+        if not exact_match:
+            await self.resolve_related_terms(self.all_search_terms, True)
+            await self.resolve_related_terms(self.all_predicate_search_terms, False)
+            await self.resolve_related_terms(self.all_scope_search_terms, False)
+
+        return GroupSearchResultsExpr(query)
+
+    async def compile_message_query(
+        self,
+        knowledge: (
+            IQueryOpExpr[dict[KnowledgeType, SemanticRefSearchResult]]
+            | dict[KnowledgeType, SemanticRefSearchResult]
+        ),
+        options: SearchOptions | None = None,
+        raw_query_text: str | None = None,
+    ) -> GetScoredMessagesExpr:
+        query: IQueryOpExpr = MessagesFromKnowledgeExpr(knowledge)
+        if options is not None:
+            query = await self.compile_message_re_rank(
+                query,
+                raw_query_text,
+                options,
+            )
+            if options.max_chars_in_budget and options.max_chars_in_budget > 0:
+                query = SelectMessagesInCharBudget(
+                    query,
+                    options.max_chars_in_budget,
+                )
+
+        return GetScoredMessagesExpr(query)
+
+    # TODO: compile_message_similarity_query
+
+    async def compile_query(
+        self,
+        search_term_group: SearchTermGroup,
+        filter: WhenFilter | None = None,
+        options: SearchOptions | None = None,
+    ) -> IQueryOpExpr[dict[KnowledgeType, SemanticRefAccumulator]]:
+        select_expr = self.compile_select(
+            search_term_group,
+            await self.compile_scope(search_term_group, filter),
+            options,
+        )
+        # Constrain the select with scopes and 'where'.
+        if filter:
+            select_expr = WhereSemanticRefExpr(
+                select_expr,
+                self.compile_where(filter),
+            )
+        # And lastly, select 'TopN' and group knowledge by type.
+        tmp = GroupByKnowledgeTypeExpr(select_expr)
+        return SelectTopNKnowledgeGroupExpr(
+            tmp,
+            (
+                options.max_knowledge_matches
+                if options
+                and hasattr(options, "max_knowledge_matches")
+                and options.max_knowledge_matches
+                else None
+            ),
+        )
+
+    def compile_select(
+        self,
+        term_group: SearchTermGroup,
+        scope_expr: GetScopeExpr | None = None,
+        options: SearchOptions | None = None,
+    ) -> IQueryOpExpr[SemanticRefAccumulator]:
+        search_terms_used, select_expr = self.compile_search_group_terms(
+            term_group, scope_expr
+        )
+        self.all_search_terms.extend(search_terms_used)
+        return select_expr
+
+    def compile_search_group_terms(
+        self,
+        search_group: SearchTermGroup,
+        scope_expr: GetScopeExpr | None = None,
+    ) -> tuple[list[CompiledTermGroup], IQueryOpExpr[SemanticRefAccumulator]]:
+        return self.compile_search_group(
+            search_group,
+            create_match_terms_boolean_expr,
+            scope_expr,
+        )
+
+    def compile_search_group_messages(
+        self,
+        search_group: SearchTermGroup,
+    ) -> tuple[list[CompiledTermGroup], IQueryOpExpr[MessageAccumulator]]:
+        return self.compile_search_group(
+            search_group, create_match_messages_boolean_expr
+        )
+
+    def compile_search_group(
+        self,
+        search_group: SearchTermGroup,
+        create_op: Callable[
+            [list[IQueryOpExpr], BooleanOp, GetScopeExpr | None],
+            IQueryOpExpr[SemanticRefAccumulator | MessageAccumulator],
+        ],
+        scope_expr: GetScopeExpr | None = None,
+    ) -> tuple[list[CompiledTermGroup], IQueryOpExpr]:
+        t0_terms: list[CompiledSearchTerm] = []
+        compiled_terms: list[CompiledTermGroup] = [
+            CompiledTermGroup(boolean_op=search_group.boolean_op, terms=t0_terms)
+        ]
+        term_expressions: list[IQueryOpExpr[SemanticRefAccumulator | None]] = []
+        for term in search_group.terms:
+            if isinstance(term, PropertySearchTerm):
+                term_expressions.append(self.compile_property_term(term))
+                if not isinstance(term.property_name, str):
+                    t0_terms.append(to_required_search_term(term.property_name))
+                t0_terms.append(to_required_search_term(term.property_value))
+            elif isinstance(term, SearchTermGroup):
+                nested_terms, group_expr = self.compile_search_group(term, create_op)
+                compiled_terms.extend(nested_terms)
+                term_expressions.append(group_expr)
+            else:
+                term_expressions.append(self.compile_search_term(term))
+                t0_terms.append(to_non_required_search_term(term))
+        bool_expr = create_op(term_expressions, search_group.boolean_op, scope_expr)
+        return (compiled_terms, bool_expr)
+
+    def compile_search_term(
+        self,
+        term: SearchTerm,
+    ) -> IQueryOpExpr[SemanticRefAccumulator | None]:
+        boost_weight = self.entity_term_match_weight / self.default_term_match_weight
+        return MatchSearchTermExpr(
+            term,
+            lambda term, sr, scored: self.boost_entities(
+                term, sr, scored, boost_weight
+            ),
+        )
+
+    def compile_property_term(
+        self,
+        term: PropertySearchTerm,
+    ) -> IQueryOpExpr[SemanticRefAccumulator | None]:
+        match term.property_name:
+            case "tag":
+                return MatchTagExpr(term.property_value)
+            case "topic":
+                return MatchTopicExpr(term.property_value)
+            case _:
+                if term.property_name in ("name", "type"):
+                    tpvt = term.property_value.term
+                    if tpvt.weight is None:
+                        tpvt.weight = self.entity_term_match_weight
+                return MatchPropertySearchTermExpr(term)
+
+    async def compile_scope(
+        self,
+        term_group: SearchTermGroup | None = None,
+        filter: WhenFilter | None = None,
+    ) -> GetScopeExpr | None:
+        scope_selectors: list[IQueryTextRangeSelector] = []
+
+        # First, use any provided date ranges to select scope
+        if filter and filter.date_range:
+            scope_selectors.append(TextRangesInDateRangeSelector(filter.date_range))
+
+        # Apply 'OUTER' scope
+        # If specific scoping terms were provided
+        if filter and filter.scope_defining_terms is not None:
+            self.add_terms_scope_selector(filter.scope_defining_terms, scope_selectors)
+        elif term_group is not None:
+            # Treat any actions as inherently scope selecting
+            action_terms_group = self.get_action_terms_from_search_group(term_group)
+            if action_terms_group is not None:
+                self.add_terms_scope_selector(action_terms_group, scope_selectors)
+
+        # Include any ranges directly provided by the caller
+        if filter and filter.text_ranges_in_scope:
+            scope_selectors.append(TextRangeSelector(filter.text_ranges_in_scope))
+
+        # Tags...
+        if filter and filter.tags:
+            self.add_terms_scope_selector(
+                create_tag_search_term_group(filter.tags), scope_selectors
+            )
+
+        # If a thread index is available...
+        threads = None
+        if self.secondary_indexes:
+            threads = self.secondary_indexes.threads
+        if filter and filter.thread_description and threads:
+            threads_in_scope = await threads.lookup_thread(filter.thread_description)
+            if threads_in_scope:
+                scope_selectors.append(
+                    ThreadSelector(
+                        [threads.threads[t.thread_ordinal] for t in threads_in_scope]
+                    )
+                )
+
+        return GetScopeExpr(scope_selectors) if scope_selectors else None
+
+    def add_terms_scope_selector(
+        self,
+        term_group: SearchTermGroup,
+        scope_selectors: list[IQueryTextRangeSelector],
+    ) -> None:
+        if term_group.terms:
+            search_terms_used, select_expr = self.compile_search_group_messages(
+                term_group
+            )
+            scope_selectors.append(TextRangesFromMessagesSelector(select_expr))
+            self.all_scope_search_terms.extend(search_terms_used)
+
+    def compile_where(self, filter: WhenFilter) -> list[IQuerySemanticRefPredicate]:
+        predicates: list[IQuerySemanticRefPredicate] = []
+        if filter.knowledge_type:
+            predicates.append(KnowledgeTypePredicate(filter.knowledge_type))
+        return predicates
+
+    async def compile_message_re_rank(
+        self,
+        src_expr: IQueryOpExpr,
+        raw_query_text: str | None = None,
+        options: SearchOptions | None = None,
+    ) -> IQueryOpExpr:
+        message_index = (
+            self.conversation.secondary_indexes.message_index
+            if self.conversation.secondary_indexes
+            else None
+        )
+        if (
+            raw_query_text is not None
+            and isinstance(message_index, IMessageTextEmbeddingIndex)
+            and not await message_index.is_empty()
+        ):
+            embedding = await message_index.generate_embedding(raw_query_text)
+            return RankMessagesBySimilarityExpr(
+                src_expr,
+                embedding,
+                options.max_message_matches if options else None,
+                options.threshold_score if options else None,
+            )
+        elif (
+            options
+            and options.max_message_matches is not None
+            and options.max_message_matches > 0
+        ):
+            return SelectTopNExpr(src_expr, options.max_message_matches)
+        else:
+            return NoOpExpr(src_expr)
+
+    # TODO: compile_message_similarity
+
+    def get_action_terms_from_search_group(
+        self,
+        search_group: SearchTermGroup,
+    ) -> SearchTermGroup | None:
+        action_group: SearchTermGroup | None = None
+        for term in search_group.terms:
+            if isinstance(term, PropertySearchTerm) and is_action_property_term(term):
+                action_group = action_group or SearchTermGroup(boolean_op="and")
+                action_group.terms.append(term)
+        return action_group
+
+    async def resolve_related_terms(
+        self,
+        compiled_terms: list[CompiledTermGroup],
+        dedupe: bool,
+        filter: WhenFilter | None = None,
+    ) -> None:
+        if not compiled_terms:
+            return
+        for ct in compiled_terms:
+            self.validate_and_prepare_search_terms(ct.terms)
+        if (
+            self.secondary_indexes is not None
+            and self.secondary_indexes.term_to_related_terms_index is not None
+        ):
+            await resolve_related_terms(
+                self.secondary_indexes.term_to_related_terms_index,
+                compiled_terms,
+                dedupe,
+            )
+            # Ensure that the resolved terms are valid etc.
+            for ct in compiled_terms:
+                self.validate_and_prepare_search_terms(ct.terms)
+
+    def validate_and_prepare_search_terms(
+        self, terms: list[CompiledSearchTerm]
+    ) -> None:
+        for term in terms:
+            self.validate_and_prepare_search_term(term)
+
+    def validate_and_prepare_search_term(self, search_term: CompiledSearchTerm) -> bool:
+        if not self.validate_and_prepare_term(search_term.term):
+            return False
+        # Matching the term - exact match - counts for more than matching related terms
+        # Therefore, we boost any matches where the term matches directly...
+        if search_term.term.weight is None:
+            search_term.term.weight = self.default_term_match_weight
+        if search_term.related_terms is not None:
+            for related_term in search_term.related_terms:
+                if not self.validate_and_prepare_term(related_term):
+                    return False
+                # If related term is *really* similar to the main term, score it the same
+                if (
+                    related_term.weight is not None
+                    and related_term.weight >= self.related_is_exact_threshold
+                ):
+                    related_term.weight = self.default_term_match_weight
+        return True
+
+    # Currently, just changes the case of a term
+    #  But here, we may do other things like:
+    # - Check for noise terms
+    # - Do additional rewriting
+    # - Additional checks that *reject* certain search terms
+    # Return false if the term should be rejected
+    def validate_and_prepare_term(self, term: Term | None) -> bool:
+        if term:
+            term.text = term.text.lower()
+        return True
+
+    def boost_entities(
+        self,
+        search_term: SearchTerm,
+        sr: SemanticRef,
+        scored_ref: ScoredSemanticRefOrdinal,
+        boost_weight: float,
+    ) -> ScoredSemanticRefOrdinal:
+        if sr.knowledge.knowledge_type == "entity" and match_entity_name_or_type(
+            search_term, cast(ConcreteEntity, sr.knowledge)
+        ):
+            return ScoredSemanticRefOrdinal(
+                scored_ref.semantic_ref_ordinal,
+                scored_ref.score * boost_weight,
+            )
+        else:
+            return scored_ref
+
+
+def has_conversation_results(results: list[ConversationSearchResult]) -> bool:
+    return any(r.knowledge_matches or r.message_matches for r in results)
+
+
+def has_conversation_result(result: ConversationSearchResult) -> bool:
+    return bool(result.knowledge_matches or result.message_matches)
+
+
+# TODO: Move to compilelib.py
+def create_match_terms_boolean_expr(
+    term_expressions: list[IQueryOpExpr[SemanticRefAccumulator | None]],
+    boolean_op: BooleanOp,
+    scope_expr: GetScopeExpr | None = None,
+) -> MatchTermsBooleanExpr:
+    match boolean_op:
+        case "and":
+            return MatchTermsAndExpr(term_expressions, scope_expr)
+        case "or":
+            return MatchTermsOrExpr(term_expressions, scope_expr)
+        case "or_max":
+            return MatchTermsOrMaxExpr(term_expressions, scope_expr)
+        case _:
+            raise ValueError(f"Unknown boolean op: {boolean_op}")
+
+
+# TODO: Move to compilelib.py
+def create_match_messages_boolean_expr(
+    term_expressions: list[
+        IQueryOpExpr[SemanticRefAccumulator | MessageAccumulator | None]
+    ],
+    boolean_op: BooleanOp,
+    scope_expr: GetScopeExpr | None = None,
+) -> MatchMessagesBooleanExpr:
+    match boolean_op:
+        case "and":
+            return MatchMessagesAndExpr(term_expressions)
+        case "or":
+            return MatchMessagesOrExpr(term_expressions)
+        case "or_max":
+            return MatchMessagesOrMaxExpr(term_expressions)
+        case _:
+            raise ValueError(f"Unknown boolean op: {boolean_op}")
+
+
+# TODO: Move to compilelib.py
+# TODO: Just call isinstance!
+def is_property_term(term: SearchTerm) -> TypeGuard[PropertySearchTerm]:
+    return isinstance(term, PropertySearchTerm)
+
+
+# TODO: Move to compilelib.py
+def is_action_property_term(term: PropertySearchTerm) -> bool:
+    return term.property_name in ("subject", "verb", "object", "indirectObject")