typeagent-py 0.1.0__py3-none-any.whl

typeagent/knowpro/query.py

@@ -0,0 +1,1128 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from re import search
+from typing import Literal, Protocol, cast
+
+from ..aitools.embeddings import NormalizedEmbedding
+
+from .collections import (
+    Match,
+    MatchAccumulator,
+    MessageAccumulator,
+    PropertyTermSet,
+    SemanticRefAccumulator,
+    TermSet,
+    TextRangeCollection,
+    TextRangesInScope,
+)
+from .common import is_search_term_wildcard
+from .interfaces import (
+    Datetime,
+    DateRange,
+    IConversation,
+    IMessage,
+    IMessageCollection,
+    IPropertyToSemanticRefIndex,
+    ISemanticRefCollection,
+    ITermToSemanticRefIndex,
+    ITimestampToTextRangeIndex,
+    KnowledgeType,
+    MessageOrdinal,
+    PropertySearchTerm,
+    ScoredMessageOrdinal,
+    ScoredSemanticRefOrdinal,
+    SearchTerm,
+    SearchTermGroup,
+    SemanticRef,
+    SemanticRefOrdinal,
+    SemanticRefSearchResult,
+    Term,
+    TextLocation,
+    TextRange,
+    Thread,
+)
+from .kplib import ConcreteEntity
+from ..storage.memory.messageindex import IMessageTextEmbeddingIndex
+from ..storage.memory.propindex import PropertyNames, lookup_property_in_property_index
+from .searchlib import create_property_search_term, create_tag_search_term_group
+
+
+# TODO: Move to compilelib.py
+type BooleanOp = Literal["and", "or", "or_max"]
+
+
+# TODO: Move to compilelib.py
+@dataclass
+class CompiledSearchTerm(SearchTerm):
+    related_terms_required: bool = False
+
+
+# TODO: Move to compilelib.py
+def to_required_search_term(term: SearchTerm) -> CompiledSearchTerm:
+    # NOTE: We must cast since the output must alias the input.
+    # If not, assignments to related_terms will be lost.
+    cst = cast(CompiledSearchTerm, term)
+    cst.related_terms_required = True
+    return cst
+
+
+# TODO: Move to compilelib.py
+def to_non_required_search_term(term: SearchTerm) -> CompiledSearchTerm:
+    # NOTE: We must cast since the output must alias the input.
+    # If not, assignments to related_terms will be lost.
+    cst = cast(CompiledSearchTerm, term)
+    cst.related_terms_required = False
+    return cst
+
+
+# TODO: Move to compilelib.py
+@dataclass
+class CompiledTermGroup:
+    boolean_op: BooleanOp
+    terms: list[CompiledSearchTerm]
+
+
+def is_conversation_searchable(conversation: IConversation) -> bool:
+    """Determine if a conversation is searchable.
+
+    A conversation is searchable if it has a semantic reference index
+    and semantic references initialized.
+    """
+    # TODO: also require secondary indices, once we have removed non-index based retrieval to test.
+    return (
+        conversation.semantic_ref_index is not None
+        and conversation.semantic_refs is not None
+    )
+
+
+async def get_text_range_for_date_range(
+    conversation: IConversation,
+    date_range: DateRange,
+) -> TextRange | None:
+    messages = conversation.messages
+    range_start_ordinal: MessageOrdinal = -1
+    range_end_ordinal = range_start_ordinal
+    async for message in messages:
+        if Datetime.fromisoformat(message.timestamp) in date_range:
+            if range_start_ordinal < 0:
+                range_start_ordinal = message.ordinal
+            range_end_ordinal = message.ordinal
+        else:
+            if range_start_ordinal >= 0:
+                # We have a range, so break.
+                break
+    if range_start_ordinal >= 0:
+        return TextRange(
+            start=TextLocation(range_start_ordinal),
+            end=TextLocation(range_end_ordinal + 1),
+        )
+    return None
+
+
+def get_matching_term_for_text(search_term: SearchTerm, text: str) -> Term | None:
+    # Do case-INSENSITIVE comparisons, since stored entities may have different case.
+    if text.lower() == search_term.term.text.lower():
+        return search_term.term
+    if search_term.related_terms:
+        for related_term in search_term.related_terms:
+            if text.lower() == related_term.text.lower():
+                return related_term
+    return None
+
+
+def match_search_term_to_text(search_term: SearchTerm, text: str | None) -> bool:
+    if text:
+        return get_matching_term_for_text(search_term, text) is not None
+    return False
+
+
+def match_search_term_to_one_of_text(
+    search_term: SearchTerm, texts: list[str] | None
+) -> bool:
+    if texts:
+        for text in texts:
+            if match_search_term_to_text(search_term, text):
+                return True
+    return False
+
+
+# TODO: match_search_term_to_entity
+# TODO: match_property_search_term_to_entity
+# TODO: match_concrete_entity
+
+
+def match_entity_name_or_type(
+    property_value: SearchTerm,
+    entity: ConcreteEntity,
+) -> bool:
+    return match_search_term_to_text(
+        property_value, entity.name
+    ) or match_search_term_to_one_of_text(property_value, entity.type)
+
+
+# TODO: match_property_name_to_facet_name
+# TODO: match_property_name_to_facet_value
+# TODO: match_property_search_term_to_action
+# TODO: match_property_search_term_to_tag
+# TODO: match_property_search_term_to_semantic_ref
+
+
+async def lookup_term_filtered(
+    semantic_ref_index: ITermToSemanticRefIndex,
+    term: Term,
+    semantic_refs: ISemanticRefCollection,
+    filter: Callable[[SemanticRef, ScoredSemanticRefOrdinal], bool],
+) -> list[ScoredSemanticRefOrdinal] | None:
+    """Look up a term in the semantic reference index and filter the results."""
+    scored_refs = await semantic_ref_index.lookup_term(term.text)
+    if scored_refs:
+        filtered = []
+        for sr in scored_refs:
+            semantic_ref = await semantic_refs.get_item(sr.semantic_ref_ordinal)
+            if filter(semantic_ref, sr):
+                filtered.append(sr)
+        return filtered
+    return None
+
+
+async def lookup_term(
+    semantic_ref_index: ITermToSemanticRefIndex,
+    term: Term,
+    semantic_refs: ISemanticRefCollection,
+    ranges_in_scope: TextRangesInScope | None = None,
+    knowledge_type: KnowledgeType | None = None,
+) -> list[ScoredSemanticRefOrdinal] | None:
+    """Look up a term in the semantic reference index, optionally filtering by ranges in scope."""
+    if ranges_in_scope is not None:
+        # If ranges_in_scope has no actual text ranges, lookups can't possibly match.
+        return await lookup_term_filtered(
+            semantic_ref_index,
+            term,
+            semantic_refs,
+            lambda sr, _: (
+                not knowledge_type or sr.knowledge.knowledge_type == knowledge_type
+            )
+            and ranges_in_scope.is_range_in_scope(sr.range),
+        )
+    return await semantic_ref_index.lookup_term(term.text)
+
+
+# TODO: lookup_property
+# TODO: lookup_knowledge_type
+
+
+@dataclass
+class QueryEvalContext[TMessage: IMessage, TIndex: ITermToSemanticRefIndex]:
+    """Context for evaluating a query within a conversation.
+
+    This class provides the necessary context for query evaluation, including
+    the conversation being queried, optional indexes for properties and timestamps,
+    and structures for tracking matched terms and text ranges in scope.
+    """
+
+    # TODO: Make property and timestamp indexes NON-OPTIONAL
+    # TODO: Move non-index based code to test
+
+    conversation: IConversation[TMessage, TIndex]
+    # If a property secondary index is available, the query processor will use it
+    property_index: IPropertyToSemanticRefIndex | None = None
+    # If a timestamp secondary index is available, the query processor will use it
+    timestamp_index: ITimestampToTextRangeIndex | None = None
+    matched_terms: TermSet = field(init=False, default_factory=TermSet)
+    matched_property_terms: PropertyTermSet = field(
+        init=False, default_factory=PropertyTermSet
+    )
+    text_ranges_in_scope: TextRangesInScope | None = field(
+        init=False, default_factory=TextRangesInScope
+    )
+
+    def __post_init__(self):
+        if not is_conversation_searchable(self.conversation):
+            raise ValueError(
+                f"{self.conversation.name_tag} "
+                + "is not initialized and cannot be searched."
+            )
+
+    @property
+    def semantic_ref_index(self) -> ITermToSemanticRefIndex:
+        assert self.conversation.semantic_ref_index is not None
+        return self.conversation.semantic_ref_index
+
+    @property
+    def semantic_refs(self) -> ISemanticRefCollection:
+        assert self.conversation.semantic_refs is not None
+        return self.conversation.semantic_refs
+
+    @property
+    def messages(self) -> IMessageCollection:
+        return self.conversation.messages
+
+    async def get_semantic_ref(
+        self, semantic_ref_ordinal: SemanticRefOrdinal
+    ) -> SemanticRef:
+        """Retrieve a semantic reference by its ordinal."""
+        assert self.conversation.semantic_refs is not None
+        return await self.conversation.semantic_refs.get_item(semantic_ref_ordinal)
+
+    async def get_message_for_ref(self, semantic_ref: SemanticRef) -> TMessage:
+        """Retrieve the message associated with a semantic reference."""
+        message_ordinal = semantic_ref.range.start.message_ordinal
+        return await self.conversation.messages.get_item(message_ordinal)
+
+    async def get_message(self, message_ordinal: MessageOrdinal) -> TMessage:
+        """Retrieve a message by its ordinal."""
+        return await self.messages.get_item(message_ordinal)
+
+    def clear_matched_terms(self) -> None:
+        """Clear all matched terms and property terms."""
+        self.matched_terms.clear()
+        self.matched_property_terms.clear()
+
+
+async def lookup_knowledge_type(
+    semantic_refs: ISemanticRefCollection, knowledge_type: KnowledgeType
+) -> list[ScoredSemanticRefOrdinal]:
+    return [
+        ScoredSemanticRefOrdinal(sr.semantic_ref_ordinal, 1.0)
+        async for sr in semantic_refs
+        if sr.knowledge.knowledge_type == knowledge_type
+    ]
+
+
+class IQueryOpExpr[T](Protocol):
+    """Protocol for query operation expressions that can be evaluated in a context."""
+
+    async def eval(self, context: QueryEvalContext) -> T: ...
+
+
+class QueryOpExpr[T](IQueryOpExpr[T]):
+    """Base class for query operation expressions."""
+
+
+@dataclass
+class SelectTopNExpr[T: MatchAccumulator](QueryOpExpr[T]):
+    """Expression for selecting the top N matches from a query."""
+
+    source_expr: IQueryOpExpr[T]
+    max_matches: int | None = None
+    min_hit_count: int | None = None
+
+    async def eval(self, context: QueryEvalContext) -> T:
+        """Evaluate the expression and return the top N matches."""
+        matches = await self.source_expr.eval(context)
+        matches.select_top_n_scoring(self.max_matches, self.min_hit_count)
+        return matches
+
+
+# Abstract base class.
+class MatchTermsBooleanExpr(QueryOpExpr[SemanticRefAccumulator]):
+    """Expression for matching terms in a boolean query.
+
+    Subclasses implement 'OR', 'OR MAX' and 'AND' logic.
+    """
+
+    get_scope_expr: "GetScopeExpr | None" = None
+
+    async def begin_match(self, context: QueryEvalContext) -> None:
+        """Prepare for matching terms in the context by resetting some things."""
+        if self.get_scope_expr is not None:
+            context.text_ranges_in_scope = await self.get_scope_expr.eval(context)
+        context.clear_matched_terms()
+
+
+@dataclass
+class MatchTermsOrExpr(MatchTermsBooleanExpr):
+    """Expression for matching terms with an OR condition."""
+
+    term_expressions: list[IQueryOpExpr[SemanticRefAccumulator | None]] = field(
+        default_factory=list
+    )
+    get_scope_expr: "GetScopeExpr | None" = None
+
+    async def eval(self, context: QueryEvalContext) -> SemanticRefAccumulator:
+        await self.begin_match(context)
+        all_matches: SemanticRefAccumulator | None = None
+        for match_expr in self.term_expressions:
+            term_matches = await match_expr.eval(context)
+            if term_matches:
+                if all_matches is None:
+                    all_matches = term_matches
+                else:
+                    all_matches.add_union(term_matches)
+        if all_matches is not None:
+            all_matches.calculate_total_score()
+        return all_matches or SemanticRefAccumulator()
+
+
+@dataclass
+class MatchTermsOrMaxExpr(MatchTermsOrExpr):
+    """OR-MAX returns the union if there are no common matches, else the maximum scoring match."""
+
+    term_expressions: list[IQueryOpExpr[SemanticRefAccumulator | None]] = field(
+        default_factory=list
+    )
+    get_scope_expr: "GetScopeExpr | None" = None
+
+    async def eval(self, context: QueryEvalContext) -> SemanticRefAccumulator:
+        matches = await super().eval(context)
+        max_hit_count = matches.get_max_hit_count()
+        if max_hit_count > 1:
+            matches.select_with_hit_count(max_hit_count)
+        return matches
+
+
+@dataclass
+class MatchTermsAndExpr(MatchTermsBooleanExpr):
+    term_expressions: list[IQueryOpExpr[SemanticRefAccumulator | None]] = field(
+        default_factory=list
+    )
+    get_scope_expr: "GetScopeExpr | None" = None
+
+    async def eval(self, context: QueryEvalContext) -> SemanticRefAccumulator:
+        await self.begin_match(context)
+        all_matches: SemanticRefAccumulator | None = None
+        for match_expr in self.term_expressions:
+            term_matches = await match_expr.eval(context)
+            if not term_matches:
+                if all_matches is not None:
+                    all_matches.clear_matches()
+                break
+            if all_matches is None:
+                all_matches = term_matches
+            else:
+                all_matches = all_matches.intersect(term_matches)
+        if all_matches is not None:
+            all_matches.calculate_total_score()
+            all_matches.select_with_hit_count(len(self.term_expressions))
+        else:
+            all_matches = SemanticRefAccumulator()
+        return all_matches
+
+
+class MatchTermExpr(QueryOpExpr[SemanticRefAccumulator | None], ABC):
+    """Expression for matching terms in a query.
+
+    Subclasses need to define accumulate_matches(), which must add
+    matches to its SemanticRefAccumulator argument.
+    """
+
+    async def eval(self, context: QueryEvalContext) -> SemanticRefAccumulator | None:
+        matches = SemanticRefAccumulator()
+        await self.accumulate_matches(context, matches)
+        if len(matches) > 0:
+            return matches
+        return None
+
+    @abstractmethod
+    async def accumulate_matches(
+        self, context: QueryEvalContext, matches: SemanticRefAccumulator
+    ) -> None: ...
+
+
+type ScoreBoosterType = Callable[
+    [SearchTerm, SemanticRef, ScoredSemanticRefOrdinal],
+    ScoredSemanticRefOrdinal,
+]
+
+
+@dataclass
+class MatchSearchTermExpr(MatchTermExpr):
+    search_term: SearchTerm
+    score_booster: ScoreBoosterType | None = None
+
+    async def accumulate_matches(
+        self, context: QueryEvalContext, matches: SemanticRefAccumulator
+    ) -> None:
+        """Accumulate matches for the search term and its related terms."""
+        # Match the search term
+        await self.accumulate_matches_for_term(context, matches, self.search_term.term)
+
+        # And any related terms
+        if self.search_term.related_terms is not None:
+            for related_term in self.search_term.related_terms:
+                await self.accumulate_matches_for_term(
+                    context, matches, self.search_term.term, related_term
+                )
+
+    async def lookup_term(
+        self, context: QueryEvalContext, term: Term
+    ) -> list[ScoredSemanticRefOrdinal] | None:
+        """Look up a term in the semantic reference index."""
+        matches = await lookup_term(
+            context.semantic_ref_index,
+            term,
+            context.semantic_refs,
+            context.text_ranges_in_scope,
+        )
+        if matches and self.score_booster:
+            for i in range(len(matches)):
+                matches[i] = self.score_booster(
+                    self.search_term,
+                    await context.get_semantic_ref(matches[i].semantic_ref_ordinal),
+                    matches[i],
+                )
+        return matches
+
+    async def accumulate_matches_for_term(
+        self,
+        context: QueryEvalContext,
+        matches: SemanticRefAccumulator,
+        term: Term,
+        related_term: Term | None = None,
+    ) -> None:
+        """Accumulate matches for a term or a related term."""
+        if related_term is None:
+            if term not in context.matched_terms:
+                semantic_refs = await self.lookup_term(context, term)
+                matches.add_term_matches(term, semantic_refs, True)
+                context.matched_terms.add(term)
+        else:
+            if related_term not in context.matched_terms:
+                # If this related term had not already matched as a related term for some other term
+                # Minimize over counting
+                semantic_refs = await self.lookup_term(context, related_term)
+                # This will only consider semantic refs that have not already matched this expression.
+                # In other words, if a semantic ref already matched due to the term 'novel',
+                # don't also match it because it matched the related term 'book'
+                matches.add_term_matches_if_new(
+                    term, semantic_refs, False, related_term.weight
+                )
+                context.matched_terms.add(related_term)
+
+
+@dataclass
+class MatchPropertySearchTermExpr(MatchTermExpr):
+    property_search_term: PropertySearchTerm
+
+    async def accumulate_matches(
+        self, context: QueryEvalContext, matches: SemanticRefAccumulator
+    ) -> None:
+        if isinstance(self.property_search_term.property_name, str):
+            await self.accumulate_matches_for_property(
+                context,
+                self.property_search_term.property_name,
+                self.property_search_term.property_value,
+                matches,
+            )
+        else:
+            await self.accumulate_matches_for_facets(
+                context,
+                self.property_search_term.property_name,
+                self.property_search_term.property_value,
+                matches,
+            )
+
+    async def accumulate_matches_for_facets(
+        self,
+        context: QueryEvalContext,
+        property_name: SearchTerm,
+        property_value: SearchTerm,
+        matches: SemanticRefAccumulator,
+    ):
+        await self.accumulate_matches_for_property(
+            context,
+            PropertyNames.FacetName.value,
+            property_name,
+            matches,
+        )
+        if not is_search_term_wildcard(property_value):
+            await self.accumulate_matches_for_property(
+                context,
+                PropertyNames.FacetValue.value,
+                property_value,
+                matches,
+            )
+
+    async def accumulate_matches_for_property(
+        self,
+        context: QueryEvalContext,
+        property_name: str,
+        property_value: SearchTerm,
+        matches: SemanticRefAccumulator,
+    ):
+        await self.accumulate_matches_for_property_value(
+            context,
+            matches,
+            property_name,
+            property_value.term,
+        )
+        if property_value.related_terms:
+            for related_property_value in property_value.related_terms:
+                await self.accumulate_matches_for_property_value(
+                    context,
+                    matches,
+                    property_name,
+                    property_value.term,
+                    related_property_value,
+                )
+
+    async def accumulate_matches_for_property_value(
+        self,
+        context: QueryEvalContext,
+        matches: SemanticRefAccumulator,
+        property_name: str,
+        property_value: Term,
+        related_prop_val: Term | None = None,
+    ) -> None:
+        if related_prop_val is None:
+            if not context.matched_property_terms.has(property_name, property_value):
+                semantic_refs = await self.lookup_property(
+                    context,
+                    property_name,
+                    property_value.text,
+                )
+                if semantic_refs:
+                    matches.add_term_matches(property_value, semantic_refs, True)
+                    context.matched_property_terms.add(property_name, property_value)
+        else:
+            # To prevent over-counting, ensure this related_prop_val was not already used to match terms earlier
+            if not context.matched_property_terms.has(property_name, related_prop_val):
+                semantic_refs = await self.lookup_property(
+                    context,
+                    property_name,
+                    related_prop_val.text,
+                )
+                if semantic_refs:
+                    # This will only consider semantic refs that were not already matched by this expression.
+                    # In other words, if a semantic ref already matched due to the term 'novel',
+                    # don't also match it because it matched the related term 'book'
+                    matches.add_term_matches_if_new(
+                        property_value,
+                        semantic_refs,
+                        False,
+                        related_prop_val.weight,
+                    )
+                    context.matched_property_terms.add(property_name, related_prop_val)
+
+    async def lookup_property(
+        self,
+        context: QueryEvalContext,
+        property_name: str,
+        property_value: str,
+    ) -> list[ScoredSemanticRefOrdinal] | None:
+        if context.property_index is not None:
+            return await lookup_property_in_property_index(
+                context.property_index,
+                property_name,
+                property_value,
+                context.semantic_refs,
+                context.text_ranges_in_scope,
+            )
+
+
+class MatchTagExpr(MatchSearchTermExpr):
+    def __init__(self, tag_term: SearchTerm):
+        self.tag_term = tag_term
+        super().__init__(tag_term)
+
+    async def lookup_term(
+        self, context: QueryEvalContext, term: Term
+    ) -> list[ScoredSemanticRefOrdinal] | None:
+        if self.tag_term.term.text == "*":
+            return await lookup_knowledge_type(context.semantic_refs, "tag")
+        else:
+            return await lookup_term(
+                context.semantic_ref_index,
+                term,
+                context.semantic_refs,
+                context.text_ranges_in_scope,
+                "tag",
+            )
+
+
+class MatchTopicExpr(MatchSearchTermExpr):
+    def __init__(self, topic: SearchTerm):
+        self.topic = topic
+        super().__init__(topic)
+
+    async def lookup_term(
+        self, context: QueryEvalContext, term: Term
+    ) -> list[ScoredSemanticRefOrdinal] | None:
+        if self.topic.term.text == "*":
+            return await lookup_knowledge_type(context.semantic_refs, "topic")
+        else:
+            return await lookup_term(
+                context.semantic_ref_index,
+                term,
+                context.semantic_refs,
+                context.text_ranges_in_scope,
+                "topic",
+            )
+
+
+@dataclass
+class GroupByKnowledgeTypeExpr(
+    QueryOpExpr[dict[KnowledgeType, SemanticRefAccumulator]]
+):
+    matches: IQueryOpExpr[SemanticRefAccumulator]
+
+    async def eval(
+        self, context: QueryEvalContext
+    ) -> dict[KnowledgeType, SemanticRefAccumulator]:
+        semantic_ref_matches = await self.matches.eval(context)
+        return await semantic_ref_matches.group_matches_by_type(context.semantic_refs)
+
+
+@dataclass
+class SelectTopNKnowledgeGroupExpr(
+    QueryOpExpr[dict[KnowledgeType, SemanticRefAccumulator]]
+):
+    source_expr: IQueryOpExpr[dict[KnowledgeType, SemanticRefAccumulator]]
+    max_matches: int | None = None
+    min_hit_count: int | None = None
+
+    async def eval(
+        self, context: QueryEvalContext
+    ) -> dict[KnowledgeType, SemanticRefAccumulator]:
+        groups_accumulators = await self.source_expr.eval(context)
+        for accumulator in groups_accumulators.values():
+            accumulator.select_top_n_scoring(self.max_matches, self.min_hit_count)
+        return groups_accumulators
+
+
+@dataclass
+class GroupSearchResultsExpr(QueryOpExpr[dict[KnowledgeType, SemanticRefSearchResult]]):
+    src_expr: IQueryOpExpr[dict[KnowledgeType, SemanticRefAccumulator]]
+
+    async def eval(
+        self, context: QueryEvalContext
+    ) -> dict[KnowledgeType, SemanticRefSearchResult]:
+        return to_grouped_search_results(await self.src_expr.eval(context))
+
+
+@dataclass
+class WhereSemanticRefExpr(QueryOpExpr[SemanticRefAccumulator]):
+    source_expr: IQueryOpExpr[SemanticRefAccumulator]
+    predicates: list["IQuerySemanticRefPredicate"]
+
+    async def eval(self, context: QueryEvalContext) -> SemanticRefAccumulator:
+        accumulator = await self.source_expr.eval(context)
+        filtered = SemanticRefAccumulator(accumulator.search_term_matches)
+
+        # Filter matches asynchronously
+        filtered_matches = []
+        for match in accumulator.get_matches():
+            if await self._eval_predicates(context, self.predicates, match):
+                filtered_matches.append(match)
+
+        filtered.set_matches(filtered_matches)
+        return filtered
+
+    async def _eval_predicates(
+        self,
+        context: QueryEvalContext,
+        predicates: list["IQuerySemanticRefPredicate"],
+        match: Match[SemanticRefOrdinal],
+    ) -> bool:
+        for predicate in predicates:
+            semantic_ref = await context.get_semantic_ref(match.value)
+            if not await predicate.eval(context, semantic_ref):
+                return False
+
+        return True
+
+
+class IQuerySemanticRefPredicate(Protocol):
+    async def eval(
+        self, context: QueryEvalContext, semantic_ref: SemanticRef
+    ) -> bool: ...
+
+
+# TODO: match_predicates
+# TODO: knowledge_type_predicate
+# TODO: property_match_predicate
+
+
+# NOTE: GetScopeExpr is moved after TextRangeSelector to avoid circular references.
+
+
+class IQueryTextRangeSelector(Protocol):
+    """Protocol for a selector that can evaluate to a text range."""
+
+    async def eval(
+        self,
+        context: QueryEvalContext,
+        semantic_refs: SemanticRefAccumulator | None = None,
+    ) -> TextRangeCollection | None:
+        """Evaluate the selector and return the text range."""
+        ...
+
+
+class TextRangeSelector(IQueryTextRangeSelector):
+    """A selector that evaluates to a pre-computed TextRangeCollection."""
+
+    text_ranges_in_scope: TextRangeCollection
+
+    def __init__(self, ranges_in_scope: list[TextRange]) -> None:
+        self.text_ranges_in_scope = TextRangeCollection(ranges_in_scope, True)
+
+    async def eval(
+        self,
+        context: QueryEvalContext,
+        semantic_refs: SemanticRefAccumulator | None = None,
+    ) -> TextRangeCollection | None:
+        return self.text_ranges_in_scope
+
+
+@dataclass
+class GetScopeExpr(QueryOpExpr[TextRangesInScope]):
+    """Expression for getting the scope of a query."""
+
+    range_selectors: list[IQueryTextRangeSelector]
+
+    async def eval(self, context: QueryEvalContext) -> TextRangesInScope:
+        """Evaluate the expression and return the text ranges in scope."""
+        ranges_in_scope = TextRangesInScope()
+        for selector in self.range_selectors:
+            range_collection = await selector.eval(context)
+            if range_collection is not None:
+                ranges_in_scope.add_text_ranges(range_collection)
+        return ranges_in_scope
+
+
+# TODO: SelectInScopeExpr
+
+
+@dataclass
+class TextRangesInDateRangeSelector(IQueryTextRangeSelector):
+    date_range_in_scope: DateRange
+
+    async def eval(
+        self,
+        context: QueryEvalContext,
+        semantic_refs: SemanticRefAccumulator | None = None,
+    ) -> TextRangeCollection | None:
+        """Evaluate the selector and return text ranges in the specified date range."""
+        text_ranges_in_scope = TextRangeCollection()
+
+        if context.timestamp_index is not None:
+            text_ranges = await context.timestamp_index.lookup_range(
+                self.date_range_in_scope,
+            )
+            for time_range in text_ranges:
+                text_ranges_in_scope.add_range(time_range.range)
+        else:
+            text_range = await get_text_range_for_date_range(
+                context.conversation,
+                self.date_range_in_scope,
+            )
+            if text_range is not None:
+                text_ranges_in_scope.add_range(text_range)
+
+        return text_ranges_in_scope
+
+
+# TODO: TextRangesPredicateSelector
+# TODO: TextRangesWithTagSelector
+# TODO: TextRangesFromSemanticRefsSelector
+
+
+@dataclass
+class TextRangesFromMessagesSelector(IQueryTextRangeSelector):
+    source_expr: IQueryOpExpr[MessageAccumulator]
+
+    async def eval(
+        self,
+        context: QueryEvalContext,
+        semantic_refs: SemanticRefAccumulator | None = None,
+    ) -> TextRangeCollection | None:
+        matches = await self.source_expr.eval(context)
+        ranges_in_scope: list[TextRange] | None = None
+        if matches:
+            all_ordinals = sorted(matches.get_matched_values())
+            ranges_in_scope = text_ranges_from_message_ordinals(all_ordinals)
+        return TextRangeCollection(ranges_in_scope)
+
+
+# TODO: Move to messagelib.py
+def text_ranges_from_message_ordinals(
+    message_ordinals: list[MessageOrdinal],
+) -> list[TextRange]:
+    return [text_range_from_message(ordinal) for ordinal in message_ordinals]
+
+
+# TODO: Move to messagelib.py
+def text_range_from_message(message_ordinal: MessageOrdinal) -> TextRange:
+    return TextRange(start=TextLocation(message_ordinal))
+
+
+# TODO: ThreadSelector
+# TODO: to_grouped_search_results
+
+
+def to_grouped_search_results(
+    eval_results: dict[KnowledgeType, SemanticRefAccumulator],
+) -> dict[KnowledgeType, SemanticRefSearchResult]:
+    semantic_ref_matches: dict[KnowledgeType, SemanticRefSearchResult] = {}
+    for typ, accumulator in eval_results.items():
+        if len(accumulator) > 0:
+            semantic_ref_matches[typ] = SemanticRefSearchResult(
+                term_matches=accumulator.search_term_matches,
+                semantic_ref_matches=accumulator.to_scored_semantic_refs(),
+            )
+    return semantic_ref_matches
+
+
+@dataclass
+class MessagesFromKnowledgeExpr(QueryOpExpr[MessageAccumulator]):
+    src_expr: (
+        IQueryOpExpr[dict[KnowledgeType, SemanticRefSearchResult]]
+        | dict[KnowledgeType, SemanticRefSearchResult]
+    )
+
+    async def eval(self, context: QueryEvalContext) -> MessageAccumulator:
+        knowledge = (
+            self.src_expr
+            if isinstance(self.src_expr, dict)
+            else await self.src_expr.eval(context)
+        )
+        return await message_matches_from_knowledge_matches(
+            context.semantic_refs, knowledge
+        )
+
+
+# TODO: SelectMessagesInCharBudget
+
+
+@dataclass
+class RankMessagesBySimilarityExpr(QueryOpExpr[MessageAccumulator]):
+    src_expr: IQueryOpExpr[MessageAccumulator]
+    embedding: NormalizedEmbedding
+    max_messages: int | None = None
+    threshold_score: float | None = None
+
+    async def eval(self, context: QueryEvalContext) -> MessageAccumulator:
+        matches = await self.src_expr.eval(context)
+        if self.max_messages is not None and len(matches) <= self.max_messages:
+            return matches
+
+        # Try to use the message embedding index for re-ranking if available.
+        message_index = (
+            None
+            if context.conversation.secondary_indexes is None
+            else context.conversation.secondary_indexes.message_index
+        )
+        if isinstance(message_index, IMessageTextEmbeddingIndex):
+            message_ordinals = await self._get_message_ordinals_in_index(
+                message_index, matches
+            )
+            if len(message_ordinals) == len(matches):
+                matches.clear_matches()
+                ranked_messages = message_index.lookup_in_subset_by_embedding(
+                    self.embedding,
+                    message_ordinals,
+                    self.max_messages,
+                    self.threshold_score,
+                )
+                for match in ranked_messages:
+                    matches.add(match.message_ordinal, match.score)
+                return matches
+
+        if self.max_messages is not None:
+            # Can't re rank, so just take the top K from what we already have.
+            matches.select_top_n_scoring(self.max_messages)
+        return matches
+
+    async def _get_message_ordinals_in_index(
+        self, message_index, matches: MessageAccumulator
+    ):
+        message_ordinals: list[MessageOrdinal] = []
+        index_size = await message_index.size()
+        for message_ordinal in matches.get_matched_values():
+            if message_ordinal >= index_size:
+                break
+            message_ordinals.append(message_ordinal)
+        return message_ordinals
+
+
+@dataclass
+class GetScoredMessagesExpr(QueryOpExpr[list[ScoredMessageOrdinal]]):
+    src_expr: IQueryOpExpr[MessageAccumulator]
+
+    async def eval(self, context: QueryEvalContext) -> list[ScoredMessageOrdinal]:
+        matches = await self.src_expr.eval(context)
+        return matches.to_scored_message_ordinals()
+
+
+@dataclass
+class MatchMessagesBooleanExpr(IQueryOpExpr[MessageAccumulator]):
+    term_expressions: list[
+        IQueryOpExpr[SemanticRefAccumulator | MessageAccumulator | None]
+    ]
+
+    def _begin_match(self, context: QueryEvalContext) -> None:
+        context.clear_matched_terms()
+
+    async def _accumulate_messages(
+        self,
+        context: QueryEvalContext,
+        semantic_ref_matches: SemanticRefAccumulator,
+    ) -> MessageAccumulator:
+        message_matches = MessageAccumulator()
+        for semantic_ref_match in semantic_ref_matches:
+            semantic_ref = await context.get_semantic_ref(semantic_ref_match.value)
+            message_matches.add_messages_for_semantic_ref(
+                semantic_ref,
+                semantic_ref_match.score,
+            )
+        return message_matches
+
+
+@dataclass
+class MatchMessagesOrExpr(MatchMessagesBooleanExpr):
+
+    async def eval(self, context: QueryEvalContext) -> MessageAccumulator:
+        self._begin_match(context)
+
+        all_matches: MessageAccumulator | None = None
+        for match_expr in self.term_expressions:
+            matches = await match_expr.eval(context)
+            if not matches:
+                continue
+            if isinstance(matches, SemanticRefAccumulator):
+                message_matches = await self._accumulate_messages(context, matches)
+            else:
+                message_matches = matches
+            if all_matches is not None:
+                all_matches.add_union(message_matches)
+            else:
+                all_matches = message_matches
+        if all_matches is not None:
+            all_matches.calculate_total_score()
+        else:
+            all_matches = MessageAccumulator()
+        return all_matches
+
+
+@dataclass
+class MatchMessagesAndExpr(MatchMessagesBooleanExpr):
+
+    async def eval(self, context: QueryEvalContext) -> MessageAccumulator:
+        self._begin_match(context)
+
+        all_matches: MessageAccumulator | None = None
+        all_done = False
+        for match_expr in self.term_expressions:
+            matches = await match_expr.eval(context)
+            if not matches:
+                # If any expr does not match, the AND fails.
+                break
+            if isinstance(matches, SemanticRefAccumulator):
+                message_matches = await self._accumulate_messages(context, matches)
+            else:
+                message_matches = matches
+            if all_matches is None:
+                all_matches = message_matches
+            else:
+                # Intersect the message matches
+                all_matches = all_matches.intersect(message_matches)
+                if not all_matches:
+                    # If the intersection is empty, we can stop early.
+                    break
+        else:
+            # If we did not break, all terms matched.
+            all_done = True
+
+        if all_matches is not None:
+            if all_done:
+                all_matches.calculate_total_score()
+                all_matches.select_with_hit_count(len(self.term_expressions))
+            else:
+                all_matches.clear_matches()
+        else:
+            all_matches = MessageAccumulator()
+        return all_matches
+
+
+@dataclass
+class MatchMessagesOrMaxExpr(MatchMessagesOrExpr):
+
+    async def eval(self, context: QueryEvalContext) -> MessageAccumulator:
+        matches = await super().eval(context)
+        max_hit_count = matches.get_max_hit_count()
+        if max_hit_count > 1:
+            matches.select_with_hit_count(max_hit_count)
+        return matches
+
+
+# TODO: class MatchMessagesBySimilarityExpr(QueryOpExpr[list[ScoredMessageOrdinal]]):
+
+
+class NoOpExpr[T](QueryOpExpr[T]):
+    def __init__(self, src_expr: IQueryOpExpr[T]) -> None:
+        self.src_expr = src_expr
+        super().__init__()
+
+    async def eval(self, context: QueryEvalContext) -> T:
+        return await self.src_expr.eval(context)
+
+
+async def message_matches_from_knowledge_matches(
+    semantic_refs: ISemanticRefCollection,
+    knowledge_matches: dict[KnowledgeType, SemanticRefSearchResult],
+    intersect_across_knowledge_types: bool = True,
+) -> MessageAccumulator:
+    message_matches = MessageAccumulator()
+    knowledge_type_hit_count = 0  # How many types of knowledge matched?
+    for knowledge_type, matches_by_type in knowledge_matches.items():
+        if matches_by_type and matches_by_type.semantic_ref_matches:
+            knowledge_type_hit_count += 1
+            for match in matches_by_type.semantic_ref_matches:
+                message_matches.add_messages_for_semantic_ref(
+                    await semantic_refs.get_item(match.semantic_ref_ordinal),
+                    match.score,
+                )
+    if intersect_across_knowledge_types and knowledge_type_hit_count > 0:
+        # Intersect the sets of messages that matched each knowledge type
+        relevant_messages = message_matches.get_with_hit_count(knowledge_type_hit_count)
+        if relevant_messages:
+            message_matches = MessageAccumulator(relevant_messages)
+    message_matches.smooth_scores()
+    return message_matches
+
+
+# TODO: Implement proper SelectMessagesInCharBudget functionality
+@dataclass
+class SelectMessagesInCharBudget(QueryOpExpr[MessageAccumulator]):
+    """Selects messages within a character budget."""
+
+    src_expr: IQueryOpExpr[MessageAccumulator]
+    max_chars: int
+
+    async def eval(self, context: QueryEvalContext) -> MessageAccumulator:
+        matches = await self.src_expr.eval(context)
+        await matches.select_messages_in_budget(context.messages, self.max_chars)
+        return matches
+
+
+# TODO: Implement proper KnowledgeTypePredicate functionality
+@dataclass
+class KnowledgeTypePredicate(IQuerySemanticRefPredicate):
+    """Predicate to filter by knowledge type."""
+
+    knowledge_type: KnowledgeType
+
+    async def eval(self, context: QueryEvalContext, semantic_ref: SemanticRef) -> bool:
+        return semantic_ref.knowledge.knowledge_type == self.knowledge_type
+
+
+# TODO: Implement proper ThreadSelector functionality
+@dataclass
+class ThreadSelector(IQueryTextRangeSelector):
+    """Selector for thread-based text ranges."""
+
+    threads: list[Thread]
+
+    async def eval(
+        self,
+        context: QueryEvalContext,
+        semantic_refs: SemanticRefAccumulator | None = None,
+    ) -> TextRangeCollection | None:
+        text_ranges = TextRangeCollection()
+        for thread in self.threads:
+            text_ranges.add_ranges(list(thread.ranges))
+        return text_ranges