Misc changes

TODO.md

@@ -158,10 +158,6 @@ The TODO items include (in no particular order):
 
 # Older TODO action items
 
-## Retries for embeddings
-
-For robustness -- TypeChat already retries, but my embeddings don't.
-
 ## Refactoring implementations
 
 - Change inconsistent module names (Claude uses different naming style)

pyproject.toml

@@ -35,23 +35,21 @@ dependencies = [
   "mcp[cli]>=1.12.1",
   "numpy>=2.2.6",
   "openai>=1.81.0",
-  "opentelemetry-instrumentation-httpx>=0.57b0",
   "pydantic>=2.11.4",
   "pydantic-ai-slim[openai]>=0.5.0",
-  "pydantic-core>=2.41.1",
   "python-dotenv>=1.1.0",
   "typechat>=0.0.4",
   "webvtt-py>=0.5.1",
 ]
 
 [project.optional-dependencies]
 dev = [
-  "build>=1.2.2.post1",
   "coverage[toml]>=7.9.1",
   "google-api-python-client>=2.184.0",
   "google-auth-httplib2>=0.2.0",
   "google-auth-oauthlib>=1.2.2",
-  "pyright>=1.1.405",
+  "opentelemetry-instrumentation-httpx>=0.57b0",
+  "pyright==1.1.406",  # 407 has a regression
   "pytest>=8.3.5",
   "pytest-asyncio>=0.26.0",
   "pytest-mock>=3.14.0",

test/test_embeddings.py

@@ -141,15 +141,19 @@ async def test_set_endpoint(monkeypatch):
         1024, "custom_model", "INFINITY_EMBEDDING_URL"
     )
 
+    assert embedding_model.embedding_size == 1024
+    assert embedding_model.model_name == "custom_model"
+
+    # NOTE: checking openai.AsyncOpenAI internals
     assert embedding_model.async_client is not None
     assert embedding_model.async_client.base_url == "http://localhost:7997"
     assert embedding_model.async_client.api_key == "does-not-matter"
 
     with pytest.raises(
         ValueError,
         match="Environment variable for embedding endpoint WRONG_ENDPOINT does not match required environment"
-        " variable AZURE_OPENAI_ENDPOINT_EMBEDDING_3_SMALL for embedding model text-embedding-small.",
+        " variable AZURE_OPENAI_ENDPOINT_EMBEDDING_3_SMALL for embedding model text-embedding-3-small.",
     ):
         embedding_model = AsyncEmbeddingModel(
-            2000, "text-embedding-small", "WRONG_ENDPOINT"
+            2000, "text-embedding-3-small", "WRONG_ENDPOINT"
         )

test/test_factory.py

@@ -10,6 +10,8 @@
 from typeagent.knowpro.convsettings import ConversationSettings
 from typeagent.transcripts.transcript import TranscriptMessage, TranscriptMessageMeta
 
+from fixtures import really_needs_auth
+
 
 @pytest.mark.asyncio
 async def test_create_conversation_minimal():
@@ -51,7 +53,7 @@ async def test_create_conversation_with_tags():
 
 
 @pytest.mark.asyncio
-async def test_create_conversation_and_add_messages():
+async def test_create_conversation_and_add_messages(really_needs_auth):
     """Test the complete workflow: create conversation and add messages."""
     # 1. Create empty conversation
     test_model = AsyncEmbeddingModel(model_name=TEST_MODEL_NAME)

test/test_podcasts.py

@@ -1,12 +1,11 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
-import asyncio
 import os
 import pytest
 from datetime import timezone
 
-from fixtures import needs_auth, temp_dir, embedding_model  # type: ignore  # Yes they are used!
+from fixtures import really_needs_auth, temp_dir, embedding_model  # type: ignore  # Yes they are used!
 
 from typeagent.podcasts.podcast import Podcast
 from typeagent.knowpro.convsettings import ConversationSettings
@@ -18,7 +17,7 @@
 
 @pytest.mark.asyncio
 async def test_ingest_podcast(
-    needs_auth: None, temp_dir: str, embedding_model: AsyncEmbeddingModel
+    really_needs_auth: None, temp_dir: str, embedding_model: AsyncEmbeddingModel
 ):
     # Import the podcast
     settings = ConversationSettings(embedding_model)

tools/ingest_vtt.py

@@ -370,7 +370,7 @@ def save_current_message():
             )
 
             # Process messages in batches
-            batch_size = 10
+            batch_size = settings.semantic_ref_index_settings.batch_size
             successful_count = 0
             start_time = time.time()
 

typeagent/aitools/embeddings.py

@@ -3,11 +3,10 @@
 
 import asyncio
 import os
-import re
 
 import numpy as np
 from numpy.typing import NDArray
-from openai import AsyncOpenAI, AsyncAzureOpenAI, OpenAIError
+from openai import AsyncOpenAI, AsyncAzureOpenAI, DEFAULT_MAX_RETRIES, OpenAIError
 
 from .auth import get_shared_token_provider, AzureTokenProvider
 from .utils import timelog
@@ -23,8 +22,8 @@
 
 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"),
+    "text-embedding-3-small": (None, "AZURE_OPENAI_ENDPOINT_EMBEDDING_3_SMALL"),
+    "text-embedding-3-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"),
 }
@@ -46,6 +45,7 @@ def __init__(
         embedding_size: int | None = None,
         model_name: str | None = None,
         endpoint_envvar: str | None = None,
+        max_retries: int = DEFAULT_MAX_RETRIES,
     ):
         if model_name is None:
             model_name = DEFAULT_MODEL_NAME
@@ -90,7 +90,7 @@ def __init__(
                 endpoint = os.getenv(self.endpoint_envvar)
                 with timelog(f"Using OpenAI"):
                     self.async_client = AsyncOpenAI(
-                        base_url=endpoint, api_key=openai_key
+                        base_url=endpoint, api_key=openai_key, max_retries=max_retries
                     )
             elif azure_api_key := os.getenv(azure_key_name):
                 with timelog("Using Azure OpenAI"):

typeagent/aitools/utils.py

@@ -51,7 +51,7 @@ def format_code(text: str, line_width=None) -> str:
         # Use the terminal width, but cap it to 200 characters.
         line_width = min(200, shutil.get_terminal_size().columns)
     formatted_text = black.format_str(
-        text, mode=black.FileMode(line_length=line_width)
+        text, mode=black.Mode(line_length=line_width)
     ).rstrip()
     return reindent(formatted_text)
 

typeagent/aitools/vectorbase.py

@@ -5,6 +5,7 @@
 from dataclasses import dataclass
 
 import numpy as np
+from openai import DEFAULT_MAX_RETRIES
 
 from .embeddings import AsyncEmbeddingModel, NormalizedEmbedding, NormalizedEmbeddings
 
@@ -18,27 +19,34 @@ class ScoredInt:
 @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
+    embedding_size: int  # Set to embedding_model.embedding_size
+    min_score: float  # Between 0.0 and 1.0
+    max_matches: int | None  # >= 1; None means no limit
+    batch_size: int  # >= 1
+    max_retries: int
 
     def __init__(
         self,
         embedding_model: AsyncEmbeddingModel | None = None,
         embedding_size: int | None = None,
         min_score: float | None = None,
         max_matches: int | None = None,
+        batch_size: int | None = None,
+        max_retries: int | None = None,
     ):
-        self.embedding_model = embedding_model or AsyncEmbeddingModel(embedding_size)
+        self.min_score = min_score if min_score is not None else 0.85
+        self.max_matches = max_matches if max_matches and max_matches >= 1 else None
+        self.batch_size = batch_size if batch_size and batch_size >= 1 else 10
+        self.max_retries = (
+            max_retries if max_retries is not None else DEFAULT_MAX_RETRIES
+        )
+        self.embedding_model = embedding_model or AsyncEmbeddingModel(
+            embedding_size, max_retries=self.max_retries
+        )
         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:

typeagent/knowpro/answers.py

@@ -128,7 +128,7 @@ def create_context_prompt(context: AnswerContext) -> str:
     prompt = [
         "[ANSWER CONTEXT]",
         "===",
-        black.format_str(str(dictify(context)), mode=black.FileMode(line_length=200)),
+        black.format_str(str(dictify(context)), mode=black.Mode(line_length=200)),
         "===",
     ]
     return "\n".join(prompt)

typeagent/knowpro/convsettings.py

@@ -54,7 +54,7 @@ def __init__(
         )
         self.semantic_ref_index_settings = SemanticRefIndexSettings(
             batch_size=10,
-            auto_extract_knowledge=False,
+            auto_extract_knowledge=True,  # The high-level API wants this
         )
 
         # Storage provider will be created lazily if not provided