# typeagent Performance Optimization

Upstream performance improvements to [microsoft/typeagent-py](https://github.com/microsoft/typeagent-py), a structured RAG library (ingest, index, query) for Python.

## Background

typeagent-py is Microsoft's Python library for structured knowledge processing — ingesting conversations, building semantic indexes, and querying them with LLM-backed answer generation. Profiling `import typeagent` revealed unnecessary eager imports pulling in heavy dependencies (like `black`, a code formatter) at module load time, even when they're only used in cold debug/formatting paths.

In addition to optimizing the core library, we optimized the vector search hot paths in a community contributor's open optimization PR ([microsoft/typeagent-py#228](https://github.com/microsoft/typeagent-py/pull/228) — "Auto-tune Embedding Model Parameters & Add Benchmarking Tool"). That PR was itself a performance effort — tuning embedding model parameters and adding benchmarking tooling — and we further improved its code with 1.4x–14.2x speedups on the search paths it touches.

## Results

### Import Time (hyperfine, 30 runs, Standard_D2s_v5)

| Benchmark | Before | After | Speedup |
|---|---:|---:|---:|
| `import typeagent` | 791 ms | 683 ms | **1.16x** |

Import-time breakdown by PR:

| PR | Before | After | Δ | Cumulative |
|---|---:|---:|---:|---:|
| #229 defer-black | 791 ms | 713 ms | −78 ms | 713 ms |
| #235 optional-black | 713 ms | 734 ms* | — | 734 ms |
| #236 defer-query-imports | 734 ms | 683 ms | −51 ms | 683 ms |

\* PR #235 replaces `black` with stdlib `pprint` — no import-time change expected; the 713→734 delta is measurement noise within hyperfine variance.

### Offline E2E Test Suite (hyperfine, 10 runs, Standard_D2s_v5)

| Benchmark | Before | After | Speedup |
|---|---:|---:|---:|
| 69 offline tests | 5.72 s | 5.60 s | **1.02x** |

### Indexing Pipeline (pytest-async-benchmark pedantic, 20 rounds, Standard_D2s_v5)

| Benchmark | Before (min) | After (min) | Speedup |
|---|---:|---:|---:|
| `add_messages_with_indexing` (200 msgs) | 28.8 ms | 25.0 ms | **1.16x** |
| `add_messages_with_indexing` (50 msgs) | 7.8 ms | 6.7 ms | **1.16x** |
| VTT ingest (40 msgs) | 6.9 ms | 6.1 ms | **1.14x** |

> Consistent ~14-16% improvement across all message counts. Only the hot path is timed — setup (DB creation, storage init) and teardown (close, delete) are excluded via `async_benchmark.pedantic()`. All 69 tests pass before and after.

### Query (pytest-async-benchmark pedantic, 200 rounds, Standard_D2s_v5)

| Benchmark | Before (median) | After (median) | Speedup |
|---|---:|---:|---:|
| `lookup_term_filtered` (200 matches) | 2.652 ms | 1.260 ms | **2.10x** |
| `group_matches_by_type` (200 matches) | 2.453 ms | 992 μs | **2.47x** |
| `get_scored_semantic_refs_from_ordinals_iter` (200 matches) | 2.511 ms | 2.979 ms | 0.84x |
| `lookup_property_in_property_index` (200 matches) | 24.484 ms | 9.376 ms | **2.61x** |
| `get_matches_in_scope` (200 matches) | 24.062 ms | 9.185 ms | **2.62x** |

> 200 matches against a 200-message indexed SQLite transcript. Only the function under test is timed. Includes batch metadata query, binary search in `contains_range`, inline tuple comparisons in `TextRange`, and skipping pydantic validation in `get_metadata_multiple`.

### Vector Search (pytest-async-benchmark pedantic, 200 rounds, Standard_D2s_v5)

| Benchmark | Before (min) | After (min) | Speedup |
|---|---:|---:|---:|
| `fuzzy_lookup_embedding` (1K vecs) | 257 us | 70 us | **3.7x** |
| `fuzzy_lookup_embedding` (10K vecs) | 5.72 ms | 559 us | **10.2x** |
| `fuzzy_lookup_embedding` (10K + predicate) | 4.79 ms | 3.41 ms | **1.4x** |
| `fuzzy_lookup_embedding_in_subset` (1K of 10K) | 3.45 ms | 243 us | **14.2x** |

> 384-dim embeddings, normalized. The no-predicate path (most common in practice) sees the largest gains by staying entirely in numpy. The subset lookup benefits from computing dot products only for subset indices instead of all vectors. This optimization applies to code from [microsoft/typeagent-py#228](https://github.com/microsoft/typeagent-py/pull/228) (not yet merged upstream); the PR was opened against the contributor's fork at [shreejaykurhade/typeagent-py#1](https://github.com/shreejaykurhade/typeagent-py/pull/1) and has been merged there.

## What We Changed

### Startup / Import

**Defer `black` import to first use** ([KRRT7/typeagent-py#4](https://github.com/KRRT7/typeagent-py/pull/4))

`black` (the code formatter) was imported at module level in two files but only used in cold formatting paths:

- **`knowpro/answers.py`** — `black.format_str()` called only in `create_context_prompt()` to pretty-print debug context
- **`aitools/utils.py`** — `black.format_str()` called only in `format_code()` for terminal output formatting

Moved `import black` inside each function. `black` pulls in `pathspec`, `platformdirs`, `tomli`, and its own parser — none of which are needed until someone actually formats code.

**Replace `black` with stdlib `pprint` for runtime formatting** ([microsoft/typeagent-py#235](https://github.com/microsoft/typeagent-py/pull/235))

After deferring `black` in #229, we went further: removed `black` from runtime dependencies entirely. The two call sites (`create_context_prompt` in `answers.py` and `format_code`/`pretty_print` in `utils.py`) only format Python data structures — `pprint.pformat` produces equivalent output with zero external dependencies. `format_code` now uses `ast.literal_eval` to round-trip `repr()` strings back to objects for `pprint` formatting. `black` moved to dev-only dependencies.

**Defer query-time imports in `conversation_base`** ([microsoft/typeagent-py#236](https://github.com/microsoft/typeagent-py/pull/236))

`conversation_base.py` eagerly imported `answers`, `searchlang`, `search_query_schema`, and `answer_response_schema` at module level. These modules are only needed when `query()` is called — not during ingestion or indexing. Moved these imports into the `query()` method body and added `TYPE_CHECKING` guards for type annotations. Saves ~51ms (734ms → 683ms).

The largest remaining import-time cost is `pydantic_ai` at ~161ms — this is an upstream issue outside typeagent's control, noted in the PR body.

**Reproducer** (upstream-friendly, uses `$RUNNER`):
```bash
# Before (main)
$RUNNER -X importtime -c 'import typeagent' 2>&1 | sort -t'|' -k1 -rn | head -20

# After (optimization branch)
# black no longer appears in the import chain at all
```

### Runtime / Indexing

**Batch SQLite INSERTs for indexing pipeline** ([KRRT7/typeagent-py#5](https://github.com/KRRT7/typeagent-py/pull/5))

The indexing pipeline (`add_messages_with_indexing`) was issuing individual `cursor.execute()` calls for every semantic ref term and property — over 1000 individual INSERT calls for 200 messages. Added `add_terms_batch` and `add_properties_batch` to the interface protocols, with SQLite backends using `executemany` to batch all inserts into 2–3 calls. Restructured the callers (`add_metadata_to_index_from_list`, `add_to_property_index`) to collect data via pure functions first, then batch-insert.

**Reverted: Batch schema DDL + pre-compile regex** — no measurable gain, reverted.

### Vector Search

**Numpy vectorized fuzzy lookup** ([shreejaykurhade/typeagent-py#1](https://github.com/shreejaykurhade/typeagent-py/pull/1))

The `fuzzy_lookup_embedding` hot path was building Python `ScoredInt` objects for every vector, then sorting the full list. `fuzzy_lookup_embedding_in_subset` delegated to the full-scan method with a `lambda i: i in ordinals_set` predicate, computing dot products for *all* vectors then filtering.

Replaced with numpy-native operations:
- **No-predicate path**: `np.flatnonzero` for score filtering + `np.argpartition` for O(n) top-k selection — `ScoredInt` only created for final k results
- **Predicate path**: numpy pre-filters by score threshold, applies predicate only to passing candidates
- **Subset lookup**: `self._vectors[subset]` fancy indexing computes dot products only for subset indices, then same fast top-k path

### Query

**Batch metadata query across 5 N+1 call sites** ([KRRT7/typeagent-py#7](https://github.com/KRRT7/typeagent-py/pull/7))

Five call sites used `get_item()` per scored ref — one SELECT and full deserialization per match (N+1 pattern). Profiling showed 64% of per-row cost was deserializing `knowledge_json`, which the filters never use — they only check `knowledge_type` (a plain DB column) and/or `range` (needs `json.loads(range_json)` only).

Added `get_metadata_multiple` to `ISemanticRefCollection` that fetches only `semref_id, range_json, knowledge_type` in a single batch query, skipping `json.loads(knowledge_json)` and `deserialize_knowledge()` entirely. Replaced the N+1 loop with one `get_metadata_multiple` call at each site.

Further optimized the scope-filtering path (benchmarks #4, #5 went from 1.08x/1.06x to 2.61x/2.62x):
- **Binary search in `TextRangeCollection.contains_range`**: replaced O(n) linear scan with `bisect_right` keyed on `start` — O(1) for the common case of non-overlapping point ranges
- **Inline tuple comparisons in `TextRange`**: replaced `TextLocation` object allocation in `__eq__`/`__lt__`/`__contains__` with a shared `_effective_end` returning tuples
- **Skip pydantic validation in `get_metadata_multiple`**: construct `TextLocation`/`TextRange` directly from JSON instead of going through `__pydantic_validator__`

Call sites:
1. `lookup_term_filtered` — batch metadata, filter by knowledge_type/range (2.10x)
2. `group_matches_by_type` — batch metadata, group by knowledge_type (2.47x)
3. `get_scored_semantic_refs_from_ordinals_iter` — two-phase: metadata filter then batch fetch
4. `lookup_property_in_property_index` — batch metadata + bisect + inline comparisons (2.61x)
5. `get_matches_in_scope` — batch metadata + bisect + inline comparisons (2.62x)

### Bugfix

**Fix parse_azure_endpoint passing query string to AsyncAzureOpenAI** ([microsoft/typeagent-py#231](https://github.com/microsoft/typeagent-py/pull/231))

`parse_azure_endpoint` returned the full URL including `?api-version=...`. `AsyncAzureOpenAI` appends `/openai/` to `azure_endpoint`, producing a mangled URL. Now strips the query string with `str.split("?", 1)[0]`. Added 6 unit tests.

### Remaining Targets (identified, not yet implemented)

From `python -X importtime` profiling on the optimization branch (683ms baseline), sorted by self-time:

| Module | Self-time | Notes |
|---|---:|---|
| `pydantic_ai` (total) | ~161 ms | Upstream dependency — largest single cost, outside typeagent control |
| `pydantic_ai.messages` | 22 ms | Heavy pydantic model definitions |
| `knowledge_schema` | 17 ms | Schema initialization |
| `griffe` | 11 ms | Used for code introspection |
| `annotated_types` | 9 ms | Type annotation overhead |

## Upstream Contributions

Stacked PRs — merge in order (#231 first, each subsequent PR builds on the previous):

| Order | PR | Status | Description |
|---|---|---|---|
| 1 | [microsoft/typeagent-py#231](https://github.com/microsoft/typeagent-py/pull/231) | **Merged** | Fix parse_azure_endpoint passing query string to AsyncAzureOpenAI |
| 2 | [microsoft/typeagent-py#229](https://github.com/microsoft/typeagent-py/pull/229) | **Merged** | Defer black import to first use |
| 3 | [microsoft/typeagent-py#230](https://github.com/microsoft/typeagent-py/pull/230) | Open | Batch SQLite INSERTs for indexing pipeline |
| 4 | [microsoft/typeagent-py#232](https://github.com/microsoft/typeagent-py/pull/232) | Open | Batch metadata query to avoid N+1 across 5 call sites |
| 5 | [microsoft/typeagent-py#234](https://github.com/microsoft/typeagent-py/pull/234) | Open | Vectorize fuzzy_lookup_embedding with numpy ops |
| 6 | [microsoft/typeagent-py#235](https://github.com/microsoft/typeagent-py/pull/235) | Open | Replace black with stdlib pprint for runtime formatting |
| 7 | [microsoft/typeagent-py#236](https://github.com/microsoft/typeagent-py/pull/236) | Open | Defer query-time imports in conversation_base |
| — | [shreejaykurhade/typeagent-py#1](https://github.com/shreejaykurhade/typeagent-py/pull/1) | **Merged** | Numpy vectorized fuzzy lookup (against PR #228's fork) |

### Fork PRs

| PR | Type | Description |
|---|---|---|
| [KRRT7/typeagent-py#3](https://github.com/KRRT7/typeagent-py/pull/3) | Stacked draft | Cumulative optimizations (all changes) |
| [KRRT7/typeagent-py#4](https://github.com/KRRT7/typeagent-py/pull/4) | Individual | Defer black import in answers.py and utils.py |
| [KRRT7/typeagent-py#5](https://github.com/KRRT7/typeagent-py/pull/5) | Individual | Batch SQLite INSERTs for indexing pipeline |
| [KRRT7/typeagent-py#6](https://github.com/KRRT7/typeagent-py/pull/6) | Individual | Benchmark tests for indexing pipeline |
| KRRT7/typeagent-py `perf/vectorbase-lookup` | Branch | Numpy vectorized fuzzy lookup (PR opened against contributor fork) |
| [KRRT7/typeagent-py#7](https://github.com/KRRT7/typeagent-py/pull/7) | Individual | Batch metadata query across 5 N+1 call sites + parse_azure_endpoint bugfix |

## Methodology

### Environment

- **VM**: Azure Standard_D2s_v5 (2 vCPU, 8 GB RAM, non-burstable)
- **OS**: Ubuntu 24.04 LTS
- **Region**: westus2
- **Python**: 3.13 via uv
- **Tooling**: hyperfine (warmup 5, min-runs 30), `python -X importtime`, pytest-async-benchmark (pedantic mode, 20 rounds, 3 warmup)

Non-burstable VM chosen for consistent CPU performance — no thermal throttling or turbo variability.

### Profiling approach

1. `python -X importtime -c 'import typeagent'` — identified heaviest imports by self-time
2. Traced each heavy import to its call site — checked if it's needed at module level or only in a cold path
3. hyperfine A/B comparison (`bench_ab.sh main optimization`) — validated every change end-to-end
4. pytest-async-benchmark for isolated runtime benchmarks (indexing pipeline, avoids import-time noise)
5. Full offline test suite (69 tests) run before and after every change to catch regressions

### Runner convention

Benchmark scripts use `.venv/bin/python` directly for accuracy (`uv run` adds ~50% overhead and 2.5x variance). Upstream reproducers use `uv run python` for portability.

### Benchmark harness

All scripts provisioned via cloud-init on the VM:

| Script | Purpose |
|---|---|
| `bench_import.sh` | `import typeagent` time via hyperfine |
| `bench_tests.sh` | Offline E2E test suite via hyperfine |
| `bench_baseline.sh` | Run both import + test benchmarks |
| `bench_compare.sh` | Single branch benchmark (checkout + rebuild + measure) |
| `bench_ab.sh` | Side-by-side A/B comparison of two branches |

### Raw data

Results tracked in [`data/results.tsv`](data/results.tsv).

## Repo Structure

```
.
├── README.md              # This file
├── bench/                 # Benchmark scripts
├── data/                  # Raw benchmark data
│   └── results.tsv
└── infra/                 # VM provisioning
    ├── cloud-init.yaml
    └── vm-manage.sh
```
