# 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.

## 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`.

## 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.

### 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) | In review | Batch SQLite INSERTs for indexing pipeline (reviewer: @bmerkle, comments addressed) |
| 4 | [microsoft/typeagent-py#232](https://github.com/microsoft/typeagent-py/pull/232) | Blocked on #230 | Batch metadata query to avoid N+1 across 5 call sites (reviewer: @bmerkle, 4 comments — will rebase after #230 merges) |
| 5 | [microsoft/typeagent-py#234](https://github.com/microsoft/typeagent-py/pull/234) | **Merged** | Vectorize fuzzy_lookup_embedding with numpy ops (see [Lessons Learned](#lessons-learned)) |
| 6 | [microsoft/typeagent-py#235](https://github.com/microsoft/typeagent-py/pull/235) | **Merged** | Replace black with stdlib pprint for runtime formatting |
| 7 | [microsoft/typeagent-py#236](https://github.com/microsoft/typeagent-py/pull/236) | **Closed** | Defer query-time imports in conversation_base (rejected — maintainer said savings not worth complexity) |
| — | [shreejaykurhade/typeagent-py#1](https://github.com/shreejaykurhade/typeagent-py/pull/1) | **Merged** | Numpy vectorized fuzzy lookup (against PR #228's fork — see [Lessons Learned](#lessons-learned)) |

### 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).

## Lessons Learned

### Failed experiment: optimizing a third-party contributor's PR

We optimized the vector search hot paths in [microsoft/typeagent-py#228](https://github.com/microsoft/typeagent-py/pull/228), a community contributor's open PR ("Auto-tune Embedding Model Parameters & Add Benchmarking Tool"). The micro-benchmarks looked impressive (3.7x–14.2x on `fuzzy_lookup_embedding`), but this was a bad target for two reasons:

1. **The bottleneck doesn't matter in practice.** Guido van Rossum (typeagent maintainer) pointed out that model API calls and SQL queries dominate real-world latency — the numpy iteration time in vector search is negligible by comparison. The 10.2x speedup on 10K vectors saves ~5ms, which is noise next to a 200ms embedding API call. Optimizing code that isn't the bottleneck in actual use is wasted effort, no matter how good the numbers look on synthetic benchmarks.

2. **We were optimizing someone else's unmerged code.** PR #228 was third-party work that hadn't been reviewed or accepted by the maintainers. The contributor's PR itself had issues — Guido later showed with real evals that the proposed min_score tuning made results *worse* (twice as many zero scores vs. main). We picked a target based on what was convenient to optimize, not what the project actually needed.

The vectorization code was technically correct (Guido approved and merged [#234](https://github.com/microsoft/typeagent-py/pull/234) without comments), but it was the wrong work to prioritize. We also accidentally leaked benchmark code into #228, adding noise to an already messy PR.

**Takeaways for target selection:**
- Profile the real workload, not synthetic benchmarks — a 14x speedup on a non-bottleneck is worth 0x
- Don't optimize unmerged third-party code as a proxy for the project's actual needs
- Validate with the maintainer that the bottleneck matters before investing in optimization

### Vector search benchmarks (for the record)

Included here for completeness — these numbers are real but don't represent meaningful optimization for typeagent.

| 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, pytest-async-benchmark pedantic 200 rounds on Standard_D2s_v5. PRs: [shreejaykurhade/typeagent-py#1](https://github.com/shreejaykurhade/typeagent-py/pull/1), [microsoft/typeagent-py#234](https://github.com/microsoft/typeagent-py/pull/234).

## Repo Structure

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