Add concurrency and fast-stop to generate_answers

.env-template

@@ -22,4 +22,12 @@ OUTLOOK_CLIENT_ID=user@email.com
 # Azure AD app registration client ID (GUID)
 OUTLOOK_APPLICATION_CLIENT_ID=...
 # Azure AD tenant ID (GUID)
-OUTLOOK_TENANT_ID=...
+OUTLOOK_TENANT_ID=...
+
+#######################
+# generate_answers() (src/typeagent/knowpro/answers.py)
+#######################
+# How many search results to answer concurrently. Default: 1 (sequential).
+#TYPEAGENT_ANSWER_CONCURRENCY=2
+# Stop starting new answers once a good answer is found. Default: false.
+#TYPEAGENT_ANSWER_FAST_STOP=true

src/typeagent/knowpro/answers.py

@@ -1,8 +1,10 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
+import asyncio
+import os
 from collections.abc import Iterable
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any
 
 import typechat
@@ -43,18 +45,56 @@ class AnswerContextOptions:
     debug: bool = False
 
 
+# Environment variables that override the AnswerGeneratorSettings defaults
+# below, so callers that don't pass `settings` explicitly can still opt in to
+# concurrency/fast-stop without a code change. See .env-template.
+CONCURRENCY_ENVVAR = "TYPEAGENT_ANSWER_CONCURRENCY"
+FAST_STOP_ENVVAR = "TYPEAGENT_ANSWER_FAST_STOP"
+
+
+@dataclass
+class AnswerGeneratorSettings:
+    """Settings controlling how generate_answers() processes search results.
+
+    Mirrors (the relevant subset of) TypeAgent's AnswerGeneratorSettings
+    in answerGenerator.ts.
+
+    Defaults preserve the pre-existing sequential, run-everything behavior
+    (concurrency=1, fast_stop=False) so that callers who don't pass
+    `settings` see no behavior change. Set the TYPEAGENT_ANSWER_CONCURRENCY
+    and/or TYPEAGENT_ANSWER_FAST_STOP environment variables to opt every
+    caller into concurrency/fast-stop without changing call sites, or pass
+    `settings` explicitly to override per call.
+    """
+
+    # How many search results to answer concurrently.
+    concurrency: int = field(
+        default_factory=lambda: int(os.getenv(CONCURRENCY_ENVVAR, "1"))
+    )
+    # Stop processing further search results once a good answer is found.
+    fast_stop: bool = field(
+        default_factory=lambda: os.getenv(FAST_STOP_ENVVAR, "false").lower() == "true"
+    )
+
+
 async def generate_answers(
     translator: typechat.TypeChatJsonTranslator[AnswerResponse],
     search_results: list[ConversationSearchResult],
     conversation: IConversation,
     orig_query_text: str,
     options: AnswerContextOptions | None = None,
-) -> tuple[list[AnswerResponse], AnswerResponse]:  # (all answers, combined answer)
-    all_answers: list[AnswerResponse] = []
+    settings: AnswerGeneratorSettings | None = None,
+) -> tuple[list[AnswerResponse], AnswerResponse]:
+    # Returns (answers, combined_answer). `answers` holds one AnswerResponse
+    # per search result that was actually run -- with settings.fast_stop
+    # enabled, results not yet started when a good answer is found are
+    # skipped, so `answers` may be shorter than `search_results`.
+    settings = settings or AnswerGeneratorSettings()
+    all_answers = await _generate_answers_concurrently(
+        translator, search_results, conversation, options, settings
+    )
     good_answers: list[str] = []
-    for result in search_results:
-        answer = await generate_answer(translator, result, conversation, options)
-        all_answers.append(answer)
+    for answer in all_answers:
         match answer.type:
             case "Answered":
                 assert answer.answer is not None, "Answered answer must not be None"
@@ -81,6 +121,36 @@ async def generate_answers(
     return all_answers, combined_answer
 
 
+async def _generate_answers_concurrently(
+    translator: typechat.TypeChatJsonTranslator[AnswerResponse],
+    search_results: list[ConversationSearchResult],
+    conversation: IConversation,
+    options: AnswerContextOptions | None,
+    settings: AnswerGeneratorSettings,
+) -> list[AnswerResponse]:
+    """Run generate_answer() over search_results, bounded by settings.concurrency.
+
+    If settings.fast_stop is set, search results that haven't started yet are
+    skipped as soon as a good answer has been found by another one -- results
+    already in flight are still allowed to finish.
+    """
+    semaphore = asyncio.Semaphore(max(1, settings.concurrency))
+    found_answer = asyncio.Event()
+
+    async def run_one(result: ConversationSearchResult) -> AnswerResponse | None:
+        async with semaphore:
+            if settings.fast_stop and found_answer.is_set():
+                return None
+            answer = await generate_answer(translator, result, conversation, options)
+            if settings.fast_stop and answer.type == "Answered" and answer.answer:
+                if answer.answer.strip():
+                    found_answer.set()
+            return answer
+
+    results = await asyncio.gather(*(run_one(result) for result in search_results))
+    return [answer for answer in results if answer is not None]
+
+
 async def generate_answer[TMessage: IMessage, TIndex: ITermToSemanticRefIndex](
     translator: typechat.TypeChatJsonTranslator[AnswerResponse],
     search_result: ConversationSearchResult,

tests/test_answers.py

@@ -1,20 +1,27 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
+import asyncio
 from collections.abc import AsyncGenerator
 
 import pytest
 import pytest_asyncio
 
+import typechat
+
+from typeagent.knowpro.answer_response_schema import AnswerResponse
 from typeagent.knowpro.answers import (
+    AnswerGeneratorSettings,
     facets_to_merged_facets,
+    generate_answers,
     get_enclosing_date_range_for_text_range,
     get_enclosing_text_range,
     merged_facets_to_facets,
     text_range_from_message_range,
 )
 from typeagent.knowpro.interfaces import TextLocation, TextRange
 from typeagent.knowpro.knowledge_schema import Facet
+from typeagent.knowpro.search import ConversationSearchResult
 
 from conftest import FakeMessage, FakeMessageCollection
 
@@ -190,3 +197,110 @@ def test_range(self) -> None:
     def test_invalid_raises(self) -> None:
         with pytest.raises(ValueError, match="Expect message ordinal range"):
             text_range_from_message_range(5, 2)
+
+
+# ---------------------------------------------------------------------------
+# generate_answers: concurrency and fast_stop
+# ---------------------------------------------------------------------------
+
+
+class FakeAnswerTranslator:
+    """Records concurrency/call info; returns canned responses in call order."""
+
+    def __init__(
+        self,
+        responses: list[AnswerResponse],
+        delay: float = 0.0,
+    ) -> None:
+        self._responses = responses
+        self._delay = delay
+        self.calls = 0
+        self.max_concurrent = 0
+        self._concurrent = 0
+
+    async def translate(self, request: str) -> typechat.Result[AnswerResponse]:
+        self._concurrent += 1
+        self.max_concurrent = max(self.max_concurrent, self._concurrent)
+        try:
+            if self._delay:
+                await asyncio.sleep(self._delay)
+            index = self.calls
+            self.calls += 1
+            return typechat.Success(self._responses[index])
+        finally:
+            self._concurrent -= 1
+
+
+def make_search_results(count: int) -> list[ConversationSearchResult]:
+    return [
+        ConversationSearchResult(
+            message_matches=[], knowledge_matches={}, raw_query_text=f"question {i}"
+        )
+        for i in range(count)
+    ]
+
+
+class TestGenerateAnswersConcurrency:
+    @pytest.mark.asyncio
+    async def test_respects_concurrency_limit(self) -> None:
+        translator = FakeAnswerTranslator(
+            responses=[
+                AnswerResponse(type="NoAnswer", why_no_answer="none") for _ in range(5)
+            ],
+            delay=0.01,
+        )
+        search_results = make_search_results(5)
+
+        all_answers, _combined = await generate_answers(
+            translator,  # type: ignore[arg-type]
+            search_results,
+            None,  # type: ignore[arg-type]  # conversation is unused: matches are empty
+            "orig question",
+            settings=AnswerGeneratorSettings(concurrency=2, fast_stop=False),
+        )
+
+        assert translator.calls == 5
+        assert len(all_answers) == 5
+        assert translator.max_concurrent <= 2
+
+    @pytest.mark.asyncio
+    async def test_fast_stop_skips_unstarted_results(self) -> None:
+        responses = [AnswerResponse(type="Answered", answer="the answer")] + [
+            AnswerResponse(type="NoAnswer", why_no_answer="none") for _ in range(4)
+        ]
+        translator = FakeAnswerTranslator(responses=responses)
+        search_results = make_search_results(5)
+
+        all_answers, combined = await generate_answers(
+            translator,  # type: ignore[arg-type]
+            search_results,
+            None,  # type: ignore[arg-type]
+            "orig question",
+            # concurrency=1 makes execution deterministically sequential, so
+            # the first result answers before any others are started.
+            settings=AnswerGeneratorSettings(concurrency=1, fast_stop=True),
+        )
+
+        assert translator.calls == 1
+        assert len(all_answers) == 1
+        assert combined.type == "Answered"
+        assert combined.answer == "the answer"
+
+    @pytest.mark.asyncio
+    async def test_fast_stop_false_processes_all_results(self) -> None:
+        responses = [AnswerResponse(type="Answered", answer="the answer")] + [
+            AnswerResponse(type="NoAnswer", why_no_answer="none") for _ in range(4)
+        ]
+        translator = FakeAnswerTranslator(responses=responses)
+        search_results = make_search_results(5)
+
+        all_answers, _combined = await generate_answers(
+            translator,  # type: ignore[arg-type]
+            search_results,
+            None,  # type: ignore[arg-type]
+            "orig question",
+            settings=AnswerGeneratorSettings(concurrency=1, fast_stop=False),
+        )
+
+        assert translator.calls == 5
+        assert len(all_answers) == 5