typeagent-py 0.1.0__py3-none-any.whl

typeagent/aitools/auth.py

@@ -0,0 +1,61 @@
+#!/usr/bin/env python
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from dataclasses import dataclass
+import time
+from typing import Protocol
+
+from azure.identity import DefaultAzureCredential
+
+
+class IAccessToken(Protocol):
+    @property
+    def token(self) -> str: ...
+    @property
+    def expires_on(self) -> int:  # Posix timestamp
+        ...
+
+
+@dataclass
+class AzureTokenProvider:
+    # Note that the Python library has no async support!
+
+    def __init__(self):
+        self.credential = DefaultAzureCredential()
+        self.access_token: IAccessToken | None = None
+
+    def get_token(self) -> str:
+        if self.needs_refresh():
+            return self.refresh_token()
+        else:
+            assert self.access_token is not None
+            return self.access_token.token
+
+    def refresh_token(self) -> str:
+        self.access_token = self.credential.get_token(
+            "https://cognitiveservices.azure.com/.default"
+        )
+        assert self.access_token is not None
+        return self.access_token.token
+
+    def needs_refresh(self) -> bool:
+        return (
+            self.access_token is None
+            or self.access_token.expires_on - time.time() <= 300
+        )
+
+
+_shared_token_provider: AzureTokenProvider | None = None
+
+
+def get_shared_token_provider() -> AzureTokenProvider:
+    global _shared_token_provider
+    if _shared_token_provider is None:
+        _shared_token_provider = AzureTokenProvider()
+    return _shared_token_provider
+
+
+if __name__ == "__main__":
+    # Usage: eval `./typeagent/aitools/auth.py`
+    print(f"export AZURE_OPENAI_API_KEY={AzureTokenProvider().get_token()}")

typeagent/aitools/embeddings.py

@@ -0,0 +1,232 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import asyncio
+import os
+import re
+
+import numpy as np
+from numpy.typing import NDArray
+from openai import AsyncOpenAI, AsyncAzureOpenAI, OpenAIError
+
+from .auth import get_shared_token_provider, AzureTokenProvider
+from .utils import timelog
+
+type NormalizedEmbedding = NDArray[np.float32]  # A single embedding
+type NormalizedEmbeddings = NDArray[np.float32]  # An array of embeddings
+
+
+DEFAULT_MODEL_NAME = "text-embedding-ada-002"
+DEFAULT_EMBEDDING_SIZE = 1536  # Default embedding size (required for ada-002)
+DEFAULT_ENVVAR = "AZURE_OPENAI_ENDPOINT_EMBEDDING"
+TEST_MODEL_NAME = "test"
+
+model_to_embedding_size_and_envvar: dict[str, tuple[int | None, str]] = {
+    DEFAULT_MODEL_NAME: (DEFAULT_EMBEDDING_SIZE, DEFAULT_ENVVAR),
+    "text-embedding-small": (None, "AZURE_OPENAI_ENDPOINT_EMBEDDING_3_SMALL"),
+    "text-embedding-large": (None, "AZURE_OPENAI_ENDPOINT_EMBEDDING_3_LARGE"),
+    # For testing only, not a real model (insert real embeddings above)
+    TEST_MODEL_NAME: (3, "SIR_NOT_APPEARING_IN_THIS_FILM"),
+}
+
+
+class AsyncEmbeddingModel:
+    model_name: str
+    embedding_size: int
+    endpoint_var: str
+    azure_token_provider: AzureTokenProvider | None
+    async_client: AsyncOpenAI | None
+    azure_endpoint: str
+    azure_api_version: str
+
+    _embedding_cache: dict[str, NormalizedEmbedding]
+
+    def __init__(
+        self, embedding_size: int | None = None, model_name: str | None = None
+    ):
+        if model_name is None:
+            model_name = DEFAULT_MODEL_NAME
+        self.model_name = model_name
+
+        required_embedding_size, endpoint_envvar = (
+            model_to_embedding_size_and_envvar.get(model_name, (None, None))
+        )
+        if required_embedding_size is not None:
+            if embedding_size is not None and embedding_size != required_embedding_size:
+                raise ValueError(
+                    f"Embedding size {embedding_size} does not match "
+                    f"required size {required_embedding_size} for model {model_name}."
+                )
+            embedding_size = required_embedding_size
+        if embedding_size is None or embedding_size <= 0:
+            embedding_size = DEFAULT_EMBEDDING_SIZE
+        self.embedding_size = embedding_size
+
+        if not endpoint_envvar:
+            raise ValueError(
+                f"Model {model_name} is not supported. "
+                f"Supported models are: {', '.join(model_to_embedding_size_and_envvar.keys())}"
+            )
+        self.endpoint_envvar = endpoint_envvar
+
+        self.azure_token_provider = None
+
+        if self.model_name == TEST_MODEL_NAME:
+            self.async_client = None
+        else:
+            openai_key_name = "OPENAI_API_KEY"
+            azure_key_name = "AZURE_OPENAI_API_KEY"
+            if os.getenv(openai_key_name):
+                with timelog(f"Using OpenAI"):
+                    self.async_client = AsyncOpenAI()
+            elif azure_api_key := os.getenv(azure_key_name):
+                with timelog("Using Azure OpenAI"):
+                    self._setup_azure(azure_api_key)
+            else:
+                raise ValueError(
+                    f"Neither {openai_key_name} nor {azure_key_name} found in environment."
+                )
+
+        self._embedding_cache = {}
+
+    def _setup_azure(self, azure_api_key: str) -> None:
+        # TODO: support different endpoint names
+        endpoint_envvar = self.endpoint_envvar
+        azure_endpoint = os.environ.get(endpoint_envvar)
+        if not azure_endpoint:
+            raise ValueError(f"Environment variable {endpoint_envvar} not found.")
+        m = re.search(r"[?,]api-version=([^,]+)$", azure_endpoint)
+        if not m:
+            raise ValueError(
+                f"{endpoint_envvar}={azure_endpoint} "
+                f"doesn't end in api-version=<version>"
+            )
+        self.azure_endpoint = azure_endpoint
+        self.azure_api_version = m.group(1)
+        if azure_api_key.lower() == "identity":
+            self.azure_token_provider = get_shared_token_provider()
+            azure_api_key = self.azure_token_provider.get_token()
+            # print("Using shared TokenProvider")
+        self.async_client = AsyncAzureOpenAI(
+            api_version=self.azure_api_version,
+            azure_endpoint=self.azure_endpoint,
+            api_key=azure_api_key,
+        )
+
+    async def refresh_auth(self):
+        """Update client when using a token provider and it's nearly expired."""
+        # refresh_token is synchronous and slow -- run it in a separate thread
+        assert self.azure_token_provider
+        refresh_token = self.azure_token_provider.refresh_token
+        loop = asyncio.get_running_loop()
+        azure_api_key = await loop.run_in_executor(None, refresh_token)
+        assert self.azure_api_version
+        assert self.azure_endpoint
+        self.async_client = AsyncAzureOpenAI(
+            api_version=self.azure_api_version,
+            azure_endpoint=self.azure_endpoint,
+            api_key=azure_api_key,
+        )
+
+    def add_embedding(self, key: str, embedding: NormalizedEmbedding) -> None:
+        existing = self._embedding_cache.get(key)
+        if existing is not None:
+            assert np.array_equal(existing, embedding)
+        else:
+            self._embedding_cache[key] = embedding
+
+    async def get_embedding_nocache(self, input: str) -> NormalizedEmbedding:
+        embeddings = await self.get_embeddings_nocache([input])
+        return embeddings[0]
+
+    async def get_embeddings_nocache(self, input: list[str]) -> NormalizedEmbeddings:
+        if not input:
+            empty = np.array([], dtype=np.float32)
+            empty.shape = (0, self.embedding_size)
+            return empty
+        if self.azure_token_provider and self.azure_token_provider.needs_refresh():
+            await self.refresh_auth()
+        extra_args = {}
+        if self.model_name != DEFAULT_MODEL_NAME:
+            extra_args["dimensions"] = self.embedding_size
+        if self.async_client is None:
+            # Compute a random embedding for testing purposes.
+
+            def hashish(s: str) -> int:
+                # Primitive deterministic hash function (hash() varies per run)
+                h = 0
+                for ch in s:
+                    h = (h * 31 + ord(ch)) & 0xFFFFFFFF
+                return h
+
+            prime = 1961
+            fake_data: list[NormalizedEmbedding] = []
+            for item in input:
+                if not item:
+                    raise OpenAIError
+                length = len(item)
+                floats = []
+                for i in range(self.embedding_size):
+                    cut = i % length
+                    scrambled = item[cut:] + item[:cut]
+                    hashed = hashish(scrambled)
+                    reduced = (hashed % prime) / prime
+                    floats.append(reduced)
+                array = np.array(floats, dtype=np.float64)
+                normalized = array / np.sqrt(np.dot(array, array))
+                dot = np.dot(normalized, normalized)
+                assert (
+                    abs(dot - 1.0) < 1e-15
+                ), f"Embedding {normalized} is not normalized: {dot}"
+                fake_data.append(normalized)
+            assert len(fake_data) == len(input), (len(fake_data), "!=", len(input))
+            result = np.array(fake_data, dtype=np.float32)
+            return result
+        else:
+            # TODO: Split in batches of 2048 inputs if too long;
+            # or smaller if inputs are large.
+            data = (
+                await self.async_client.embeddings.create(
+                    input=input,
+                    model=self.model_name,
+                    encoding_format="float",
+                    **extra_args,
+                )
+            ).data
+            assert len(data) == len(input), (len(data), "!=", len(input))
+            return np.array([d.embedding for d in data], dtype=np.float32)
+
+    async def get_embedding(self, key: str) -> NormalizedEmbedding:
+        """Retrieve an embedding, using the cache."""
+        if key in self._embedding_cache:
+            return self._embedding_cache[key]
+        embedding = await self.get_embedding_nocache(key)
+        self._embedding_cache[key] = embedding
+        return embedding
+
+    async def get_embeddings(self, keys: list[str]) -> NormalizedEmbeddings:
+        """Retrieve embeddings for multiple keys, using the cache."""
+        embeddings: list[NormalizedEmbedding | None] = []
+        missing_keys: list[str] = []
+
+        # Collect cached embeddings and identify missing keys
+        for key in keys:
+            if key in self._embedding_cache:
+                embeddings.append(self._embedding_cache[key])
+            else:
+                embeddings.append(None)  # Placeholder for missing keys
+                missing_keys.append(key)
+
+        # Retrieve embeddings for missing keys
+        if missing_keys:
+            new_embeddings = await self.get_embeddings_nocache(missing_keys)
+            for key, embedding in zip(missing_keys, new_embeddings):
+                self._embedding_cache[key] = embedding
+
+            # Replace placeholders with retrieved embeddings
+            for i, key in enumerate(keys):
+                if embeddings[i] is None:
+                    embeddings[i] = self._embedding_cache[key]
+        return np.array(embeddings, dtype=np.float32).reshape(
+            (len(keys), self.embedding_size)
+        )

typeagent/aitools/utils.py

@@ -0,0 +1,244 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Utilities that are hard to fit in any specific module."""
+
+from contextlib import contextmanager
+import difflib
+import os
+import re
+import shutil
+import time
+
+import black
+import colorama
+import dotenv
+import typechat
+
+from pydantic_ai import Agent
+
+cap = min  # More readable name for capping a value at some limit.
+
+
+@contextmanager
+def timelog(label: str, verbose: bool = True):
+    """Context manager to log the time taken by a block of code.
+
+    With verbose=False it prints nothing."""
+    start_time = time.time()
+    try:
+        yield
+    finally:
+        elapsed_time = time.time() - start_time
+        if verbose:
+            dim = colorama.Style.DIM
+            reset = colorama.Style.RESET_ALL
+            print(f"{dim}{elapsed_time:.3f}s -- {label}{reset}")
+
+
+def pretty_print(obj: object, prefix: str = "", suffix: str = "") -> None:
+    """Pretty-print an object using black.
+
+    NOTE: Only works if its repr() is a valid Python expression.
+    """
+    print(prefix + format_code(repr(obj)) + suffix)
+
+
+def format_code(text: str, line_width=None) -> str:
+    """Format a block of code using black, then reindent to 2 spaces.
+
+    NOTE: The text must be a valid Python expression or code block.
+    """
+    if line_width is None:
+        # Use the terminal width, but cap it to 200 characters.
+        line_width = cap(200, shutil.get_terminal_size().columns)
+    formatted_text = black.format_str(
+        text, mode=black.FileMode(line_length=line_width)
+    ).rstrip()
+    return reindent(formatted_text)
+
+
+def reindent(text: str) -> str:
+    """Reindent a block of text from 4 to 2 spaces per indent level."""
+    lines = text.splitlines()
+    reindented_lines = []
+    for line in lines:
+        stripped_line = line.lstrip()
+        twice_indent_level = (len(line) - len(stripped_line) + 1) // 2  # Round up
+        reindented_lines.append(" " * twice_indent_level + stripped_line)
+    return "\n".join(reindented_lines)
+
+
+def load_dotenv() -> None:
+    """Load environment variables from '<repo_root>/ta/.env'."""
+    paths = []
+    # Look for <repo_root>/ts/.env first.
+    repo_root = os.popen("git rev-parse --show-toplevel").read().strip()
+    if repo_root:
+        env_path = os.path.join(repo_root, "ts", ".env")
+        if os.path.exists(env_path):
+            paths.append(env_path)
+
+    # Also look in current directory and going up.
+    cur_dir = os.path.abspath(os.getcwd())
+    while True:
+        paths.append(os.path.join(cur_dir, ".env"))
+        parent_dir = os.path.dirname(cur_dir)
+        if parent_dir == cur_dir:
+            break  # Reached filesystem root ('/').
+        cur_dir = parent_dir
+
+    env_path = None
+    for path in paths:
+        # Filter out non-existing paths.
+        if os.path.exists(path):
+            env_path = path
+            break
+    if env_path:
+        dotenv.load_dotenv(env_path)
+
+
+def create_translator[T](
+    model: typechat.TypeChatLanguageModel,
+    schema_class: type[T],
+) -> typechat.TypeChatJsonTranslator[T]:
+    """Create a TypeChat translator for a given model and schema."""
+    validator = typechat.TypeChatValidator[T](schema_class)
+    translator = typechat.TypeChatJsonTranslator[T](model, validator, schema_class)
+    return translator
+
+
+# Vibe-coded by o4-mini-high
+def list_diff(label_a, a, label_b, b, max_items):
+    """Print colorized diff between two sorted list of numbers."""
+    sm = difflib.SequenceMatcher(None, a, b)
+    a_out, b_out = [], []
+    for tag, i1, i2, j1, j2 in sm.get_opcodes():
+        a_slice, b_slice = a[i1:i2], b[j1:j2]
+        L = max(len(a_slice), len(b_slice))
+        for k in range(L):
+            a_out.append(str(a_slice[k]) if k < len(a_slice) else "")
+            b_out.append(str(b_slice[k]) if k < len(b_slice) else "")
+
+    # color helpers
+    def color_a(val, other):
+        return (
+            colorama.Fore.RED + val + colorama.Style.RESET_ALL
+            if val and val != other
+            else val
+        )
+
+    def color_b(val, other):
+        return (
+            colorama.Fore.GREEN + val + colorama.Style.RESET_ALL
+            if val and val != other
+            else val
+        )
+
+    # apply color
+    a_cols = [color_a(a_out[i], b_out[i]) for i in range(len(a_out))]
+    b_cols = [color_b(b_out[i], a_out[i]) for i in range(len(b_out))]
+
+    # compute column widths
+    widths = [max(len(a_out[i]), len(b_out[i])) for i in range(len(a_out))]
+
+    # prepare labels
+    max_label = max(len(label_a), len(label_b))
+    la = label_a.ljust(max_label)
+    lb = label_b.ljust(max_label)
+
+    # split into segments
+    if max_items and max_items > 0:
+        segments = [
+            (i, min(i + max_items, len(a_cols)))
+            for i in range(0, len(a_cols), max_items)
+        ]
+    else:
+        segments = [(0, len(a_cols))]
+
+    # formatter for a row segment
+    def fmt(row, seg_widths):
+        return " ".join(f"{cell:>{w}}" for cell, w in zip(row, seg_widths))
+
+    # print each segment
+    for start, end in segments:
+        seg_widths = widths[start:end]
+        print(la, fmt(a_cols[start:end], seg_widths))
+        print(lb, fmt(b_cols[start:end], seg_widths))
+
+
+def setup_logfire():
+    """Configure logfire for pydantic_ai and httpx."""
+
+    import logfire
+
+    def scrubbing_callback(m: logfire.ScrubMatch):
+        """Instructions: Uncomment any block where you deem it safe to not scrub."""
+        # if m.path == ('attributes', 'http.request.header.authorization'):
+        #     return m.value
+
+        # if m.path == ('attributes', 'http.request.header.api-key'):
+        #     return m.value
+
+        if (
+            m.path == ("attributes", "http.request.body.text", "messages", 0, "content")
+            and m.pattern_match.group(0) == "secret"
+        ):
+            return m.value
+
+        # if m.path == ('attributes', 'http.response.header.azureml-model-session'):
+        #     return m.value
+
+    logfire.configure(scrubbing=logfire.ScrubbingOptions(callback=scrubbing_callback))
+    logfire.instrument_pydantic_ai()
+    logfire.instrument_httpx(capture_all=True)
+
+
+def make_agent[T](cls: type[T]) -> Agent[None, T]:
+    """Create Pydantic AI agent using hardcoded preferences."""
+    from pydantic_ai import NativeOutput, ToolOutput
+    from pydantic_ai.models.openai import OpenAIModel
+    from pydantic_ai.providers.azure import AzureProvider
+    from .auth import get_shared_token_provider
+
+    # Prefer straight OpenAI over Azure OpenAI.
+    if os.getenv("OPENAI_API_KEY"):
+        Wrapper = NativeOutput
+        print(f"## Using OpenAI with {Wrapper.__name__} ##")
+        model = OpenAIModel("gpt-4o")  # Retrieves OPENAI_API_KEY again.
+
+    elif azure_openai_api_key := os.getenv("AZURE_OPENAI_API_KEY"):
+        # This section is rather specific to our team's setup  at Microsoft.
+        if azure_openai_api_key == "identity":
+            token_provider = get_shared_token_provider()
+            azure_openai_api_key = token_provider.get_token()
+
+        azure_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
+        if not azure_endpoint:
+            raise RuntimeError("AZURE_OPENAI_ENDPOINT not found")
+
+        print(f"## {azure_endpoint} ##")
+        m = re.search(r"api-version=([\d-]+(?:preview)?)", azure_endpoint)
+        if not m:
+            raise RuntimeError(
+                f"AZURE_OPENAI_ENDPOINT has no valid api-version field: {azure_endpoint}"
+            )
+        api_version = m.group(1)
+        Wrapper = ToolOutput
+
+        print(f"## Using Azure {api_version} with {Wrapper.__name__} ##")
+        model = OpenAIModel(
+            "gpt-4o",
+            provider=AzureProvider(
+                azure_endpoint=azure_endpoint,
+                api_version=api_version,
+                api_key=azure_openai_api_key,
+            ),
+        )
+
+    else:
+        raise RuntimeError(
+            "Neither OPENAI_API_KEY nor AZURE_OPENAI_API_KEY was provided."
+        )
+
+    return Agent(model, output_type=Wrapper(cls, strict=True), retries=3)