From 9a3faebea281be1cfac25dd2a4f1b21a16296015 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Thu, 9 Apr 2026 22:45:02 -0500 Subject: [PATCH 01/20] Bump uv_build upper bound to <0.11.0 uv 0.10.x is current; the <0.10.0 constraint caused build warnings. --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 01470ecf..1339e34e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,5 +1,5 @@ [build-system] -requires = ["uv_build>=0.9.10,<0.10.0"] +requires = ["uv_build>=0.9.10,<0.11.0"] build-backend = "uv_build" [project] From ecbf6f5ee1bcaf9fd0ab555b4a6dd4b1eae3ffa4 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Thu, 9 Apr 2026 23:17:37 -0500 Subject: [PATCH 02/20] Defer black import to first use black is only used in create_context_prompt() and format_code() -- both cold paths. Moving the import inside the functions avoids loading black and its transitive deps (pathspec, black.nodes, etc.) on every import typeagent. --- src/typeagent/aitools/utils.py | 3 ++- src/typeagent/knowpro/answers.py | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/src/typeagent/aitools/utils.py b/src/typeagent/aitools/utils.py index adba401a..f3b49d65 100644 --- a/src/typeagent/aitools/utils.py +++ b/src/typeagent/aitools/utils.py @@ -11,7 +11,6 @@ import sys import time -import black import colorama import typechat @@ -57,6 +56,8 @@ def format_code(text: str, line_width=None) -> str: NOTE: The text must be a valid Python expression or code block. """ + import black + if line_width is None: # Use the terminal width, but cap it to 200 characters. line_width = min(200, shutil.get_terminal_size().columns) diff --git a/src/typeagent/knowpro/answers.py b/src/typeagent/knowpro/answers.py index 2c300506..58a536ed 100644 --- a/src/typeagent/knowpro/answers.py +++ b/src/typeagent/knowpro/answers.py @@ -5,8 +5,6 @@ from dataclasses import dataclass from typing import Any -import black - import typechat from .answer_context_schema import AnswerContext, RelevantKnowledge, RelevantMessage @@ -127,6 +125,8 @@ def create_question_prompt(question: str) -> str: def create_context_prompt(context: AnswerContext) -> str: # TODO: Use a more compact representation of the context than JSON. + import black + prompt = [ "[ANSWER CONTEXT]", "===", From d4bc7444c91f8bb1357fb069da3435ef510cddeb Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Fri, 10 Apr 2026 00:11:02 -0500 Subject: [PATCH 03/20] Batch schema DDL into executescript and pre-compile regex - Combine 16 separate cursor.execute() calls in init_db_schema into a single db.executescript() call, reducing SQLite round-trips during database initialization. - Pre-compile the whitespace regex in _prepare_term to avoid re-compiling on every call (552 calls during indexing). --- src/typeagent/storage/sqlite/schema.py | 42 ++++++++++----------- src/typeagent/storage/sqlite/semrefindex.py | 4 +- 2 files changed, 24 insertions(+), 22 deletions(-) diff --git a/src/typeagent/storage/sqlite/schema.py b/src/typeagent/storage/sqlite/schema.py index db6933db..2f50b7c6 100644 --- a/src/typeagent/storage/sqlite/schema.py +++ b/src/typeagent/storage/sqlite/schema.py @@ -256,29 +256,29 @@ def _set_conversation_metadata( ) +_ALL_SCHEMA_DDL = "\n".join([ + CONVERSATION_METADATA_SCHEMA, + MESSAGES_SCHEMA, + SEMANTIC_REFS_SCHEMA, + SEMANTIC_REF_INDEX_SCHEMA, + MESSAGE_TEXT_INDEX_SCHEMA, + PROPERTY_INDEX_SCHEMA, + RELATED_TERMS_ALIASES_SCHEMA, + RELATED_TERMS_FUZZY_SCHEMA, + TIMESTAMP_INDEX_SCHEMA, + INGESTED_SOURCES_SCHEMA, + SEMANTIC_REF_INDEX_TERM_INDEX, + MESSAGE_TEXT_INDEX_MESSAGE_INDEX, + MESSAGE_TEXT_INDEX_POSITION_INDEX, + RELATED_TERMS_ALIASES_TERM_INDEX, + RELATED_TERMS_ALIASES_ALIAS_INDEX, + RELATED_TERMS_FUZZY_TERM_INDEX, +]) + + def init_db_schema(db: sqlite3.Connection) -> None: """Initialize the database schema with all required tables.""" - cursor = db.cursor() - - # Create all tables - cursor.execute(CONVERSATION_METADATA_SCHEMA) - cursor.execute(MESSAGES_SCHEMA) - cursor.execute(SEMANTIC_REFS_SCHEMA) - cursor.execute(SEMANTIC_REF_INDEX_SCHEMA) - cursor.execute(MESSAGE_TEXT_INDEX_SCHEMA) - cursor.execute(PROPERTY_INDEX_SCHEMA) - cursor.execute(RELATED_TERMS_ALIASES_SCHEMA) - cursor.execute(RELATED_TERMS_FUZZY_SCHEMA) - cursor.execute(TIMESTAMP_INDEX_SCHEMA) - cursor.execute(INGESTED_SOURCES_SCHEMA) - - # Create additional indexes - cursor.execute(SEMANTIC_REF_INDEX_TERM_INDEX) - cursor.execute(MESSAGE_TEXT_INDEX_MESSAGE_INDEX) - cursor.execute(MESSAGE_TEXT_INDEX_POSITION_INDEX) - cursor.execute(RELATED_TERMS_ALIASES_TERM_INDEX) - cursor.execute(RELATED_TERMS_ALIASES_ALIAS_INDEX) - cursor.execute(RELATED_TERMS_FUZZY_TERM_INDEX) + db.executescript(_ALL_SCHEMA_DDL) def get_db_schema_version(db: sqlite3.Connection) -> int: diff --git a/src/typeagent/storage/sqlite/semrefindex.py b/src/typeagent/storage/sqlite/semrefindex.py index 682b8e7d..96a7d565 100644 --- a/src/typeagent/storage/sqlite/semrefindex.py +++ b/src/typeagent/storage/sqlite/semrefindex.py @@ -10,6 +10,8 @@ from ...knowpro import interfaces from ...knowpro.interfaces import ScoredSemanticRefOrdinal +_WHITESPACE_RE = re.compile(r"\s+") + class SqliteTermToSemanticRefIndex(interfaces.ITermToSemanticRefIndex): """SQLite-backed implementation of term to semantic ref index.""" @@ -150,7 +152,7 @@ def _prepare_term(self, term: str) -> str: term = unicodedata.normalize("NFC", term) # Collapse multiple whitespace characters to single space - term = re.sub(r"\s+", " ", term) + term = _WHITESPACE_RE.sub(" ", term) # Convert to lowercase return term.lower() From 88c8d5e1dc411ad10b0b96174cd4c7e1a4e8aad5 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Fri, 10 Apr 2026 00:15:37 -0500 Subject: [PATCH 04/20] Revert "Batch schema DDL into executescript and pre-compile regex" This reverts commit d4bc7444c91f8bb1357fb069da3435ef510cddeb. --- src/typeagent/storage/sqlite/schema.py | 42 ++++++++++----------- src/typeagent/storage/sqlite/semrefindex.py | 4 +- 2 files changed, 22 insertions(+), 24 deletions(-) diff --git a/src/typeagent/storage/sqlite/schema.py b/src/typeagent/storage/sqlite/schema.py index 2f50b7c6..db6933db 100644 --- a/src/typeagent/storage/sqlite/schema.py +++ b/src/typeagent/storage/sqlite/schema.py @@ -256,29 +256,29 @@ def _set_conversation_metadata( ) -_ALL_SCHEMA_DDL = "\n".join([ - CONVERSATION_METADATA_SCHEMA, - MESSAGES_SCHEMA, - SEMANTIC_REFS_SCHEMA, - SEMANTIC_REF_INDEX_SCHEMA, - MESSAGE_TEXT_INDEX_SCHEMA, - PROPERTY_INDEX_SCHEMA, - RELATED_TERMS_ALIASES_SCHEMA, - RELATED_TERMS_FUZZY_SCHEMA, - TIMESTAMP_INDEX_SCHEMA, - INGESTED_SOURCES_SCHEMA, - SEMANTIC_REF_INDEX_TERM_INDEX, - MESSAGE_TEXT_INDEX_MESSAGE_INDEX, - MESSAGE_TEXT_INDEX_POSITION_INDEX, - RELATED_TERMS_ALIASES_TERM_INDEX, - RELATED_TERMS_ALIASES_ALIAS_INDEX, - RELATED_TERMS_FUZZY_TERM_INDEX, -]) - - def init_db_schema(db: sqlite3.Connection) -> None: """Initialize the database schema with all required tables.""" - db.executescript(_ALL_SCHEMA_DDL) + cursor = db.cursor() + + # Create all tables + cursor.execute(CONVERSATION_METADATA_SCHEMA) + cursor.execute(MESSAGES_SCHEMA) + cursor.execute(SEMANTIC_REFS_SCHEMA) + cursor.execute(SEMANTIC_REF_INDEX_SCHEMA) + cursor.execute(MESSAGE_TEXT_INDEX_SCHEMA) + cursor.execute(PROPERTY_INDEX_SCHEMA) + cursor.execute(RELATED_TERMS_ALIASES_SCHEMA) + cursor.execute(RELATED_TERMS_FUZZY_SCHEMA) + cursor.execute(TIMESTAMP_INDEX_SCHEMA) + cursor.execute(INGESTED_SOURCES_SCHEMA) + + # Create additional indexes + cursor.execute(SEMANTIC_REF_INDEX_TERM_INDEX) + cursor.execute(MESSAGE_TEXT_INDEX_MESSAGE_INDEX) + cursor.execute(MESSAGE_TEXT_INDEX_POSITION_INDEX) + cursor.execute(RELATED_TERMS_ALIASES_TERM_INDEX) + cursor.execute(RELATED_TERMS_ALIASES_ALIAS_INDEX) + cursor.execute(RELATED_TERMS_FUZZY_TERM_INDEX) def get_db_schema_version(db: sqlite3.Connection) -> int: diff --git a/src/typeagent/storage/sqlite/semrefindex.py b/src/typeagent/storage/sqlite/semrefindex.py index 96a7d565..682b8e7d 100644 --- a/src/typeagent/storage/sqlite/semrefindex.py +++ b/src/typeagent/storage/sqlite/semrefindex.py @@ -10,8 +10,6 @@ from ...knowpro import interfaces from ...knowpro.interfaces import ScoredSemanticRefOrdinal -_WHITESPACE_RE = re.compile(r"\s+") - class SqliteTermToSemanticRefIndex(interfaces.ITermToSemanticRefIndex): """SQLite-backed implementation of term to semantic ref index.""" @@ -152,7 +150,7 @@ def _prepare_term(self, term: str) -> str: term = unicodedata.normalize("NFC", term) # Collapse multiple whitespace characters to single space - term = _WHITESPACE_RE.sub(" ", term) + term = re.sub(r"\s+", " ", term) # Convert to lowercase return term.lower() From bc9f2df440cc735f53657b0b3552d62f28e6f1a1 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Fri, 10 Apr 2026 00:20:51 -0500 Subject: [PATCH 05/20] Batch SQLite INSERTs for semantic ref and property indexing Add add_terms_batch / add_properties_batch to the index interfaces with executemany-based SQLite implementations. Restructure add_metadata_to_index_from_list and add_to_property_index to collect all items first, then batch-insert via extend() and the new batch methods. Eliminates ~1000 individual INSERT round-trips during indexing. --- src/typeagent/knowpro/interfaces_core.py | 5 ++ src/typeagent/knowpro/interfaces_indexes.py | 5 ++ src/typeagent/storage/memory/propindex.py | 75 ++++++++++++++++-- src/typeagent/storage/memory/semrefindex.py | 87 ++++++++++++++++++++- src/typeagent/storage/sqlite/propindex.py | 30 +++++++ src/typeagent/storage/sqlite/semrefindex.py | 23 ++++++ 6 files changed, 215 insertions(+), 10 deletions(-) diff --git a/src/typeagent/knowpro/interfaces_core.py b/src/typeagent/knowpro/interfaces_core.py index 105e45b6..72f11f8c 100644 --- a/src/typeagent/knowpro/interfaces_core.py +++ b/src/typeagent/knowpro/interfaces_core.py @@ -168,6 +168,11 @@ async def add_term( semantic_ref_ordinal: SemanticRefOrdinal | ScoredSemanticRefOrdinal, ) -> str: ... + async def add_terms_batch( + self, + terms: list[tuple[str, SemanticRefOrdinal | ScoredSemanticRefOrdinal]], + ) -> None: ... + async def remove_term( self, term: str, semantic_ref_ordinal: SemanticRefOrdinal ) -> None: ... diff --git a/src/typeagent/knowpro/interfaces_indexes.py b/src/typeagent/knowpro/interfaces_indexes.py index a894ab88..3ae62024 100644 --- a/src/typeagent/knowpro/interfaces_indexes.py +++ b/src/typeagent/knowpro/interfaces_indexes.py @@ -59,6 +59,11 @@ async def add_property( semantic_ref_ordinal: SemanticRefOrdinal | ScoredSemanticRefOrdinal, ) -> None: ... + async def add_properties_batch( + self, + properties: list[tuple[str, str, SemanticRefOrdinal | ScoredSemanticRefOrdinal]], + ) -> None: ... + async def lookup_property( self, property_name: str, value: str ) -> list[ScoredSemanticRefOrdinal] | None: ... diff --git a/src/typeagent/storage/memory/propindex.py b/src/typeagent/storage/memory/propindex.py index acc7b89a..1a1c7968 100644 --- a/src/typeagent/storage/memory/propindex.py +++ b/src/typeagent/storage/memory/propindex.py @@ -109,6 +109,57 @@ async def build_property_index(conversation: IConversation) -> None: await add_to_property_index(conversation, 0) +def _collect_facet_properties( + facet: kplib.Facet | None, + ordinal: SemanticRefOrdinal, +) -> list[tuple[str, str, SemanticRefOrdinal]]: + """Collect property tuples from a facet without touching any index.""" + if facet is None: + return [] + props: list[tuple[str, str, SemanticRefOrdinal]] = [ + (PropertyNames.FacetName.value, facet.name, ordinal) + ] + value = facet.value + if value is not None: + if isinstance(value, float) and value: + value = f"{value:g}" + props.append((PropertyNames.FacetValue.value, str(value), ordinal)) + return props + + +def _collect_entity_properties( + entity: kplib.ConcreteEntity, + ordinal: SemanticRefOrdinal, +) -> list[tuple[str, str, SemanticRefOrdinal]]: + """Collect all property tuples for an entity.""" + props: list[tuple[str, str, SemanticRefOrdinal]] = [ + (PropertyNames.EntityName.value, entity.name, ordinal) + ] + for t in entity.type: + props.append((PropertyNames.EntityType.value, t, ordinal)) + if entity.facets: + for facet in entity.facets: + props.extend(_collect_facet_properties(facet, ordinal)) + return props + + +def _collect_action_properties( + action: kplib.Action, + ordinal: SemanticRefOrdinal, +) -> list[tuple[str, str, SemanticRefOrdinal]]: + """Collect all property tuples for an action.""" + props: list[tuple[str, str, SemanticRefOrdinal]] = [ + (PropertyNames.Verb.value, " ".join(action.verbs), ordinal) + ] + if action.subject_entity_name != "none": + props.append((PropertyNames.Subject.value, action.subject_entity_name, ordinal)) + if action.object_entity_name != "none": + props.append((PropertyNames.Object.value, action.object_entity_name, ordinal)) + if action.indirect_object_entity_name != "none": + props.append((PropertyNames.IndirectObject.value, action.indirect_object_entity_name, ordinal)) + return props + + async def add_to_property_index( conversation: IConversation, start_at_ordinal: SemanticRefOrdinal, @@ -127,29 +178,32 @@ async def add_to_property_index( semantic_refs = conversation.semantic_refs size = await semantic_refs.size() + collected: list[tuple[str, str, SemanticRefOrdinal]] = [] for semantic_ref_ordinal, semantic_ref in enumerate( await semantic_refs.get_slice(start_at_ordinal, size), start_at_ordinal, ): assert semantic_ref.semantic_ref_ordinal == semantic_ref_ordinal if isinstance(semantic_ref.knowledge, kplib.Action): - await add_action_properties_to_index( - semantic_ref.knowledge, property_index, semantic_ref_ordinal + collected.extend( + _collect_action_properties(semantic_ref.knowledge, semantic_ref_ordinal) ) elif isinstance(semantic_ref.knowledge, kplib.ConcreteEntity): - await add_entity_properties_to_index( - semantic_ref.knowledge, property_index, semantic_ref_ordinal + collected.extend( + _collect_entity_properties(semantic_ref.knowledge, semantic_ref_ordinal) ) elif isinstance(semantic_ref.knowledge, Tag): - tag = semantic_ref.knowledge - await property_index.add_property( - PropertyNames.Tag.value, tag.text, semantic_ref_ordinal + collected.append( + (PropertyNames.Tag.value, semantic_ref.knowledge.text, semantic_ref_ordinal) ) elif isinstance(semantic_ref.knowledge, Topic): pass else: assert_never(semantic_ref.knowledge) + if collected: + await property_index.add_properties_batch(collected) + class PropertyIndex(IPropertyToSemanticRefIndex): def __init__(self): @@ -183,6 +237,13 @@ async def add_property( else: self._map[term_text] = [semantic_ref_ordinal] + async def add_properties_batch( + self, + properties: list[tuple[str, str, SemanticRefOrdinal | ScoredSemanticRefOrdinal]], + ) -> None: + for name, value, ordinal in properties: + await self.add_property(name, value, ordinal) + async def clear(self) -> None: self._map = {} diff --git a/src/typeagent/storage/memory/semrefindex.py b/src/typeagent/storage/memory/semrefindex.py index 6c42022d..7f131d9a 100644 --- a/src/typeagent/storage/memory/semrefindex.py +++ b/src/typeagent/storage/memory/semrefindex.py @@ -577,6 +577,48 @@ async def add_metadata_to_index[TMessage: IMessage]( i += 1 +def _collect_facet_terms(facet: kplib.Facet | None) -> list[str]: + """Collect terms from a facet without touching any index.""" + if facet is None: + return [] + terms = [facet.name] + if facet.value is not None: + terms.append(str(facet.value)) + return terms + + +def _collect_entity_terms(entity: kplib.ConcreteEntity) -> list[str]: + """Collect all terms an entity would add to the semantic ref index.""" + terms = [entity.name] + for t in entity.type: + terms.append(t) + if entity.facets: + for facet in entity.facets: + terms.extend(_collect_facet_terms(facet)) + return terms + + +def _collect_action_terms(action: kplib.Action) -> list[str]: + """Collect all terms an action would add to the semantic ref index.""" + terms = [" ".join(action.verbs)] + if action.subject_entity_name != "none": + terms.append(action.subject_entity_name) + if action.object_entity_name != "none": + terms.append(action.object_entity_name) + if action.indirect_object_entity_name != "none": + terms.append(action.indirect_object_entity_name) + if action.params: + for param in action.params: + if isinstance(param, str): + terms.append(param) + else: + terms.append(param.name) + if isinstance(param.value, str): + terms.append(param.value) + terms.extend(_collect_facet_terms(action.subject_entity_facet)) + return terms + + async def add_metadata_to_index_from_list[TMessage: IMessage]( messages: list[TMessage], semantic_refs: ISemanticRefCollection, @@ -585,18 +627,50 @@ async def add_metadata_to_index_from_list[TMessage: IMessage]( knowledge_validator: KnowledgeValidator | None = None, ) -> None: """Extract metadata knowledge from a list of messages starting at ordinal.""" + next_ordinal = await semantic_refs.size() + collected_refs: list[SemanticRef] = [] + collected_terms: list[tuple[str, SemanticRefOrdinal]] = [] + for i, msg in enumerate(messages, start_from_ordinal): knowledge_response = msg.get_knowledge() for entity in knowledge_response.entities: if knowledge_validator is None or knowledge_validator("entity", entity): - await add_entity_to_index(entity, semantic_refs, semantic_ref_index, i) + ref = SemanticRef( + semantic_ref_ordinal=next_ordinal, + range=text_range_from_location(i), + knowledge=entity, + ) + collected_refs.append(ref) + for term in _collect_entity_terms(entity): + collected_terms.append((term, next_ordinal)) + next_ordinal += 1 for action in knowledge_response.actions: if knowledge_validator is None or knowledge_validator("action", action): - await add_action_to_index(action, semantic_refs, semantic_ref_index, i) + ref = SemanticRef( + semantic_ref_ordinal=next_ordinal, + range=text_range_from_location(i), + knowledge=action, + ) + collected_refs.append(ref) + for term in _collect_action_terms(action): + collected_terms.append((term, next_ordinal)) + next_ordinal += 1 for topic_response in knowledge_response.topics: topic = Topic(text=topic_response) if knowledge_validator is None or knowledge_validator("topic", topic): - await add_topic_to_index(topic, semantic_refs, semantic_ref_index, i) + ref = SemanticRef( + semantic_ref_ordinal=next_ordinal, + range=text_range_from_location(i), + knowledge=topic, + ) + collected_refs.append(ref) + collected_terms.append((topic.text, next_ordinal)) + next_ordinal += 1 + + if collected_refs: + await semantic_refs.extend(collected_refs) + if collected_terms: + await semantic_ref_index.add_terms_batch(collected_terms) class TermToSemanticRefIndex(ITermToSemanticRefIndex): @@ -635,6 +709,13 @@ async def add_term( self._map[term] = [semantic_ref_ordinal] return term + async def add_terms_batch( + self, + terms: list[tuple[str, SemanticRefOrdinal | ScoredSemanticRefOrdinal]], + ) -> None: + for term, ordinal in terms: + await self.add_term(term, ordinal) + async def lookup_term(self, term: str) -> list[ScoredSemanticRefOrdinal] | None: return self._map.get(self._prepare_term(term)) or [] diff --git a/src/typeagent/storage/sqlite/propindex.py b/src/typeagent/storage/sqlite/propindex.py index 5a0fa63a..6a619869 100644 --- a/src/typeagent/storage/sqlite/propindex.py +++ b/src/typeagent/storage/sqlite/propindex.py @@ -67,6 +67,36 @@ async def add_property( (property_name, value, score, semref_id), ) + async def add_properties_batch( + self, + properties: list[tuple[str, str, interfaces.SemanticRefOrdinal | interfaces.ScoredSemanticRefOrdinal]], + ) -> None: + if not properties: + return + from ...storage.memory.propindex import ( + make_property_term_text, + split_property_term_text, + ) + rows = [] + for property_name, value, ordinal in properties: + if isinstance(ordinal, interfaces.ScoredSemanticRefOrdinal): + semref_id = ordinal.semantic_ref_ordinal + score = ordinal.score + else: + semref_id = ordinal + score = 1.0 + term_text = make_property_term_text(property_name, value) + term_text = term_text.lower() + property_name, value = split_property_term_text(term_text) + if property_name.startswith("prop."): + property_name = property_name[5:] + rows.append((property_name, value, score, semref_id)) + cursor = self.db.cursor() + cursor.executemany( + "INSERT INTO PropertyIndex (prop_name, value_str, score, semref_id) VALUES (?, ?, ?, ?)", + rows, + ) + async def clear(self) -> None: cursor = self.db.cursor() cursor.execute("DELETE FROM PropertyIndex") diff --git a/src/typeagent/storage/sqlite/semrefindex.py b/src/typeagent/storage/sqlite/semrefindex.py index 682b8e7d..0925eb45 100644 --- a/src/typeagent/storage/sqlite/semrefindex.py +++ b/src/typeagent/storage/sqlite/semrefindex.py @@ -56,6 +56,29 @@ async def add_term( return term + async def add_terms_batch( + self, + terms: list[tuple[str, interfaces.SemanticRefOrdinal | interfaces.ScoredSemanticRefOrdinal]], + ) -> None: + if not terms: + return + rows = [] + for term, ordinal in terms: + if not term: + continue + term = self._prepare_term(term) + if isinstance(ordinal, interfaces.ScoredSemanticRefOrdinal): + semref_id = ordinal.semantic_ref_ordinal + else: + semref_id = ordinal + rows.append((term, semref_id)) + if rows: + cursor = self.db.cursor() + cursor.executemany( + "INSERT OR IGNORE INTO SemanticRefIndex (term, semref_id) VALUES (?, ?)", + rows, + ) + async def remove_term( self, term: str, semantic_ref_ordinal: interfaces.SemanticRefOrdinal ) -> None: From 2140df4f363866448559c5a5707be923ea4123d8 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Fri, 10 Apr 2026 00:35:14 -0500 Subject: [PATCH 06/20] Add pytest-async-benchmark tests for indexing pipeline --- tests/benchmarks/test_benchmark_indexing.py | 260 ++++++++++++++++++++ 1 file changed, 260 insertions(+) create mode 100644 tests/benchmarks/test_benchmark_indexing.py diff --git a/tests/benchmarks/test_benchmark_indexing.py b/tests/benchmarks/test_benchmark_indexing.py new file mode 100644 index 00000000..274dd91d --- /dev/null +++ b/tests/benchmarks/test_benchmark_indexing.py @@ -0,0 +1,260 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Benchmarks for add_messages_with_indexing — the core indexing pipeline. + +Exercises: message storage, semantic ref creation, term index insertion, +property index insertion, and embedding computation. + +Only the hot path (add_messages_with_indexing) is timed — DB creation, +storage provider init, VTT parsing, and teardown are excluded. + +Run: + uv run python -m pytest tests/benchmarks/test_benchmark_indexing.py -v -s +""" + +import os +import shutil +import statistics +import tempfile +import time +from datetime import timedelta + +import pytest +import webvtt + +from typeagent.aitools.model_adapters import create_test_embedding_model +from typeagent.knowpro.convsettings import ConversationSettings +from typeagent.knowpro.universal_message import ( + UNIX_EPOCH, + ConversationMessage, + ConversationMessageMeta, + format_timestamp_utc, +) +from typeagent.storage.sqlite.provider import SqliteStorageProvider +from typeagent.transcripts.transcript import ( + Transcript, + TranscriptMessage, + TranscriptMessageMeta, +) +from typeagent.transcripts.transcript_ingest import ( + parse_voice_tags, + webvtt_timestamp_to_seconds, +) + +TESTDATA = os.path.join(os.path.dirname(__file__), "..", "testdata") +CONFUSE_A_CAT_VTT = os.path.join(TESTDATA, "Confuse-A-Cat.vtt") + +WARMUP = 3 +ROUNDS = 20 + + +def _make_settings(): + """Create conversation settings with fake embedding model (no API keys).""" + model = create_test_embedding_model() + settings = ConversationSettings(model=model) + settings.semantic_ref_index_settings.auto_extract_knowledge = False + return settings + + +def _parse_vtt(vtt_file_path: str) -> list[ConversationMessage]: + """Parse a VTT file into messages (mirrors ingest_vtt_transcript parsing).""" + vtt = webvtt.read(vtt_file_path) + base_date = UNIX_EPOCH + + messages: list[ConversationMessage] = [] + current_speaker = None + current_text_chunks: list[str] = [] + current_start_time = None + + for caption in vtt: + if not caption.text.strip(): + continue + raw_text = getattr(caption, "raw_text", caption.text) + voice_segments = parse_voice_tags(raw_text) + start_time = caption.start + + for speaker, text in voice_segments: + if not text.strip(): + continue + if speaker == current_speaker and current_text_chunks: + current_text_chunks.append(text) + else: + if current_text_chunks and current_start_time is not None: + combined_text = " ".join(current_text_chunks).strip() + if combined_text: + offset_seconds = webvtt_timestamp_to_seconds( + current_start_time + ) + timestamp = format_timestamp_utc( + base_date + timedelta(seconds=offset_seconds) + ) + messages.append( + ConversationMessage( + text_chunks=[combined_text], + metadata=ConversationMessageMeta( + speaker=current_speaker, recipients=[] + ), + timestamp=timestamp, + ) + ) + current_speaker = speaker + current_text_chunks = [text] if text.strip() else [] + current_start_time = start_time + + if current_text_chunks and current_start_time is not None: + combined_text = " ".join(current_text_chunks).strip() + if combined_text: + offset_seconds = webvtt_timestamp_to_seconds(current_start_time) + timestamp = format_timestamp_utc( + base_date + timedelta(seconds=offset_seconds) + ) + messages.append( + ConversationMessage( + text_chunks=[combined_text], + metadata=ConversationMessageMeta( + speaker=current_speaker, recipients=[] + ), + timestamp=timestamp, + ) + ) + + return messages + + +def _report(name: str, times: list[float]) -> None: + """Print benchmark stats to stdout.""" + mn = min(times) + md = statistics.median(times) + avg = statistics.mean(times) + sd = statistics.stdev(times) if len(times) > 1 else 0.0 + print(f"\n{'=' * 60}") + print(f" {name} ({ROUNDS} rounds, {WARMUP} warmup)") + print(f" min = {mn * 1000:.3f} ms") + print(f" median = {md * 1000:.3f} ms") + print(f" mean = {avg * 1000:.3f} ms") + print(f" stddev = {sd * 1000:.3f} ms") + print(f"{'=' * 60}\n") + + +@pytest.mark.asyncio +async def test_benchmark_vtt_ingest(): + """Benchmark indexing of pre-parsed VTT messages.""" + vtt_messages = _parse_vtt(CONFUSE_A_CAT_VTT) + settings = _make_settings() + tmpdir = tempfile.mkdtemp() + times: list[float] = [] + + try: + for i in range(WARMUP + ROUNDS): + db_path = os.path.join(tmpdir, f"bench_{i}.db") + storage = SqliteStorageProvider( + db_path, + message_type=ConversationMessage, + message_text_index_settings=settings.message_text_index_settings, + related_term_index_settings=settings.related_term_index_settings, + ) + settings.storage_provider = storage + transcript = await Transcript.create( + settings, + name="bench", + tags=["bench", "vtt-transcript"], + ) + + start = time.perf_counter() + await transcript.add_messages_with_indexing(vtt_messages) + elapsed = time.perf_counter() - start + + await storage.close() + os.remove(db_path) + + if i >= WARMUP: + times.append(elapsed) + + _report(f"VTT ingest ({len(vtt_messages)} msgs)", times) + finally: + shutil.rmtree(tmpdir, ignore_errors=True) + + +@pytest.mark.asyncio +async def test_benchmark_add_messages_50(): + """Benchmark add_messages_with_indexing with 50 messages.""" + messages = [ + TranscriptMessage( + text_chunks=[f"Message {i} about topic {i % 10}"], + metadata=TranscriptMessageMeta(speaker=f"Speaker{i % 3}"), + tags=[f"tag{i % 5}"], + ) + for i in range(50) + ] + settings = _make_settings() + tmpdir = tempfile.mkdtemp() + times: list[float] = [] + + try: + for i in range(WARMUP + ROUNDS): + db_path = os.path.join(tmpdir, f"bench_{i}.db") + storage = SqliteStorageProvider( + db_path, + message_type=TranscriptMessage, + message_text_index_settings=settings.message_text_index_settings, + related_term_index_settings=settings.related_term_index_settings, + ) + settings.storage_provider = storage + transcript = await Transcript.create(settings, name="bench") + + start = time.perf_counter() + await transcript.add_messages_with_indexing(messages) + elapsed = time.perf_counter() - start + + await storage.close() + os.remove(db_path) + + if i >= WARMUP: + times.append(elapsed) + + _report("add_messages (50)", times) + finally: + shutil.rmtree(tmpdir, ignore_errors=True) + + +@pytest.mark.asyncio +async def test_benchmark_add_messages_200(): + """Benchmark add_messages_with_indexing with 200 messages.""" + messages = [ + TranscriptMessage( + text_chunks=[f"Message {i} about topic {i % 10}"], + metadata=TranscriptMessageMeta(speaker=f"Speaker{i % 3}"), + tags=[f"tag{i % 5}"], + ) + for i in range(200) + ] + settings = _make_settings() + tmpdir = tempfile.mkdtemp() + times: list[float] = [] + + try: + for i in range(WARMUP + ROUNDS): + db_path = os.path.join(tmpdir, f"bench_{i}.db") + storage = SqliteStorageProvider( + db_path, + message_type=TranscriptMessage, + message_text_index_settings=settings.message_text_index_settings, + related_term_index_settings=settings.related_term_index_settings, + ) + settings.storage_provider = storage + transcript = await Transcript.create(settings, name="bench") + + start = time.perf_counter() + await transcript.add_messages_with_indexing(messages) + elapsed = time.perf_counter() - start + + await storage.close() + os.remove(db_path) + + if i >= WARMUP: + times.append(elapsed) + + _report("add_messages (200)", times) + finally: + shutil.rmtree(tmpdir, ignore_errors=True) From d0d070b998b97b570966426e3e79227661c6d1cc Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Fri, 10 Apr 2026 01:19:07 -0500 Subject: [PATCH 07/20] Rewrite benchmarks to use async_benchmark.pedantic() Replace hand-rolled time.perf_counter() loop with the pedantic fixture from pytest-async-benchmark. Setup (DB/storage/transcript creation) and teardown (close/delete) are now properly excluded from timing via the framework instead of inline timing code. --- tests/benchmarks/test_benchmark_indexing.py | 313 ++++++++------------ 1 file changed, 121 insertions(+), 192 deletions(-) diff --git a/tests/benchmarks/test_benchmark_indexing.py b/tests/benchmarks/test_benchmark_indexing.py index 274dd91d..d25c4b6e 100644 --- a/tests/benchmarks/test_benchmark_indexing.py +++ b/tests/benchmarks/test_benchmark_indexing.py @@ -7,49 +7,41 @@ property index insertion, and embedding computation. Only the hot path (add_messages_with_indexing) is timed — DB creation, -storage provider init, VTT parsing, and teardown are excluded. +storage provider init, VTT parsing, and teardown are excluded via +async_benchmark.pedantic(). Run: uv run python -m pytest tests/benchmarks/test_benchmark_indexing.py -v -s """ +import itertools import os import shutil -import statistics import tempfile -import time -from datetime import timedelta import pytest -import webvtt from typeagent.aitools.model_adapters import create_test_embedding_model from typeagent.knowpro.convsettings import ConversationSettings -from typeagent.knowpro.universal_message import ( - UNIX_EPOCH, - ConversationMessage, - ConversationMessageMeta, - format_timestamp_utc, -) +from typeagent.knowpro.universal_message import ConversationMessage from typeagent.storage.sqlite.provider import SqliteStorageProvider from typeagent.transcripts.transcript import ( Transcript, TranscriptMessage, TranscriptMessageMeta, ) -from typeagent.transcripts.transcript_ingest import ( - parse_voice_tags, - webvtt_timestamp_to_seconds, -) +from typeagent.transcripts.transcript_ingest import ingest_vtt_transcript TESTDATA = os.path.join(os.path.dirname(__file__), "..", "testdata") CONFUSE_A_CAT_VTT = os.path.join(TESTDATA, "Confuse-A-Cat.vtt") -WARMUP = 3 -ROUNDS = 20 + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- -def _make_settings(): +def make_settings() -> ConversationSettings: """Create conversation settings with fake embedding model (no API keys).""" model = create_test_embedding_model() settings = ConversationSettings(model=model) @@ -57,204 +49,141 @@ def _make_settings(): return settings -def _parse_vtt(vtt_file_path: str) -> list[ConversationMessage]: - """Parse a VTT file into messages (mirrors ingest_vtt_transcript parsing).""" - vtt = webvtt.read(vtt_file_path) - base_date = UNIX_EPOCH - - messages: list[ConversationMessage] = [] - current_speaker = None - current_text_chunks: list[str] = [] - current_start_time = None - - for caption in vtt: - if not caption.text.strip(): - continue - raw_text = getattr(caption, "raw_text", caption.text) - voice_segments = parse_voice_tags(raw_text) - start_time = caption.start - - for speaker, text in voice_segments: - if not text.strip(): - continue - if speaker == current_speaker and current_text_chunks: - current_text_chunks.append(text) - else: - if current_text_chunks and current_start_time is not None: - combined_text = " ".join(current_text_chunks).strip() - if combined_text: - offset_seconds = webvtt_timestamp_to_seconds( - current_start_time - ) - timestamp = format_timestamp_utc( - base_date + timedelta(seconds=offset_seconds) - ) - messages.append( - ConversationMessage( - text_chunks=[combined_text], - metadata=ConversationMessageMeta( - speaker=current_speaker, recipients=[] - ), - timestamp=timestamp, - ) - ) - current_speaker = speaker - current_text_chunks = [text] if text.strip() else [] - current_start_time = start_time - - if current_text_chunks and current_start_time is not None: - combined_text = " ".join(current_text_chunks).strip() - if combined_text: - offset_seconds = webvtt_timestamp_to_seconds(current_start_time) - timestamp = format_timestamp_utc( - base_date + timedelta(seconds=offset_seconds) - ) - messages.append( - ConversationMessage( - text_chunks=[combined_text], - metadata=ConversationMessageMeta( - speaker=current_speaker, recipients=[] - ), - timestamp=timestamp, - ) - ) - +async def extract_vtt_messages(vtt_path: str) -> list[ConversationMessage]: + """Parse a VTT file via ingest_vtt_transcript and return the messages.""" + settings = make_settings() + with tempfile.TemporaryDirectory() as tmpdir: + db_path = os.path.join(tmpdir, "parse.db") + transcript = await ingest_vtt_transcript(vtt_path, settings, dbname=db_path) + n = await transcript.messages.size() + messages = await transcript.messages.get_slice(0, n) + await settings.storage_provider.close() return messages -def _report(name: str, times: list[float]) -> None: - """Print benchmark stats to stdout.""" - mn = min(times) - md = statistics.median(times) - avg = statistics.mean(times) - sd = statistics.stdev(times) if len(times) > 1 else 0.0 - print(f"\n{'=' * 60}") - print(f" {name} ({ROUNDS} rounds, {WARMUP} warmup)") - print(f" min = {mn * 1000:.3f} ms") - print(f" median = {md * 1000:.3f} ms") - print(f" mean = {avg * 1000:.3f} ms") - print(f" stddev = {sd * 1000:.3f} ms") - print(f"{'=' * 60}\n") +def synthetic_messages(n: int) -> list[TranscriptMessage]: + """Build n synthetic TranscriptMessages.""" + return [ + TranscriptMessage( + text_chunks=[f"Message {i} about topic {i % 10}"], + metadata=TranscriptMessageMeta(speaker=f"Speaker{i % 3}"), + tags=[f"tag{i % 5}"], + ) + for i in range(n) + ] + + +# --------------------------------------------------------------------------- +# Test cases +# --------------------------------------------------------------------------- @pytest.mark.asyncio -async def test_benchmark_vtt_ingest(): - """Benchmark indexing of pre-parsed VTT messages.""" - vtt_messages = _parse_vtt(CONFUSE_A_CAT_VTT) - settings = _make_settings() +async def test_benchmark_vtt_ingest(async_benchmark): + """Benchmark indexing of pre-parsed VTT messages (Confuse-A-Cat, 40 msgs).""" + messages = await extract_vtt_messages(CONFUSE_A_CAT_VTT) + settings = make_settings() tmpdir = tempfile.mkdtemp() - times: list[float] = [] + counter = itertools.count() + + async def setup(): + i = next(counter) + db_path = os.path.join(tmpdir, f"bench_{i}.db") + storage = SqliteStorageProvider( + db_path, + message_type=ConversationMessage, + message_text_index_settings=settings.message_text_index_settings, + related_term_index_settings=settings.related_term_index_settings, + ) + settings.storage_provider = storage + transcript = await Transcript.create(settings, name="bench") + return transcript, storage, db_path + + async def teardown(setup_rv): + _, storage, db_path = setup_rv + await storage.close() + os.remove(db_path) + + async def target(transcript, storage, db_path): + await transcript.add_messages_with_indexing(messages) try: - for i in range(WARMUP + ROUNDS): - db_path = os.path.join(tmpdir, f"bench_{i}.db") - storage = SqliteStorageProvider( - db_path, - message_type=ConversationMessage, - message_text_index_settings=settings.message_text_index_settings, - related_term_index_settings=settings.related_term_index_settings, - ) - settings.storage_provider = storage - transcript = await Transcript.create( - settings, - name="bench", - tags=["bench", "vtt-transcript"], - ) - - start = time.perf_counter() - await transcript.add_messages_with_indexing(vtt_messages) - elapsed = time.perf_counter() - start - - await storage.close() - os.remove(db_path) - - if i >= WARMUP: - times.append(elapsed) - - _report(f"VTT ingest ({len(vtt_messages)} msgs)", times) + await async_benchmark.pedantic( + target, setup=setup, teardown=teardown, rounds=20, warmup_rounds=3 + ) finally: shutil.rmtree(tmpdir, ignore_errors=True) @pytest.mark.asyncio -async def test_benchmark_add_messages_50(): - """Benchmark add_messages_with_indexing with 50 messages.""" - messages = [ - TranscriptMessage( - text_chunks=[f"Message {i} about topic {i % 10}"], - metadata=TranscriptMessageMeta(speaker=f"Speaker{i % 3}"), - tags=[f"tag{i % 5}"], - ) - for i in range(50) - ] - settings = _make_settings() +async def test_benchmark_add_messages_50(async_benchmark): + """Benchmark add_messages_with_indexing with 50 synthetic messages.""" + messages = synthetic_messages(50) + settings = make_settings() tmpdir = tempfile.mkdtemp() - times: list[float] = [] + counter = itertools.count() + + async def setup(): + i = next(counter) + db_path = os.path.join(tmpdir, f"bench_{i}.db") + storage = SqliteStorageProvider( + db_path, + message_type=TranscriptMessage, + message_text_index_settings=settings.message_text_index_settings, + related_term_index_settings=settings.related_term_index_settings, + ) + settings.storage_provider = storage + transcript = await Transcript.create(settings, name="bench") + return transcript, storage, db_path + + async def teardown(setup_rv): + _, storage, db_path = setup_rv + await storage.close() + os.remove(db_path) + + async def target(transcript, storage, db_path): + await transcript.add_messages_with_indexing(messages) try: - for i in range(WARMUP + ROUNDS): - db_path = os.path.join(tmpdir, f"bench_{i}.db") - storage = SqliteStorageProvider( - db_path, - message_type=TranscriptMessage, - message_text_index_settings=settings.message_text_index_settings, - related_term_index_settings=settings.related_term_index_settings, - ) - settings.storage_provider = storage - transcript = await Transcript.create(settings, name="bench") - - start = time.perf_counter() - await transcript.add_messages_with_indexing(messages) - elapsed = time.perf_counter() - start - - await storage.close() - os.remove(db_path) - - if i >= WARMUP: - times.append(elapsed) - - _report("add_messages (50)", times) + await async_benchmark.pedantic( + target, setup=setup, teardown=teardown, rounds=20, warmup_rounds=3 + ) finally: shutil.rmtree(tmpdir, ignore_errors=True) @pytest.mark.asyncio -async def test_benchmark_add_messages_200(): - """Benchmark add_messages_with_indexing with 200 messages.""" - messages = [ - TranscriptMessage( - text_chunks=[f"Message {i} about topic {i % 10}"], - metadata=TranscriptMessageMeta(speaker=f"Speaker{i % 3}"), - tags=[f"tag{i % 5}"], - ) - for i in range(200) - ] - settings = _make_settings() +async def test_benchmark_add_messages_200(async_benchmark): + """Benchmark add_messages_with_indexing with 200 synthetic messages.""" + messages = synthetic_messages(200) + settings = make_settings() tmpdir = tempfile.mkdtemp() - times: list[float] = [] + counter = itertools.count() + + async def setup(): + i = next(counter) + db_path = os.path.join(tmpdir, f"bench_{i}.db") + storage = SqliteStorageProvider( + db_path, + message_type=TranscriptMessage, + message_text_index_settings=settings.message_text_index_settings, + related_term_index_settings=settings.related_term_index_settings, + ) + settings.storage_provider = storage + transcript = await Transcript.create(settings, name="bench") + return transcript, storage, db_path + + async def teardown(setup_rv): + _, storage, db_path = setup_rv + await storage.close() + os.remove(db_path) + + async def target(transcript, storage, db_path): + await transcript.add_messages_with_indexing(messages) try: - for i in range(WARMUP + ROUNDS): - db_path = os.path.join(tmpdir, f"bench_{i}.db") - storage = SqliteStorageProvider( - db_path, - message_type=TranscriptMessage, - message_text_index_settings=settings.message_text_index_settings, - related_term_index_settings=settings.related_term_index_settings, - ) - settings.storage_provider = storage - transcript = await Transcript.create(settings, name="bench") - - start = time.perf_counter() - await transcript.add_messages_with_indexing(messages) - elapsed = time.perf_counter() - start - - await storage.close() - os.remove(db_path) - - if i >= WARMUP: - times.append(elapsed) - - _report("add_messages (200)", times) + await async_benchmark.pedantic( + target, setup=setup, teardown=teardown, rounds=20, warmup_rounds=3 + ) finally: shutil.rmtree(tmpdir, ignore_errors=True) From 93a05a7ec964fa4c848e1dc9296e990c13a93e49 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Fri, 10 Apr 2026 01:27:36 -0500 Subject: [PATCH 08/20] Remove section-divider comments from benchmark tests --- tests/benchmarks/test_benchmark_indexing.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/tests/benchmarks/test_benchmark_indexing.py b/tests/benchmarks/test_benchmark_indexing.py index d25c4b6e..ad64b46b 100644 --- a/tests/benchmarks/test_benchmark_indexing.py +++ b/tests/benchmarks/test_benchmark_indexing.py @@ -36,11 +36,6 @@ CONFUSE_A_CAT_VTT = os.path.join(TESTDATA, "Confuse-A-Cat.vtt") -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - def make_settings() -> ConversationSettings: """Create conversation settings with fake embedding model (no API keys).""" model = create_test_embedding_model() @@ -73,11 +68,6 @@ def synthetic_messages(n: int) -> list[TranscriptMessage]: ] -# --------------------------------------------------------------------------- -# Test cases -# --------------------------------------------------------------------------- - - @pytest.mark.asyncio async def test_benchmark_vtt_ingest(async_benchmark): """Benchmark indexing of pre-parsed VTT messages (Confuse-A-Cat, 40 msgs).""" From dd5d738ee3b67fb6a72ac6cbdce4736b907cfbbc Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Fri, 10 Apr 2026 01:38:42 -0500 Subject: [PATCH 09/20] Extract shared benchmark harness to reduce duplication Move repeated setup/teardown/target pattern into run_indexing_benchmark() helper. Each test now delegates with just messages and message_type. --- tests/benchmarks/test_benchmark_indexing.py | 85 ++++----------------- 1 file changed, 16 insertions(+), 69 deletions(-) diff --git a/tests/benchmarks/test_benchmark_indexing.py b/tests/benchmarks/test_benchmark_indexing.py index ad64b46b..d730ca5c 100644 --- a/tests/benchmarks/test_benchmark_indexing.py +++ b/tests/benchmarks/test_benchmark_indexing.py @@ -68,10 +68,8 @@ def synthetic_messages(n: int) -> list[TranscriptMessage]: ] -@pytest.mark.asyncio -async def test_benchmark_vtt_ingest(async_benchmark): - """Benchmark indexing of pre-parsed VTT messages (Confuse-A-Cat, 40 msgs).""" - messages = await extract_vtt_messages(CONFUSE_A_CAT_VTT) +async def run_indexing_benchmark(async_benchmark, messages, message_type): + """Shared benchmark harness: fresh DB per round, only hot path timed.""" settings = make_settings() tmpdir = tempfile.mkdtemp() counter = itertools.count() @@ -81,7 +79,7 @@ async def setup(): db_path = os.path.join(tmpdir, f"bench_{i}.db") storage = SqliteStorageProvider( db_path, - message_type=ConversationMessage, + message_type=message_type, message_text_index_settings=settings.message_text_index_settings, related_term_index_settings=settings.related_term_index_settings, ) @@ -106,74 +104,23 @@ async def target(transcript, storage, db_path): @pytest.mark.asyncio -async def test_benchmark_add_messages_50(async_benchmark): - """Benchmark add_messages_with_indexing with 50 synthetic messages.""" - messages = synthetic_messages(50) - settings = make_settings() - tmpdir = tempfile.mkdtemp() - counter = itertools.count() - - async def setup(): - i = next(counter) - db_path = os.path.join(tmpdir, f"bench_{i}.db") - storage = SqliteStorageProvider( - db_path, - message_type=TranscriptMessage, - message_text_index_settings=settings.message_text_index_settings, - related_term_index_settings=settings.related_term_index_settings, - ) - settings.storage_provider = storage - transcript = await Transcript.create(settings, name="bench") - return transcript, storage, db_path - - async def teardown(setup_rv): - _, storage, db_path = setup_rv - await storage.close() - os.remove(db_path) +async def test_benchmark_vtt_ingest(async_benchmark): + """Benchmark indexing of pre-parsed VTT messages (Confuse-A-Cat, 40 msgs).""" + messages = await extract_vtt_messages(CONFUSE_A_CAT_VTT) + await run_indexing_benchmark(async_benchmark, messages, ConversationMessage) - async def target(transcript, storage, db_path): - await transcript.add_messages_with_indexing(messages) - try: - await async_benchmark.pedantic( - target, setup=setup, teardown=teardown, rounds=20, warmup_rounds=3 - ) - finally: - shutil.rmtree(tmpdir, ignore_errors=True) +@pytest.mark.asyncio +async def test_benchmark_add_messages_50(async_benchmark): + """Benchmark add_messages_with_indexing with 50 synthetic messages.""" + await run_indexing_benchmark( + async_benchmark, synthetic_messages(50), TranscriptMessage + ) @pytest.mark.asyncio async def test_benchmark_add_messages_200(async_benchmark): """Benchmark add_messages_with_indexing with 200 synthetic messages.""" - messages = synthetic_messages(200) - settings = make_settings() - tmpdir = tempfile.mkdtemp() - counter = itertools.count() - - async def setup(): - i = next(counter) - db_path = os.path.join(tmpdir, f"bench_{i}.db") - storage = SqliteStorageProvider( - db_path, - message_type=TranscriptMessage, - message_text_index_settings=settings.message_text_index_settings, - related_term_index_settings=settings.related_term_index_settings, - ) - settings.storage_provider = storage - transcript = await Transcript.create(settings, name="bench") - return transcript, storage, db_path - - async def teardown(setup_rv): - _, storage, db_path = setup_rv - await storage.close() - os.remove(db_path) - - async def target(transcript, storage, db_path): - await transcript.add_messages_with_indexing(messages) - - try: - await async_benchmark.pedantic( - target, setup=setup, teardown=teardown, rounds=20, warmup_rounds=3 - ) - finally: - shutil.rmtree(tmpdir, ignore_errors=True) + await run_indexing_benchmark( + async_benchmark, synthetic_messages(200), TranscriptMessage + ) From 4148b7a7313362e74b5fc82eecb2d0a329b8a107 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Fri, 10 Apr 2026 01:41:33 -0500 Subject: [PATCH 10/20] Remove underscore prefix from collect helper functions Rename _collect_{facet,entity,action}_{terms,properties} to drop the leading underscore in propindex.py and semrefindex.py. --- src/typeagent/storage/memory/propindex.py | 12 ++++++------ src/typeagent/storage/memory/semrefindex.py | 14 +++++++------- 2 files changed, 13 insertions(+), 13 deletions(-) diff --git a/src/typeagent/storage/memory/propindex.py b/src/typeagent/storage/memory/propindex.py index 1a1c7968..7ff49e29 100644 --- a/src/typeagent/storage/memory/propindex.py +++ b/src/typeagent/storage/memory/propindex.py @@ -109,7 +109,7 @@ async def build_property_index(conversation: IConversation) -> None: await add_to_property_index(conversation, 0) -def _collect_facet_properties( +def collect_facet_properties( facet: kplib.Facet | None, ordinal: SemanticRefOrdinal, ) -> list[tuple[str, str, SemanticRefOrdinal]]: @@ -127,7 +127,7 @@ def _collect_facet_properties( return props -def _collect_entity_properties( +def collect_entity_properties( entity: kplib.ConcreteEntity, ordinal: SemanticRefOrdinal, ) -> list[tuple[str, str, SemanticRefOrdinal]]: @@ -139,11 +139,11 @@ def _collect_entity_properties( props.append((PropertyNames.EntityType.value, t, ordinal)) if entity.facets: for facet in entity.facets: - props.extend(_collect_facet_properties(facet, ordinal)) + props.extend(collect_facet_properties(facet, ordinal)) return props -def _collect_action_properties( +def collect_action_properties( action: kplib.Action, ordinal: SemanticRefOrdinal, ) -> list[tuple[str, str, SemanticRefOrdinal]]: @@ -186,11 +186,11 @@ async def add_to_property_index( assert semantic_ref.semantic_ref_ordinal == semantic_ref_ordinal if isinstance(semantic_ref.knowledge, kplib.Action): collected.extend( - _collect_action_properties(semantic_ref.knowledge, semantic_ref_ordinal) + collect_action_properties(semantic_ref.knowledge, semantic_ref_ordinal) ) elif isinstance(semantic_ref.knowledge, kplib.ConcreteEntity): collected.extend( - _collect_entity_properties(semantic_ref.knowledge, semantic_ref_ordinal) + collect_entity_properties(semantic_ref.knowledge, semantic_ref_ordinal) ) elif isinstance(semantic_ref.knowledge, Tag): collected.append( diff --git a/src/typeagent/storage/memory/semrefindex.py b/src/typeagent/storage/memory/semrefindex.py index 7f131d9a..a98a3b21 100644 --- a/src/typeagent/storage/memory/semrefindex.py +++ b/src/typeagent/storage/memory/semrefindex.py @@ -577,7 +577,7 @@ async def add_metadata_to_index[TMessage: IMessage]( i += 1 -def _collect_facet_terms(facet: kplib.Facet | None) -> list[str]: +def collect_facet_terms(facet: kplib.Facet | None) -> list[str]: """Collect terms from a facet without touching any index.""" if facet is None: return [] @@ -587,18 +587,18 @@ def _collect_facet_terms(facet: kplib.Facet | None) -> list[str]: return terms -def _collect_entity_terms(entity: kplib.ConcreteEntity) -> list[str]: +def collect_entity_terms(entity: kplib.ConcreteEntity) -> list[str]: """Collect all terms an entity would add to the semantic ref index.""" terms = [entity.name] for t in entity.type: terms.append(t) if entity.facets: for facet in entity.facets: - terms.extend(_collect_facet_terms(facet)) + terms.extend(collect_facet_terms(facet)) return terms -def _collect_action_terms(action: kplib.Action) -> list[str]: +def collect_action_terms(action: kplib.Action) -> list[str]: """Collect all terms an action would add to the semantic ref index.""" terms = [" ".join(action.verbs)] if action.subject_entity_name != "none": @@ -615,7 +615,7 @@ def _collect_action_terms(action: kplib.Action) -> list[str]: terms.append(param.name) if isinstance(param.value, str): terms.append(param.value) - terms.extend(_collect_facet_terms(action.subject_entity_facet)) + terms.extend(collect_facet_terms(action.subject_entity_facet)) return terms @@ -641,7 +641,7 @@ async def add_metadata_to_index_from_list[TMessage: IMessage]( knowledge=entity, ) collected_refs.append(ref) - for term in _collect_entity_terms(entity): + for term in collect_entity_terms(entity): collected_terms.append((term, next_ordinal)) next_ordinal += 1 for action in knowledge_response.actions: @@ -652,7 +652,7 @@ async def add_metadata_to_index_from_list[TMessage: IMessage]( knowledge=action, ) collected_refs.append(ref) - for term in _collect_action_terms(action): + for term in collect_action_terms(action): collected_terms.append((term, next_ordinal)) next_ordinal += 1 for topic_response in knowledge_response.topics: From aa7e958eac095b4e6d6b0fd4f196171b95f61df7 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Fri, 10 Apr 2026 01:46:50 -0500 Subject: [PATCH 11/20] Add pytest-async-benchmark as dev dependency Install from fork with pedantic mode support for benchmark tests. --- pyproject.toml | 4 ++++ uv.lock | 11 +++++++++++ 2 files changed, 15 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 1339e34e..3cb25d7a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,6 +58,9 @@ Documentation = "https://github.com/microsoft/typeagent-py/tree/main/docs/README [tool.uv.build-backend] module-root = "src" +[tool.uv.sources] +pytest-async-benchmark = { git = "https://github.com/KRRT7/pytest-async-benchmark.git", rev = "feat/pedantic-mode" } + [tool.pytest.ini_options] asyncio_default_fixture_loop_scope = "function" testpaths = ["tests"] @@ -91,6 +94,7 @@ dev = [ "opentelemetry-instrumentation-httpx>=0.57b0", "pyright>=1.1.408", # 407 has a regression "pytest>=8.3.5", + "pytest-async-benchmark", "pytest-asyncio>=0.26.0", "pytest-mock>=3.14.0", ] diff --git a/uv.lock b/uv.lock index 4eab7ead..e2b66a3d 100644 --- a/uv.lock +++ b/uv.lock @@ -1922,6 +1922,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3b/ab/b3226f0bd7cdcf710fbede2b3548584366da3b19b5021e74f5bde2a8fa3f/pytest-9.0.2-py3-none-any.whl", hash = "sha256:711ffd45bf766d5264d487b917733b453d917afd2b0ad65223959f59089f875b", size = 374801, upload-time = "2025-12-06T21:30:49.154Z" }, ] +[[package]] +name = "pytest-async-benchmark" +version = "0.2.0" +source = { git = "https://github.com/KRRT7/pytest-async-benchmark.git?rev=feat%2Fpedantic-mode#029d03634d140789baebc6c3c8f72d5c81a67f9a" } +dependencies = [ + { name = "pytest" }, + { name = "rich" }, +] + [[package]] name = "pytest-asyncio" version = "1.3.0" @@ -2398,6 +2407,7 @@ dev = [ { name = "opentelemetry-instrumentation-httpx" }, { name = "pyright" }, { name = "pytest" }, + { name = "pytest-async-benchmark" }, { name = "pytest-asyncio" }, { name = "pytest-mock" }, ] @@ -2436,6 +2446,7 @@ dev = [ { name = "opentelemetry-instrumentation-httpx", specifier = ">=0.57b0" }, { name = "pyright", specifier = ">=1.1.408" }, { name = "pytest", specifier = ">=8.3.5" }, + { name = "pytest-async-benchmark", git = "https://github.com/KRRT7/pytest-async-benchmark.git?rev=feat%2Fpedantic-mode" }, { name = "pytest-asyncio", specifier = ">=0.26.0" }, { name = "pytest-mock", specifier = ">=3.14.0" }, ] From 771e481f4a34592ebc1581019b52fb2e68092ec7 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Fri, 10 Apr 2026 01:52:09 -0500 Subject: [PATCH 12/20] Fix pyright errors: use Sequence for batch method signatures Change list to Sequence in add_terms_batch and add_properties_batch interfaces and implementations to satisfy covariance. Add missing add_terms_batch to FakeTermIndex in conftest.py. --- src/typeagent/knowpro/interfaces_core.py | 3 ++- src/typeagent/knowpro/interfaces_indexes.py | 2 +- src/typeagent/storage/memory/propindex.py | 27 +++++++++++++++++---- src/typeagent/storage/memory/semrefindex.py | 8 +++--- src/typeagent/storage/sqlite/propindex.py | 10 +++++++- src/typeagent/storage/sqlite/semrefindex.py | 7 +++++- tests/conftest.py | 9 ++++++- 7 files changed, 53 insertions(+), 13 deletions(-) diff --git a/src/typeagent/knowpro/interfaces_core.py b/src/typeagent/knowpro/interfaces_core.py index 72f11f8c..cd9e885c 100644 --- a/src/typeagent/knowpro/interfaces_core.py +++ b/src/typeagent/knowpro/interfaces_core.py @@ -4,6 +4,7 @@ from __future__ import annotations +from collections.abc import Sequence from datetime import datetime as Datetime from typing import ( Any, @@ -170,7 +171,7 @@ async def add_term( async def add_terms_batch( self, - terms: list[tuple[str, SemanticRefOrdinal | ScoredSemanticRefOrdinal]], + terms: Sequence[tuple[str, SemanticRefOrdinal | ScoredSemanticRefOrdinal]], ) -> None: ... async def remove_term( diff --git a/src/typeagent/knowpro/interfaces_indexes.py b/src/typeagent/knowpro/interfaces_indexes.py index 3ae62024..6c348a01 100644 --- a/src/typeagent/knowpro/interfaces_indexes.py +++ b/src/typeagent/knowpro/interfaces_indexes.py @@ -61,7 +61,7 @@ async def add_property( async def add_properties_batch( self, - properties: list[tuple[str, str, SemanticRefOrdinal | ScoredSemanticRefOrdinal]], + properties: Sequence[tuple[str, str, SemanticRefOrdinal | ScoredSemanticRefOrdinal]], ) -> None: ... async def lookup_property( diff --git a/src/typeagent/storage/memory/propindex.py b/src/typeagent/storage/memory/propindex.py index 7ff49e29..f9717b24 100644 --- a/src/typeagent/storage/memory/propindex.py +++ b/src/typeagent/storage/memory/propindex.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. +from collections.abc import Sequence import enum from typing import assert_never @@ -156,7 +157,13 @@ def collect_action_properties( if action.object_entity_name != "none": props.append((PropertyNames.Object.value, action.object_entity_name, ordinal)) if action.indirect_object_entity_name != "none": - props.append((PropertyNames.IndirectObject.value, action.indirect_object_entity_name, ordinal)) + props.append( + ( + PropertyNames.IndirectObject.value, + action.indirect_object_entity_name, + ordinal, + ) + ) return props @@ -186,15 +193,23 @@ async def add_to_property_index( assert semantic_ref.semantic_ref_ordinal == semantic_ref_ordinal if isinstance(semantic_ref.knowledge, kplib.Action): collected.extend( - collect_action_properties(semantic_ref.knowledge, semantic_ref_ordinal) + collect_action_properties( + semantic_ref.knowledge, semantic_ref_ordinal + ) ) elif isinstance(semantic_ref.knowledge, kplib.ConcreteEntity): collected.extend( - collect_entity_properties(semantic_ref.knowledge, semantic_ref_ordinal) + collect_entity_properties( + semantic_ref.knowledge, semantic_ref_ordinal + ) ) elif isinstance(semantic_ref.knowledge, Tag): collected.append( - (PropertyNames.Tag.value, semantic_ref.knowledge.text, semantic_ref_ordinal) + ( + PropertyNames.Tag.value, + semantic_ref.knowledge.text, + semantic_ref_ordinal, + ) ) elif isinstance(semantic_ref.knowledge, Topic): pass @@ -239,7 +254,9 @@ async def add_property( async def add_properties_batch( self, - properties: list[tuple[str, str, SemanticRefOrdinal | ScoredSemanticRefOrdinal]], + properties: Sequence[ + tuple[str, str, SemanticRefOrdinal | ScoredSemanticRefOrdinal] + ], ) -> None: for name, value, ordinal in properties: await self.add_property(name, value, ordinal) diff --git a/src/typeagent/storage/memory/semrefindex.py b/src/typeagent/storage/memory/semrefindex.py index a98a3b21..8437bacd 100644 --- a/src/typeagent/storage/memory/semrefindex.py +++ b/src/typeagent/storage/memory/semrefindex.py @@ -3,11 +3,13 @@ from __future__ import annotations # TODO: Avoid -from collections.abc import AsyncIterable, Callable +from collections.abc import AsyncIterable, Callable, Sequence from typechat import Failure -from ...knowpro import convknowledge, knowledge_schema as kplib, secindex +from ...knowpro import convknowledge +from ...knowpro import knowledge_schema as kplib +from ...knowpro import secindex from ...knowpro.convsettings import ConversationSettings, SemanticRefIndexSettings from ...knowpro.interfaces import ( # Interfaces.; Other imports. IConversation, @@ -711,7 +713,7 @@ async def add_term( async def add_terms_batch( self, - terms: list[tuple[str, SemanticRefOrdinal | ScoredSemanticRefOrdinal]], + terms: Sequence[tuple[str, SemanticRefOrdinal | ScoredSemanticRefOrdinal]], ) -> None: for term, ordinal in terms: await self.add_term(term, ordinal) diff --git a/src/typeagent/storage/sqlite/propindex.py b/src/typeagent/storage/sqlite/propindex.py index 6a619869..f9704b45 100644 --- a/src/typeagent/storage/sqlite/propindex.py +++ b/src/typeagent/storage/sqlite/propindex.py @@ -3,6 +3,7 @@ """SQLite-based property index implementation.""" +from collections.abc import Sequence import sqlite3 from ...knowpro import interfaces @@ -69,7 +70,13 @@ async def add_property( async def add_properties_batch( self, - properties: list[tuple[str, str, interfaces.SemanticRefOrdinal | interfaces.ScoredSemanticRefOrdinal]], + properties: Sequence[ + tuple[ + str, + str, + interfaces.SemanticRefOrdinal | interfaces.ScoredSemanticRefOrdinal, + ] + ], ) -> None: if not properties: return @@ -77,6 +84,7 @@ async def add_properties_batch( make_property_term_text, split_property_term_text, ) + rows = [] for property_name, value, ordinal in properties: if isinstance(ordinal, interfaces.ScoredSemanticRefOrdinal): diff --git a/src/typeagent/storage/sqlite/semrefindex.py b/src/typeagent/storage/sqlite/semrefindex.py index 0925eb45..ac68a1e0 100644 --- a/src/typeagent/storage/sqlite/semrefindex.py +++ b/src/typeagent/storage/sqlite/semrefindex.py @@ -3,6 +3,7 @@ """SQLite-based semantic reference index implementation.""" +from collections.abc import Sequence import re import sqlite3 import unicodedata @@ -58,7 +59,11 @@ async def add_term( async def add_terms_batch( self, - terms: list[tuple[str, interfaces.SemanticRefOrdinal | interfaces.ScoredSemanticRefOrdinal]], + terms: Sequence[ + tuple[ + str, interfaces.SemanticRefOrdinal | interfaces.ScoredSemanticRefOrdinal + ] + ], ) -> None: if not terms: return diff --git a/tests/conftest.py b/tests/conftest.py index dae619c1..7f0f11f5 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,7 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. -from collections.abc import AsyncGenerator, Callable, Iterator +from collections.abc import AsyncGenerator, Callable, Iterator, Sequence import os from pathlib import Path import tempfile @@ -236,6 +236,13 @@ async def add_term( self.term_to_refs[term].append(scored_ref) return term + async def add_terms_batch( + self, + terms: Sequence[tuple[str, int | ScoredSemanticRefOrdinal]], + ) -> None: + for term, ordinal in terms: + await self.add_term(term, ordinal) + async def remove_term(self, term: str, semantic_ref_ordinal: int) -> None: if term in self.term_to_refs: self.term_to_refs[term] = [ From 9e3de1c528307fb1638b489e15842f95649bb3fe Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Fri, 10 Apr 2026 02:24:57 -0500 Subject: [PATCH 13/20] Optimize fuzzy_lookup_embedding with numpy vectorized ops MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace Python-level list comprehension + sort with numpy operations: - No-predicate path: np.flatnonzero for score filtering, np.argpartition for O(n) top-k selection — avoids building ScoredInt for every vector - Predicate path: numpy pre-filters by score, applies predicate only to candidates above threshold - Subset lookup: numpy fancy indexing computes dot products only for subset indices instead of delegating to full-vector scan with predicate --- src/typeagent/aitools/vectorbase.py | 65 +++++++++++---- tests/benchmarks/test_benchmark_vectorbase.py | 83 +++++++++++++++++++ 2 files changed, 133 insertions(+), 15 deletions(-) create mode 100644 tests/benchmarks/test_benchmark_vectorbase.py diff --git a/src/typeagent/aitools/vectorbase.py b/src/typeagent/aitools/vectorbase.py index 63e2e77a..4ca8be06 100644 --- a/src/typeagent/aitools/vectorbase.py +++ b/src/typeagent/aitools/vectorbase.py @@ -1,7 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. -from collections.abc import Callable, Iterable +from collections.abc import Callable from dataclasses import dataclass import numpy as np @@ -132,17 +132,35 @@ def fuzzy_lookup_embedding( min_score = 0.0 if len(self._vectors) == 0: return [] - # 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] + scores = np.dot(self._vectors, embedding) + + if predicate is None: + # Fast numpy path: filter and top-k without Python-level iteration. + indices = np.flatnonzero(scores >= min_score) + if len(indices) == 0: + return [] + filtered_scores = scores[indices] + if len(indices) <= max_hits: + order = np.argsort(filtered_scores)[::-1] + else: + # argpartition is O(n) vs O(n log n) for full sort. + top_k = np.argpartition(filtered_scores, -max_hits)[-max_hits:] + order = top_k[np.argsort(filtered_scores[top_k])[::-1]] + return [ + ScoredInt(int(indices[i]), float(filtered_scores[i])) for i in order + ] + else: + # Predicate path: pre-filter by score in numpy, then apply predicate + # only to candidates that pass the score threshold. + candidates = np.flatnonzero(scores >= min_score) + scored_ordinals = [ + ScoredInt(int(i), float(scores[i])) + for i in candidates + if predicate(int(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, @@ -150,10 +168,27 @@ def fuzzy_lookup_embedding_in_subset( max_hits: int | None = None, min_score: float | None = None, ) -> list[ScoredInt]: - ordinals_set = set(ordinals_of_subset) - return self.fuzzy_lookup_embedding( - embedding, max_hits, min_score, lambda i: i in ordinals_set - ) + if max_hits is None: + max_hits = 10 + if min_score is None: + min_score = 0.0 + if not ordinals_of_subset or len(self._vectors) == 0: + return [] + # Compute dot products only for the subset instead of all vectors. + subset = np.asarray(ordinals_of_subset) + scores = np.dot(self._vectors[subset], embedding) + indices = np.flatnonzero(scores >= min_score) + if len(indices) == 0: + return [] + filtered_scores = scores[indices] + if len(indices) <= max_hits: + order = np.argsort(filtered_scores)[::-1] + else: + top_k = np.argpartition(filtered_scores, -max_hits)[-max_hits:] + order = top_k[np.argsort(filtered_scores[top_k])[::-1]] + return [ + ScoredInt(int(subset[indices[i]]), float(filtered_scores[i])) for i in order + ] async def fuzzy_lookup( self, diff --git a/tests/benchmarks/test_benchmark_vectorbase.py b/tests/benchmarks/test_benchmark_vectorbase.py new file mode 100644 index 00000000..b61859a8 --- /dev/null +++ b/tests/benchmarks/test_benchmark_vectorbase.py @@ -0,0 +1,83 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Benchmarks for VectorBase fuzzy lookup methods. + +Measures fuzzy_lookup_embedding and fuzzy_lookup_embedding_in_subset +with varying vector counts and result sizes. +""" + +import numpy as np +import pytest + +from typeagent.aitools.model_adapters import create_test_embedding_model +from typeagent.aitools.vectorbase import TextEmbeddingIndexSettings, VectorBase + +EMBEDDING_DIM = 384 # Typical small embedding model dimension + + +def make_populated_vector_base(n_vectors: int) -> tuple[VectorBase, np.ndarray]: + """Create a VectorBase with n_vectors random normalized embeddings.""" + settings = TextEmbeddingIndexSettings(create_test_embedding_model()) + vb = VectorBase(settings) + rng = np.random.default_rng(42) + embeddings = rng.standard_normal((n_vectors, EMBEDDING_DIM)).astype(np.float32) + # Normalize to unit vectors (as the real pipeline does). + norms = np.linalg.norm(embeddings, axis=1, keepdims=True) + embeddings = embeddings / norms + vb.add_embeddings(None, embeddings) + # Query vector: also normalized. + query = rng.standard_normal(EMBEDDING_DIM).astype(np.float32) + query = query / np.linalg.norm(query) + return vb, query + + +# --- fuzzy_lookup_embedding --- + + +@pytest.mark.asyncio +async def test_benchmark_fuzzy_lookup_1k(async_benchmark): + vb, query = make_populated_vector_base(1_000) + + async def target(): + vb.fuzzy_lookup_embedding(query, max_hits=10, min_score=0.0) + + await async_benchmark.pedantic(target, rounds=200, warmup_rounds=20) + + +@pytest.mark.asyncio +async def test_benchmark_fuzzy_lookup_10k(async_benchmark): + vb, query = make_populated_vector_base(10_000) + + async def target(): + vb.fuzzy_lookup_embedding(query, max_hits=10, min_score=0.0) + + await async_benchmark.pedantic(target, rounds=200, warmup_rounds=20) + + +@pytest.mark.asyncio +async def test_benchmark_fuzzy_lookup_10k_with_predicate(async_benchmark): + vb, query = make_populated_vector_base(10_000) + # Predicate that accepts ~50% of indices. + even_only = lambda i: i % 2 == 0 + + async def target(): + vb.fuzzy_lookup_embedding( + query, max_hits=10, min_score=0.0, predicate=even_only + ) + + await async_benchmark.pedantic(target, rounds=200, warmup_rounds=20) + + +# --- fuzzy_lookup_embedding_in_subset --- + + +@pytest.mark.asyncio +async def test_benchmark_fuzzy_lookup_subset_1k_of_10k(async_benchmark): + vb, query = make_populated_vector_base(10_000) + subset = list(range(0, 10_000, 10)) # 1000 indices + + async def target(): + vb.fuzzy_lookup_embedding_in_subset(query, subset, max_hits=10, min_score=0.0) + + await async_benchmark.pedantic(target, rounds=200, warmup_rounds=20) From cfac6572412f7d3ce905375fa542874cc13653bc Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Fri, 10 Apr 2026 03:37:56 -0500 Subject: [PATCH 14/20] Add batch metadata query to avoid N+1 in lookup_term_filtered MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit lookup_term_filtered called get_item() per scored ref — one SELECT and full deserialization per match. The filter only needs knowledge_type (a plain column) and range (json.loads of range_json), never the expensive knowledge_json deserialization (64% of per-row cost). Add get_metadata_multiple to ISemanticRefCollection that fetches only semref_id, range_json, knowledge_type in a single batch query. Replace the N+1 loop in lookup_term_filtered with one get_metadata_multiple call. Benchmark (200 matches, 200 rounds): 4.38ms → 1.32ms (3.3x speedup). --- src/typeagent/knowpro/interfaces_storage.py | 19 +++- src/typeagent/knowpro/query.py | 18 ++-- src/typeagent/storage/memory/collections.py | 13 +++ src/typeagent/storage/sqlite/collections.py | 25 +++++ tests/benchmarks/test_benchmark_query.py | 104 ++++++++++++++++++++ 5 files changed, 168 insertions(+), 11 deletions(-) create mode 100644 tests/benchmarks/test_benchmark_query.py diff --git a/src/typeagent/knowpro/interfaces_storage.py b/src/typeagent/knowpro/interfaces_storage.py index a82fe7ad..97f7b600 100644 --- a/src/typeagent/knowpro/interfaces_storage.py +++ b/src/typeagent/knowpro/interfaces_storage.py @@ -6,16 +6,18 @@ from collections.abc import AsyncIterable, Iterable from datetime import datetime as Datetime -from typing import Any, Protocol, Self +from typing import Any, NamedTuple, Protocol, Self from pydantic.dataclasses import dataclass from .interfaces_core import ( IMessage, ITermToSemanticRefIndex, + KnowledgeType, MessageOrdinal, SemanticRef, SemanticRefOrdinal, + TextRange, ) from .interfaces_indexes import ( IConversationSecondaryIndexes, @@ -57,6 +59,14 @@ class ConversationMetadata: extra: dict[str, str] | None = None +class SemanticRefMetadata(NamedTuple): + """Lightweight metadata for filtering without full knowledge deserialization.""" + + ordinal: SemanticRefOrdinal + range: TextRange + knowledge_type: KnowledgeType + + class IReadonlyCollection[T, TOrdinal](AsyncIterable[T], Protocol): async def size(self) -> int: ... @@ -91,6 +101,12 @@ class IMessageCollection[TMessage: IMessage]( class ISemanticRefCollection(ICollection[SemanticRef, SemanticRefOrdinal], Protocol): """A collection of SemanticRefs.""" + async def get_metadata_multiple( + self, ordinals: list[SemanticRefOrdinal] + ) -> list[SemanticRefMetadata]: + """Batch-fetch lightweight metadata without deserializing knowledge.""" + ... + class IStorageProvider[TMessage: IMessage](Protocol): """API spec for storage providers -- maybe in-memory or persistent.""" @@ -190,4 +206,5 @@ class IConversation[ "ISemanticRefCollection", "IStorageProvider", "STATUS_INGESTED", + "SemanticRefMetadata", ] diff --git a/src/typeagent/knowpro/query.py b/src/typeagent/knowpro/query.py index 44fa06ec..5859e3bc 100644 --- a/src/typeagent/knowpro/query.py +++ b/src/typeagent/knowpro/query.py @@ -37,6 +37,7 @@ ScoredSemanticRefOrdinal, SearchTerm, SemanticRef, + SemanticRefMetadata, SemanticRefOrdinal, SemanticRefSearchResult, Term, @@ -174,17 +175,14 @@ async def lookup_term_filtered( semantic_ref_index: ITermToSemanticRefIndex, term: Term, semantic_refs: ISemanticRefCollection, - filter: Callable[[SemanticRef, ScoredSemanticRefOrdinal], bool], + filter: Callable[[SemanticRefMetadata, ScoredSemanticRefOrdinal], bool], ) -> list[ScoredSemanticRefOrdinal] | None: """Look up a term in the semantic reference index and filter the results.""" scored_refs = await semantic_ref_index.lookup_term(term.text) if scored_refs: - filtered = [] - for sr in scored_refs: - semantic_ref = await semantic_refs.get_item(sr.semantic_ref_ordinal) - if filter(semantic_ref, sr): - filtered.append(sr) - return filtered + ordinals = [sr.semantic_ref_ordinal for sr in scored_refs] + metadata = await semantic_refs.get_metadata_multiple(ordinals) + return [sr for sr, m in zip(scored_refs, metadata) if filter(m, sr)] return None @@ -202,10 +200,10 @@ async def lookup_term( semantic_ref_index, term, semantic_refs, - lambda sr, _: ( - not knowledge_type or sr.knowledge.knowledge_type == knowledge_type + lambda m, _: ( + not knowledge_type or m.knowledge_type == knowledge_type ) - and ranges_in_scope.is_range_in_scope(sr.range), + and ranges_in_scope.is_range_in_scope(m.range), ) return await semantic_ref_index.lookup_term(term.text) diff --git a/src/typeagent/storage/memory/collections.py b/src/typeagent/storage/memory/collections.py index 9973a290..8a5b14eb 100644 --- a/src/typeagent/storage/memory/collections.py +++ b/src/typeagent/storage/memory/collections.py @@ -10,6 +10,7 @@ IMessage, MessageOrdinal, SemanticRef, + SemanticRefMetadata, SemanticRefOrdinal, ) @@ -63,6 +64,18 @@ async def extend(self, items: Iterable[T]) -> None: class MemorySemanticRefCollection(MemoryCollection[SemanticRef, SemanticRefOrdinal]): """A collection of semantic references.""" + async def get_metadata_multiple( + self, ordinals: list[SemanticRefOrdinal] + ) -> list[SemanticRefMetadata]: + return [ + SemanticRefMetadata( + ordinal=o, + range=self.items[o].range, + knowledge_type=self.items[o].knowledge.knowledge_type, + ) + for o in ordinals + ] + class MemoryMessageCollection[TMessage: IMessage]( MemoryCollection[TMessage, MessageOrdinal] diff --git a/src/typeagent/storage/sqlite/collections.py b/src/typeagent/storage/sqlite/collections.py index 9730f6d1..911fff94 100644 --- a/src/typeagent/storage/sqlite/collections.py +++ b/src/typeagent/storage/sqlite/collections.py @@ -340,6 +340,31 @@ async def get_multiple(self, arg: list[int]) -> list[interfaces.SemanticRef]: assert set(rowdict) == set(arg) return [self._deserialize_semantic_ref_from_row(rowdict[ordl]) for ordl in arg] + async def get_metadata_multiple( + self, ordinals: list[int] + ) -> list[interfaces.SemanticRefMetadata]: + if not ordinals: + return [] + cursor = self.db.cursor() + placeholders = ",".join("?" * len(ordinals)) + cursor.execute( + f""" + SELECT semref_id, range_json, knowledge_type + FROM SemanticRefs WHERE semref_id IN ({placeholders}) + """, + ordinals, + ) + rows = cursor.fetchall() + rowdict = {r[0]: r for r in rows} + return [ + interfaces.SemanticRefMetadata( + ordinal=rowdict[o][0], + range=interfaces.TextRange.deserialize(json.loads(rowdict[o][1])), + knowledge_type=rowdict[o][2], + ) + for o in ordinals + ] + async def append(self, item: interfaces.SemanticRef) -> None: cursor = self.db.cursor() semref_id, range_json, knowledge_type, knowledge_json = ( diff --git a/tests/benchmarks/test_benchmark_query.py b/tests/benchmarks/test_benchmark_query.py new file mode 100644 index 00000000..8c4dd137 --- /dev/null +++ b/tests/benchmarks/test_benchmark_query.py @@ -0,0 +1,104 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Benchmark for lookup_term_filtered — measures the N+1 query pattern. + +After indexing 200 synthetic messages, looks up a high-frequency term +and filters results via lookup_term_filtered. Each call triggers +one get_item() SELECT per matching semantic ref (N+1 pattern). + +Run: + uv run python -m pytest tests/benchmarks/test_benchmark_query.py -v -s +""" + +import os +import tempfile + +import pytest + +from typeagent.aitools.model_adapters import create_test_embedding_model +from typeagent.knowpro.convsettings import ConversationSettings +from typeagent.knowpro.interfaces_core import Term +from typeagent.knowpro.query import lookup_term_filtered +from typeagent.storage.sqlite.provider import SqliteStorageProvider +from typeagent.transcripts.transcript import ( + Transcript, + TranscriptMessage, + TranscriptMessageMeta, +) + + +def make_settings() -> ConversationSettings: + model = create_test_embedding_model() + settings = ConversationSettings(model=model) + settings.semantic_ref_index_settings.auto_extract_knowledge = False + return settings + + +def synthetic_messages(n: int) -> list[TranscriptMessage]: + return [ + TranscriptMessage( + text_chunks=[f"Message {i} about topic {i % 10}"], + metadata=TranscriptMessageMeta(speaker=f"Speaker{i % 3}"), + tags=[f"tag{i % 5}"], + ) + for i in range(n) + ] + + +async def create_indexed_transcript( + db_path: str, settings: ConversationSettings, n_messages: int +) -> Transcript: + """Create and index a transcript, returning it ready for queries.""" + storage = SqliteStorageProvider( + db_path, + message_type=TranscriptMessage, + message_text_index_settings=settings.message_text_index_settings, + related_term_index_settings=settings.related_term_index_settings, + ) + settings.storage_provider = storage + transcript = await Transcript.create(settings, name="bench") + messages = synthetic_messages(n_messages) + await transcript.add_messages_with_indexing(messages) + return transcript + + +@pytest.mark.asyncio +async def test_benchmark_lookup_term_filtered(async_benchmark): + """Benchmark lookup_term_filtered with N+1 get_item pattern.""" + settings = make_settings() + tmpdir = tempfile.mkdtemp() + db_path = os.path.join(tmpdir, "query_bench.db") + + transcript = await create_indexed_transcript(db_path, settings, 200) + + # Find a high-frequency term to look up. + semref_index = transcript.semantic_ref_index + terms = await semref_index.get_terms() + # Pick the term with the most matches. + best_term = None + best_count = 0 + for t in terms: + refs = await semref_index.lookup_term(t) + if refs and len(refs) > best_count: + best_count = len(refs) + best_term = t + + assert best_term is not None, "No terms found after indexing" + print(f"\nBenchmarking term '{best_term}' with {best_count} matches") + + term = Term(text=best_term) + semantic_refs = transcript.semantic_refs + # Filter that accepts all — isolates the get_item overhead. + accept_all = lambda sr, scored: True + + async def target(): + await lookup_term_filtered(semref_index, term, semantic_refs, accept_all) + + try: + await async_benchmark.pedantic(target, rounds=200, warmup_rounds=20) + finally: + await settings.storage_provider.close() + import shutil + + shutil.rmtree(tmpdir, ignore_errors=True) From f1537b471b59cb33c488d8f57e3dcf31b276da59 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Fri, 10 Apr 2026 03:49:46 -0500 Subject: [PATCH 15/20] Extend batch metadata query to remaining N+1 call sites Apply the same get_metadata_multiple pattern from lookup_term_filtered to four more sites that called get_item() in a loop: - propindex.lookup_property_in_property_index: filter by .range - SemanticRefAccumulator.group_matches_by_type: group by .knowledge_type - SemanticRefAccumulator.get_matches_in_scope: filter by .range - answers.get_scored_semantic_refs_from_ordinals_iter: two-phase metadata filter then batch get_multiple for matching full objects All sites now use a single batch query instead of N individual SELECTs, skipping knowledge_json deserialization where only range or knowledge_type is needed. --- src/typeagent/knowpro/answers.py | 29 +++++++++++++---------- src/typeagent/knowpro/collections.py | 23 +++++++++++------- src/typeagent/storage/memory/propindex.py | 13 +++++----- 3 files changed, 38 insertions(+), 27 deletions(-) diff --git a/src/typeagent/knowpro/answers.py b/src/typeagent/knowpro/answers.py index 58a536ed..663d5804 100644 --- a/src/typeagent/knowpro/answers.py +++ b/src/typeagent/knowpro/answers.py @@ -452,19 +452,22 @@ async def get_scored_semantic_refs_from_ordinals_iter( semantic_ref_matches: list[ScoredSemanticRefOrdinal], knowledge_type: KnowledgeType, ) -> list[Scored[SemanticRef]]: - result = [] - for semantic_ref_match in semantic_ref_matches: - semantic_ref = await semantic_refs.get_item( - semantic_ref_match.semantic_ref_ordinal - ) - if semantic_ref.knowledge.knowledge_type == knowledge_type: - result.append( - Scored( - item=semantic_ref, - score=semantic_ref_match.score, - ) - ) - return result + if not semantic_ref_matches: + return [] + ordinals = [m.semantic_ref_ordinal for m in semantic_ref_matches] + metadata = await semantic_refs.get_metadata_multiple(ordinals) + matching = [ + (sr_match, m.ordinal) + for sr_match, m in zip(semantic_ref_matches, metadata) + if m.knowledge_type == knowledge_type + ] + if not matching: + return [] + full_refs = await semantic_refs.get_multiple([o for _, o in matching]) + return [ + Scored(item=ref, score=sr_match.score) + for (sr_match, _), ref in zip(matching, full_refs) + ] def merge_scored_concrete_entities( diff --git a/src/typeagent/knowpro/collections.py b/src/typeagent/knowpro/collections.py index a2716577..6ff7b557 100644 --- a/src/typeagent/knowpro/collections.py +++ b/src/typeagent/knowpro/collections.py @@ -331,13 +331,17 @@ async def group_matches_by_type( self, semantic_refs: ISemanticRefCollection, ) -> dict[KnowledgeType, "SemanticRefAccumulator"]: + matches = list(self) + if not matches: + return {} + ordinals = [match.value for match in matches] + metadata = await semantic_refs.get_metadata_multiple(ordinals) groups: dict[KnowledgeType, SemanticRefAccumulator] = {} - for match in self: - semantic_ref = await semantic_refs.get_item(match.value) - group = groups.get(semantic_ref.knowledge.knowledge_type) + for match, m in zip(matches, metadata): + group = groups.get(m.knowledge_type) if group is None: group = SemanticRefAccumulator(self.search_term_matches) - groups[semantic_ref.knowledge.knowledge_type] = group + groups[m.knowledge_type] = group group.set_match(match) return groups @@ -346,11 +350,14 @@ async def get_matches_in_scope( semantic_refs: ISemanticRefCollection, ranges_in_scope: "TextRangesInScope", ) -> "SemanticRefAccumulator": + matches = list(self) + if not matches: + return SemanticRefAccumulator(self.search_term_matches) + ordinals = [match.value for match in matches] + metadata = await semantic_refs.get_metadata_multiple(ordinals) accumulator = SemanticRefAccumulator(self.search_term_matches) - for match in self: - if ranges_in_scope.is_range_in_scope( - (await semantic_refs.get_item(match.value)).range - ): + for match, m in zip(matches, metadata): + if ranges_in_scope.is_range_in_scope(m.range): accumulator.set_match(match) return accumulator diff --git a/src/typeagent/storage/memory/propindex.py b/src/typeagent/storage/memory/propindex.py index f9717b24..ecb3e85d 100644 --- a/src/typeagent/storage/memory/propindex.py +++ b/src/typeagent/storage/memory/propindex.py @@ -330,12 +330,13 @@ async def lookup_property_in_property_index( property_value, ) if ranges_in_scope is not None and scored_refs: - filtered_refs = [] - for sr in scored_refs: - semantic_ref = await semantic_refs.get_item(sr.semantic_ref_ordinal) - if ranges_in_scope.is_range_in_scope(semantic_ref.range): - filtered_refs.append(sr) - scored_refs = filtered_refs + ordinals = [sr.semantic_ref_ordinal for sr in scored_refs] + metadata = await semantic_refs.get_metadata_multiple(ordinals) + scored_refs = [ + sr + for sr, m in zip(scored_refs, metadata) + if ranges_in_scope.is_range_in_scope(m.range) + ] return scored_refs or None # Return None if no results From c8f69fd0400d7e66850d36066bc725e0b1513079 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Fri, 10 Apr 2026 04:13:38 -0500 Subject: [PATCH 16/20] Fix parse_azure_endpoint passing query string to AsyncAzureOpenAI MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit parse_azure_endpoint returned the raw URL including ?api-version=... which AsyncAzureOpenAI then mangled into invalid paths like ...?api-version=2024-06-01/openai/. Strip the query string before returning — api_version is already returned as a separate value and passed to the SDK independently. --- src/typeagent/aitools/utils.py | 6 ++++- tests/test_utils.py | 40 +++++++++++++++++++++++++++++++++- 2 files changed, 44 insertions(+), 2 deletions(-) diff --git a/src/typeagent/aitools/utils.py b/src/typeagent/aitools/utils.py index f3b49d65..eb21f9ec 100644 --- a/src/typeagent/aitools/utils.py +++ b/src/typeagent/aitools/utils.py @@ -198,7 +198,11 @@ def parse_azure_endpoint( f"{endpoint_envvar}={azure_endpoint} doesn't contain valid api-version field" ) - return azure_endpoint, m.group(1) + # Strip query string — AsyncAzureOpenAI expects a clean base URL and + # receives api_version as a separate parameter. + clean_endpoint = azure_endpoint.split("?", 1)[0] + + return clean_endpoint, m.group(1) def get_azure_api_key(azure_api_key: str) -> str: diff --git a/tests/test_utils.py b/tests/test_utils.py index 5966af61..7f806f74 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -67,7 +67,7 @@ def test_api_version_after_question_mark( ) endpoint, version = utils.parse_azure_endpoint("TEST_ENDPOINT") assert version == "2025-01-01-preview" - assert endpoint.startswith("https://") + assert endpoint == "https://myhost.openai.azure.com/openai/deployments/gpt-4" def test_api_version_after_ampersand(self, monkeypatch: pytest.MonkeyPatch) -> None: """api-version preceded by & (not the first query parameter).""" @@ -84,6 +84,44 @@ def test_missing_env_var_raises(self, monkeypatch: pytest.MonkeyPatch) -> None: with pytest.raises(RuntimeError, match="not found"): utils.parse_azure_endpoint("NONEXISTENT_ENDPOINT") + def test_query_string_stripped_from_endpoint( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Returned endpoint should not contain query string parameters.""" + monkeypatch.setenv( + "TEST_ENDPOINT", + "https://myhost.openai.azure.com?api-version=2024-06-01", + ) + endpoint, version = utils.parse_azure_endpoint("TEST_ENDPOINT") + assert endpoint == "https://myhost.openai.azure.com" + assert version == "2024-06-01" + + def test_query_string_stripped_with_path( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Query string stripped even when endpoint includes a path.""" + monkeypatch.setenv( + "TEST_ENDPOINT", + "https://myhost.openai.azure.com/openai/deployments/gpt-4?api-version=2025-01-01-preview", + ) + endpoint, version = utils.parse_azure_endpoint("TEST_ENDPOINT") + assert endpoint == "https://myhost.openai.azure.com/openai/deployments/gpt-4" + assert "?" not in endpoint + assert version == "2025-01-01-preview" + + def test_query_string_stripped_multiple_params( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """All query parameters stripped, not just api-version.""" + monkeypatch.setenv( + "TEST_ENDPOINT", + "https://myhost.openai.azure.com?foo=bar&api-version=2024-06-01", + ) + endpoint, version = utils.parse_azure_endpoint("TEST_ENDPOINT") + assert endpoint == "https://myhost.openai.azure.com" + assert "foo" not in endpoint + assert version == "2024-06-01" + def test_no_api_version_raises(self, monkeypatch: pytest.MonkeyPatch) -> None: """RuntimeError when the endpoint has no api-version field.""" monkeypatch.setenv( From ee86888133eca612b9391e298f1aa16641937e29 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Fri, 10 Apr 2026 04:27:03 -0500 Subject: [PATCH 17/20] Add benchmarks for all batch metadata query call sites --- tests/benchmarks/test_benchmark_query.py | 189 ++++++++++++++++++++--- 1 file changed, 169 insertions(+), 20 deletions(-) diff --git a/tests/benchmarks/test_benchmark_query.py b/tests/benchmarks/test_benchmark_query.py index 8c4dd137..7e948ad2 100644 --- a/tests/benchmarks/test_benchmark_query.py +++ b/tests/benchmarks/test_benchmark_query.py @@ -1,25 +1,35 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. -"""Benchmark for lookup_term_filtered — measures the N+1 query pattern. +"""Benchmarks for batch metadata query optimization. -After indexing 200 synthetic messages, looks up a high-frequency term -and filters results via lookup_term_filtered. Each call triggers -one get_item() SELECT per matching semantic ref (N+1 pattern). +After indexing 200 synthetic messages, exercises each function that was +converted from N+1 get_item() to batch get_metadata_multiple(). Run: uv run python -m pytest tests/benchmarks/test_benchmark_query.py -v -s """ import os +import shutil import tempfile import pytest from typeagent.aitools.model_adapters import create_test_embedding_model +from typeagent.knowpro.answers import get_scored_semantic_refs_from_ordinals_iter +from typeagent.knowpro.collections import ( + SemanticRefAccumulator, + TextRangeCollection, + TextRangesInScope, +) from typeagent.knowpro.convsettings import ConversationSettings -from typeagent.knowpro.interfaces_core import Term +from typeagent.knowpro.interfaces_core import Term, TextLocation, TextRange from typeagent.knowpro.query import lookup_term_filtered +from typeagent.storage.memory.propindex import ( + PropertyNames, + lookup_property_in_property_index, +) from typeagent.storage.sqlite.provider import SqliteStorageProvider from typeagent.transcripts.transcript import ( Transcript, @@ -63,19 +73,9 @@ async def create_indexed_transcript( return transcript -@pytest.mark.asyncio -async def test_benchmark_lookup_term_filtered(async_benchmark): - """Benchmark lookup_term_filtered with N+1 get_item pattern.""" - settings = make_settings() - tmpdir = tempfile.mkdtemp() - db_path = os.path.join(tmpdir, "query_bench.db") - - transcript = await create_indexed_transcript(db_path, settings, 200) - - # Find a high-frequency term to look up. - semref_index = transcript.semantic_ref_index +async def find_best_term(semref_index) -> tuple[str, int]: + """Find the term with the most matches in the semantic ref index.""" terms = await semref_index.get_terms() - # Pick the term with the most matches. best_term = None best_count = 0 for t in terms: @@ -83,13 +83,38 @@ async def test_benchmark_lookup_term_filtered(async_benchmark): if refs and len(refs) > best_count: best_count = len(refs) best_term = t - assert best_term is not None, "No terms found after indexing" + return best_term, best_count + + +def make_scope_first_half(n_messages: int) -> TextRangesInScope: + """Build a TextRangesInScope covering the first half of messages.""" + ranges = [ + TextRange( + start=TextLocation(i, 0), + end=TextLocation(i, 0), + ) + for i in range(n_messages // 2) + ] + scope = TextRangesInScope() + scope.add_text_ranges(TextRangeCollection(ranges)) + return scope + + +@pytest.mark.asyncio +async def test_benchmark_lookup_term_filtered(async_benchmark): + """Benchmark lookup_term_filtered with batch get_metadata_multiple.""" + settings = make_settings() + tmpdir = tempfile.mkdtemp() + db_path = os.path.join(tmpdir, "bench_ltf.db") + + transcript = await create_indexed_transcript(db_path, settings, 200) + semref_index = transcript.semantic_ref_index + best_term, best_count = await find_best_term(semref_index) print(f"\nBenchmarking term '{best_term}' with {best_count} matches") term = Term(text=best_term) semantic_refs = transcript.semantic_refs - # Filter that accepts all — isolates the get_item overhead. accept_all = lambda sr, scored: True async def target(): @@ -99,6 +124,130 @@ async def target(): await async_benchmark.pedantic(target, rounds=200, warmup_rounds=20) finally: await settings.storage_provider.close() - import shutil + shutil.rmtree(tmpdir, ignore_errors=True) + + +@pytest.mark.asyncio +async def test_benchmark_lookup_property_in_property_index(async_benchmark): + """Benchmark property lookup with range filtering.""" + settings = make_settings() + tmpdir = tempfile.mkdtemp() + db_path = os.path.join(tmpdir, "bench_prop.db") + + transcript = await create_indexed_transcript(db_path, settings, 200) + assert transcript.secondary_indexes is not None + property_index = transcript.secondary_indexes.property_to_semantic_ref_index + assert property_index is not None, "Property index not built" + + # Verify there are matches for entity type "person" + refs = await property_index.lookup_property( + PropertyNames.EntityType.value, "person" + ) + match_count = len(refs) if refs else 0 + print(f"\nBenchmarking property 'type=person' with {match_count} matches") + assert match_count > 0 + + scope = make_scope_first_half(200) + + async def target(): + await lookup_property_in_property_index( + property_index, + PropertyNames.EntityType.value, + "person", + transcript.semantic_refs, + ranges_in_scope=scope, + ) + + try: + await async_benchmark.pedantic(target, rounds=200, warmup_rounds=20) + finally: + await settings.storage_provider.close() + shutil.rmtree(tmpdir, ignore_errors=True) + + +@pytest.mark.asyncio +async def test_benchmark_group_matches_by_type(async_benchmark): + """Benchmark grouping accumulated matches by knowledge type.""" + settings = make_settings() + tmpdir = tempfile.mkdtemp() + db_path = os.path.join(tmpdir, "bench_group.db") + + transcript = await create_indexed_transcript(db_path, settings, 200) + semref_index = transcript.semantic_ref_index + best_term, best_count = await find_best_term(semref_index) + print(f"\nBenchmarking group_matches_by_type: term '{best_term}' ({best_count} matches)") + + scored_refs = await semref_index.lookup_term(best_term) + accumulator = SemanticRefAccumulator() + accumulator.add_term_matches( + Term(text=best_term), scored_refs, is_exact_match=True + ) + + async def target(): + await accumulator.group_matches_by_type(transcript.semantic_refs) + + try: + await async_benchmark.pedantic(target, rounds=200, warmup_rounds=20) + finally: + await settings.storage_provider.close() + shutil.rmtree(tmpdir, ignore_errors=True) + +@pytest.mark.asyncio +async def test_benchmark_get_matches_in_scope(async_benchmark): + """Benchmark filtering accumulated matches by range scope.""" + settings = make_settings() + tmpdir = tempfile.mkdtemp() + db_path = os.path.join(tmpdir, "bench_scope.db") + + transcript = await create_indexed_transcript(db_path, settings, 200) + semref_index = transcript.semantic_ref_index + best_term, best_count = await find_best_term(semref_index) + print(f"\nBenchmarking get_matches_in_scope: term '{best_term}' ({best_count} matches)") + + scored_refs = await semref_index.lookup_term(best_term) + accumulator = SemanticRefAccumulator() + accumulator.add_term_matches( + Term(text=best_term), scored_refs, is_exact_match=True + ) + + scope = make_scope_first_half(200) + + async def target(): + await accumulator.get_matches_in_scope(transcript.semantic_refs, scope) + + try: + await async_benchmark.pedantic(target, rounds=200, warmup_rounds=20) + finally: + await settings.storage_provider.close() + shutil.rmtree(tmpdir, ignore_errors=True) + + +@pytest.mark.asyncio +async def test_benchmark_get_scored_semantic_refs_from_ordinals_iter(async_benchmark): + """Benchmark two-phase metadata filter + batch fetch for scored refs.""" + settings = make_settings() + tmpdir = tempfile.mkdtemp() + db_path = os.path.join(tmpdir, "bench_scored.db") + + transcript = await create_indexed_transcript(db_path, settings, 200) + semref_index = transcript.semantic_ref_index + best_term, best_count = await find_best_term(semref_index) + print( + f"\nBenchmarking get_scored_semantic_refs_from_ordinals_iter: " + f"term '{best_term}' ({best_count} matches), filter=entity" + ) + + scored_refs = await semref_index.lookup_term(best_term) + assert scored_refs is not None + + async def target(): + await get_scored_semantic_refs_from_ordinals_iter( + transcript.semantic_refs, scored_refs, "entity" + ) + + try: + await async_benchmark.pedantic(target, rounds=200, warmup_rounds=20) + finally: + await settings.storage_provider.close() shutil.rmtree(tmpdir, ignore_errors=True) From 163ef51e335d42760b798c744302e9600af603d1 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Fri, 10 Apr 2026 04:54:11 -0500 Subject: [PATCH 18/20] Speed up scope-filtering: bisect in contains_range, inline tuple comparisons - Use bisect_right with key=start in TextRangeCollection.contains_range to skip O(n) linear scan (O(log n) for non-overlapping point ranges) - Replace TextLocation allocations in TextRange __eq__/__lt__/__contains__ with a shared _effective_end returning tuples - Skip pydantic validation in get_metadata_multiple by constructing TextLocation/TextRange directly from JSON --- src/typeagent/knowpro/collections.py | 16 +++++---- src/typeagent/knowpro/interfaces_core.py | 36 +++++++-------------- src/typeagent/storage/sqlite/collections.py | 33 +++++++++++++++---- 3 files changed, 48 insertions(+), 37 deletions(-) diff --git a/src/typeagent/knowpro/collections.py b/src/typeagent/knowpro/collections.py index 6ff7b557..d7c07b19 100644 --- a/src/typeagent/knowpro/collections.py +++ b/src/typeagent/knowpro/collections.py @@ -526,12 +526,16 @@ def add_ranges(self, text_ranges: "list[TextRange] | TextRangeCollection") -> No self.add_range(text_range) def contains_range(self, inner_range: TextRange) -> bool: - # Since ranges are sorted by start, once we pass inner_range's start - # no further range can contain it. - for outer_range in self._ranges: - if outer_range.start > inner_range.start: - break - if inner_range in outer_range: + if not self._ranges: + return False + # Bisect on start only to find all ranges with start <= inner.start, + # then scan backwards — the most likely containing range has the + # largest start still <= inner's. + hi = bisect.bisect_right( + self._ranges, inner_range.start, key=lambda r: r.start + ) + for i in range(hi - 1, -1, -1): + if inner_range in self._ranges[i]: return True return False diff --git a/src/typeagent/knowpro/interfaces_core.py b/src/typeagent/knowpro/interfaces_core.py index cd9e885c..4dc8fc8e 100644 --- a/src/typeagent/knowpro/interfaces_core.py +++ b/src/typeagent/knowpro/interfaces_core.py @@ -255,32 +255,24 @@ def __repr__(self) -> str: else: return f"{self.__class__.__name__}({self.start}, {self.end})" + @staticmethod + def _effective_end(tr: "TextRange") -> tuple[int, int]: + """Return (message_ordinal, chunk_ordinal) for the effective end.""" + if tr.end is not None: + return (tr.end.message_ordinal, tr.end.chunk_ordinal) + return (tr.start.message_ordinal, tr.start.chunk_ordinal + 1) + def __eq__(self, other: object) -> bool: if not isinstance(other, TextRange): return NotImplemented - if self.start != other.start: return False - - # Get the effective end for both ranges - self_end = self.end or TextLocation( - self.start.message_ordinal, self.start.chunk_ordinal + 1 - ) - other_end = other.end or TextLocation( - other.start.message_ordinal, other.start.chunk_ordinal + 1 - ) - return self_end == other_end + return TextRange._effective_end(self) == TextRange._effective_end(other) def __lt__(self, other: Self) -> bool: if self.start != other.start: return self.start < other.start - self_end = self.end or TextLocation( - self.start.message_ordinal, self.start.chunk_ordinal + 1 - ) - other_end = other.end or TextLocation( - other.start.message_ordinal, other.start.chunk_ordinal + 1 - ) - return self_end < other_end + return TextRange._effective_end(self) < TextRange._effective_end(other) def __gt__(self, other: Self) -> bool: return other.__lt__(self) @@ -292,13 +284,9 @@ def __le__(self, other: Self) -> bool: return not other.__lt__(self) def __contains__(self, other: Self) -> bool: - other_end = other.end or TextLocation( - other.start.message_ordinal, other.start.chunk_ordinal + 1 - ) - self_end = self.end or TextLocation( - self.start.message_ordinal, self.start.chunk_ordinal + 1 - ) - return self.start <= other.start and other_end <= self_end + if not (self.start <= other.start): + return False + return TextRange._effective_end(other) <= TextRange._effective_end(self) def serialize(self) -> TextRangeData: return self.__pydantic_serializer__.to_python( # type: ignore diff --git a/src/typeagent/storage/sqlite/collections.py b/src/typeagent/storage/sqlite/collections.py index 911fff94..fe394dcb 100644 --- a/src/typeagent/storage/sqlite/collections.py +++ b/src/typeagent/storage/sqlite/collections.py @@ -356,14 +356,33 @@ async def get_metadata_multiple( ) rows = cursor.fetchall() rowdict = {r[0]: r for r in rows} - return [ - interfaces.SemanticRefMetadata( - ordinal=rowdict[o][0], - range=interfaces.TextRange.deserialize(json.loads(rowdict[o][1])), - knowledge_type=rowdict[o][2], + result = [] + for o in ordinals: + row = rowdict[o] + range_data = json.loads(row[1]) + start = range_data["start"] + end_data = range_data.get("end") + result.append( + interfaces.SemanticRefMetadata( + ordinal=row[0], + range=interfaces.TextRange( + start=interfaces.TextLocation( + start["messageOrdinal"], + start.get("chunkOrdinal", 0), + ), + end=( + interfaces.TextLocation( + end_data["messageOrdinal"], + end_data.get("chunkOrdinal", 0), + ) + if end_data + else None + ), + ), + knowledge_type=row[2], + ) ) - for o in ordinals - ] + return result async def append(self, item: interfaces.SemanticRef) -> None: cursor = self.db.cursor() From c3eff5ff3773bd40c4d4d0a351c481d344f70f92 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Fri, 10 Apr 2026 16:38:30 -0500 Subject: [PATCH 19/20] perf: Replace black with stdlib pprint for runtime formatting MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit black is only used at runtime in two cold formatting paths: - create_context_prompt() in answers.py (LLM debug context) - format_code()/pretty_print() in utils.py (developer terminal output) Both format Python data structures, which is exactly what pprint does. Replace black.format_str with pprint.pformat + ast.literal_eval, eliminating the runtime dependency entirely. Move black from dependencies to dev dependency-group — it remains available for make format/check but is no longer required by library consumers. --- pyproject.toml | 2 +- src/typeagent/aitools/utils.py | 21 +++++++++------------ src/typeagent/knowpro/answers.py | 4 ++-- tests/test_utils.py | 28 +++++++++++++++++++++++++--- 4 files changed, 37 insertions(+), 18 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 3cb25d7a..92a6d388 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,7 +30,6 @@ classifiers = [ ] dependencies = [ "azure-identity>=1.22.0", - "black>=25.12.0", "colorama>=0.4.6", "mcp[cli]>=1.12.1", "numpy>=2.2.6", @@ -84,6 +83,7 @@ known_local_folder = ["conftest"] dev = [ "azure-mgmt-authorization>=4.0.0", "azure-mgmt-keyvault>=12.1.1", + "black>=25.12.0", "coverage[toml]>=7.9.1", "google-api-python-client>=2.184.0", "google-auth-httplib2>=0.2.0", diff --git a/src/typeagent/aitools/utils.py b/src/typeagent/aitools/utils.py index eb21f9ec..51621f79 100644 --- a/src/typeagent/aitools/utils.py +++ b/src/typeagent/aitools/utils.py @@ -44,27 +44,24 @@ def timelog(label: str, verbose: bool = True): def pretty_print(obj: object, prefix: str = "", suffix: str = "") -> None: - """Pretty-print an object using black. + """Pretty-print an object using pprint.""" + import pprint - NOTE: Only works if its repr() is a valid Python expression. - """ - print(prefix + format_code(repr(obj)) + suffix) + line_width = min(200, shutil.get_terminal_size().columns) + print(pprint.pformat(obj, width=line_width)) def format_code(text: str, line_width=None) -> str: - """Format a block of code using black, then reindent to 2 spaces. + """Format a Python literal expression using pprint. - NOTE: The text must be a valid Python expression or code block. + NOTE: The text must be a valid Python literal expression (as produced by repr()). """ - import black + import ast + import pprint if line_width is None: - # 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.Mode(line_length=line_width) - ).rstrip() - return reindent(formatted_text) + return pprint.pformat(ast.literal_eval(text), width=line_width) def reindent(text: str) -> str: diff --git a/src/typeagent/knowpro/answers.py b/src/typeagent/knowpro/answers.py index 663d5804..9e984e18 100644 --- a/src/typeagent/knowpro/answers.py +++ b/src/typeagent/knowpro/answers.py @@ -125,12 +125,12 @@ def create_question_prompt(question: str) -> str: def create_context_prompt(context: AnswerContext) -> str: # TODO: Use a more compact representation of the context than JSON. - import black + import pprint prompt = [ "[ANSWER CONTEXT]", "===", - black.format_str(str(dictify(context)), mode=black.Mode(line_length=200)), + pprint.pformat(dictify(context), width=200), "===", ] return "\n".join(prompt) diff --git a/tests/test_utils.py b/tests/test_utils.py index 7f806f74..4c3edbe5 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -24,14 +24,36 @@ def test_timelog(): def test_pretty_print(): - # Use a simple object and check output is formatted by black obj = {"a": 1} buf = StringIO() with redirect_stdout(buf): utils.pretty_print(obj) out = buf.getvalue() - # Should be valid Python and contain the dict - assert out == '{"a": 1}\n', out + assert out == "{'a': 1}\n", out + + +def test_pretty_print_nested(): + obj = {"b": [1, 2], "a": {"nested": True}} + buf = StringIO() + with redirect_stdout(buf): + utils.pretty_print(obj) + out = buf.getvalue() + # pprint sorts keys and formats nested structures + assert "'a'" in out + assert "'nested'" in out + + +def test_format_code_simple(): + text = repr({"a": 1}) + result = utils.format_code(text) + assert result == "{'a': 1}" + + +def test_format_code_nested(): + obj = {"b": [1, 2, 3], "a": {"nested": True}} + result = utils.format_code(repr(obj)) + parsed = eval(result) + assert parsed == obj def test_load_dotenv(really_needs_auth): From 8f858c95e0beba548bd498990af7e2cd0b733f92 Mon Sep 17 00:00:00 2001 From: Kevin Turcios Date: Fri, 10 Apr 2026 18:02:34 -0500 Subject: [PATCH 20/20] perf: Defer query-time imports in conversation_base answers, search_query_schema, searchlang, and answer_response_schema are only used in the query() method. Move their imports from module level into query() and use TYPE_CHECKING + __future__.annotations for the type hints. These modules pull in search, query, and schema initialization that isn't needed when creating or indexing conversations. --- src/typeagent/knowpro/conversation_base.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/src/typeagent/knowpro/conversation_base.py b/src/typeagent/knowpro/conversation_base.py index 4b72ac99..bbb2f9f7 100644 --- a/src/typeagent/knowpro/conversation_base.py +++ b/src/typeagent/knowpro/conversation_base.py @@ -3,19 +3,17 @@ """Base class for conversations with incremental indexing support.""" +from __future__ import annotations + from dataclasses import dataclass from datetime import datetime, timezone -from typing import Generic, Self, TypeVar +from typing import TYPE_CHECKING, Generic, Self, TypeVar import typechat from . import ( - answer_response_schema, - answers, convknowledge, knowledge_schema as kplib, - search_query_schema, - searchlang, secindex, ) from ..aitools import model_adapters, utils @@ -35,6 +33,9 @@ Topic, ) +if TYPE_CHECKING: + from . import answer_response_schema, answers, search_query_schema, searchlang + TMessage = TypeVar("TMessage", bound=IMessage) @@ -350,6 +351,8 @@ async def query( >>> answer = await conv.query("What topics were discussed?") >>> print(answer) """ + from . import answer_response_schema, answers, search_query_schema, searchlang + # Create translators lazily (once per conversation instance) if self._query_translator is None: model = model_adapters.create_chat_model()