typeagent-py 0.1.0__py3-none-any.whl

typeagent/aitools/vectorbase.py

@@ -0,0 +1,175 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from collections.abc import Callable, Iterable
+from dataclasses import dataclass
+
+import numpy as np
+
+from .embeddings import AsyncEmbeddingModel, NormalizedEmbedding, NormalizedEmbeddings
+
+
+@dataclass
+class ScoredInt:
+    item: int
+    score: float
+
+
+@dataclass
+class TextEmbeddingIndexSettings:
+    embedding_model: AsyncEmbeddingModel
+    embedding_size: int  # Always embedding_model.embedding_size
+    min_score: float
+    max_matches: int | None
+    retry_max_attempts: int = 2
+    retry_delay: float = 2.0  # Seconds
+    batch_size: int = 8
+
+    def __init__(
+        self,
+        embedding_model: AsyncEmbeddingModel | None = None,
+        embedding_size: int | None = None,
+        min_score: float | None = None,
+        max_matches: int | None = None,
+    ):
+        self.embedding_model = embedding_model or AsyncEmbeddingModel(embedding_size)
+        self.embedding_size = self.embedding_model.embedding_size
+        assert (
+            embedding_size is None or self.embedding_size == embedding_size
+        ), f"Given embedding size {embedding_size} doesn't match model's embedding size {self.embedding_size}"
+        self.min_score = min_score if min_score is not None else 0.85
+        self.max_matches = max_matches
+
+
+class VectorBase:
+    settings: TextEmbeddingIndexSettings
+    _vectors: NormalizedEmbeddings
+    _model: AsyncEmbeddingModel
+    _embedding_size: int
+
+    def __init__(self, settings: TextEmbeddingIndexSettings):
+        self.settings = settings
+        self._model = settings.embedding_model
+        self._embedding_size = self._model.embedding_size
+        self.clear()
+
+    async def get_embedding(self, key: str, cache: bool = True) -> NormalizedEmbedding:
+        if cache:
+            return await self._model.get_embedding(key)
+        else:
+            return await self._model.get_embedding_nocache(key)
+
+    async def get_embeddings(
+        self, keys: list[str], cache: bool = True
+    ) -> NormalizedEmbeddings:
+        if cache:
+            return await self._model.get_embeddings(keys)
+        else:
+            return await self._model.get_embeddings_nocache(keys)
+
+    def __len__(self) -> int:
+        return len(self._vectors)
+
+    # Needed because otherwise an empty index would be falsy.
+    def __bool__(self) -> bool:
+        return True
+
+    def add_embedding(
+        self, key: str | None, embedding: NormalizedEmbedding | list[float]
+    ) -> None:
+        if isinstance(embedding, list):
+            embedding = np.array(embedding, dtype=np.float32)
+        embeddings = embedding.reshape(1, -1)  # Make it 2D: 1xN
+        self._vectors = np.append(self._vectors, embeddings, axis=0)
+        if key is not None:
+            self._model.add_embedding(key, embedding)
+
+    def add_embeddings(self, embeddings: NormalizedEmbeddings) -> None:
+        assert embeddings.ndim == 2
+        assert embeddings.shape[1] == self._embedding_size
+        self._vectors = np.concatenate((self._vectors, embeddings), axis=0)
+
+    async def add_key(self, key: str, cache: bool = True) -> None:
+        embeddings = (await self.get_embedding(key, cache=cache)).reshape(1, -1)
+        self._vectors = np.append(self._vectors, embeddings, axis=0)
+
+    async def add_keys(self, keys: list[str], cache: bool = True) -> None:
+        embeddings = await self.get_embeddings(keys, cache=cache)
+        self._vectors = np.concatenate((self._vectors, embeddings), axis=0)
+
+    def fuzzy_lookup_embedding(
+        self,
+        embedding: NormalizedEmbedding,
+        max_hits: int | None = None,
+        min_score: float | None = None,
+        predicate: Callable[[int], bool] | None = None,
+    ) -> list[ScoredInt]:
+        if max_hits is None:
+            max_hits = 10
+        if min_score is None:
+            min_score = 0.0
+        # This line does most of the work:
+        scores: Iterable[float] = np.dot(self._vectors, embedding)
+        scored_ordinals = [
+            ScoredInt(i, score)
+            for i, score in enumerate(scores)
+            if score >= min_score and (predicate is None or predicate(i))
+        ]
+        scored_ordinals.sort(key=lambda x: x.score, reverse=True)
+        return scored_ordinals[:max_hits]
+
+    # TODO: Make this and fuzzy_lookup_embedding() more similar.
+    def fuzzy_lookup_embedding_in_subset(
+        self,
+        embedding: NormalizedEmbedding,
+        ordinals_of_subset: list[int],
+        max_hits: int | None = None,
+        min_score: float | None = None,
+    ) -> list[ScoredInt]:
+        return self.fuzzy_lookup_embedding(
+            embedding, max_hits, min_score, lambda i: i in ordinals_of_subset
+        )
+
+    async def fuzzy_lookup(
+        self,
+        key: str,
+        max_hits: int | None = None,
+        min_score: float | None = None,
+        predicate: Callable[[int], bool] | None = None,
+    ) -> list[ScoredInt]:
+        if max_hits is None:
+            max_hits = self.settings.max_matches
+        if min_score is None:
+            min_score = self.settings.min_score
+        embedding = await self.get_embedding(key)
+        return self.fuzzy_lookup_embedding(
+            embedding, max_hits=max_hits, min_score=min_score, predicate=predicate
+        )
+
+    def clear(self) -> None:
+        self._vectors = np.array([], dtype=np.float32)
+        self._vectors.shape = (0, self._embedding_size)
+
+    def get_embedding_at(self, pos: int) -> NormalizedEmbedding:
+        if 0 <= pos < len(self._vectors):
+            return self._vectors[pos]
+        raise IndexError(
+            f"Index {pos} out of bounds for embedding index of size {len(self)}"
+        )
+
+    def serialize_embedding_at(self, pos: int) -> NormalizedEmbedding | None:
+        return self._vectors[pos] if 0 <= pos < len(self._vectors) else None
+
+    def serialize(self) -> NormalizedEmbeddings:
+        assert self._vectors.shape == (len(self._vectors), self._embedding_size)
+        return self._vectors  # TODO: Should we make a copy?
+
+    def deserialize(self, data: NormalizedEmbeddings | None) -> None:
+        if data is None:
+            self.clear()
+            return
+        assert data.shape == (len(data), self._embedding_size), [
+            data.shape,
+            self._embedding_size,
+        ]
+        self._vectors = data  # TODO: Should we make a copy?

typeagent/knowpro/answer_context_schema.py

@@ -0,0 +1,49 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+# TODO: Are we sure this isn't used as a translator schema class?
+
+from dataclasses import dataclass
+from typing import Annotated, Any, Union
+from typing_extensions import Doc
+
+from ..knowpro.interfaces import DateRange
+
+EntityNames = Union[str, list[str]]
+
+
+@dataclass
+class RelevantKnowledge:
+    knowledge: Annotated[Any, Doc("The actual knowledge")]
+    origin: Annotated[
+        EntityNames | None, Doc("Entity or entities who mentioned the knowledge")
+    ] = None
+    audience: Annotated[
+        EntityNames | None,
+        Doc("Entity or entities who received or consumed this knowledge"),
+    ] = None
+    time_range: Annotated[
+        DateRange | None, Doc("Time period during which this knowledge was gathered")
+    ] = None
+
+
+@dataclass
+class RelevantMessage:
+    from_: Annotated[EntityNames | None, Doc("Sender(s) of the message")]
+    to: Annotated[EntityNames | None, Doc("Recipient(s) of the message")]
+    timestamp: Annotated[str | None, Doc("Timestamp of the message in ISO format")]
+    messageText: Annotated[str | list[str] | None, Doc("Text chunks in this message")]
+
+
+@dataclass
+class AnswerContext:
+    """Use empty lists for unneeded properties."""
+
+    entities: Annotated[
+        list[RelevantKnowledge],
+        Doc(
+            "Relevant entities. Use the 'name' and 'type' properties of entities to PRECISELY identify those that answer the user question."
+        ),
+    ]
+    topics: Annotated[list[RelevantKnowledge], Doc("Relevant topics")]
+    messages: Annotated[list[RelevantMessage], Doc("Relevant messages")]

typeagent/knowpro/answer_response_schema.py

@@ -0,0 +1,34 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from typing import Literal, Annotated
+from typing_extensions import Doc
+from pydantic.dataclasses import dataclass
+
+AnswerType = Literal[
+    "NoAnswer",  # If question cannot be accurately answered from [ANSWER CONTEXT]
+    "Answered",  # Fully answer question
+    # TODO: Add a category for outright errors, e.g. network errors
+]
+
+
+@dataclass
+class AnswerResponse:
+    type: Annotated[
+        AnswerType,
+        Doc(
+            'use "NoAnswer" if no highly relevant answer found in the [ANSWER CONTEXT]'
+        ),
+    ]
+    answer: Annotated[
+        str | None,
+        Doc(
+            "the answer to display if [ANSWER CONTEXT] is highly relevant and can be used to answer the user's question"
+        ),
+    ] = None
+    whyNoAnswer: Annotated[
+        str | None,
+        Doc(
+            "If NoAnswer, explain why..\nparticularly explain why you didn't use any supplied entities"
+        ),
+    ] = None