From 27708e5ee57619445f16793ce4f791ee9449546a Mon Sep 17 00:00:00 2001
From: wirthual <wirthra@gmail.com>
Date: Thu, 23 Oct 2025 10:00:37 -0700
Subject: [PATCH] add retry to extract_knowledge_from_text

---
 test/test_knowledge.py         | 12 ++++++++++++
 typeagent/knowpro/knowledge.py | 15 +++++++++++++--
 2 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/test/test_knowledge.py b/test/test_knowledge.py
index d4ab882a..3c244560 100644
--- a/test/test_knowledge.py
+++ b/test/test_knowledge.py
@@ -15,6 +15,7 @@
 from typeagent.knowpro import convknowledge, kplib
 
 from fixtures import really_needs_auth
+from unittest.mock import AsyncMock
 
 
 class MockKnowledgeExtractor:
@@ -45,16 +46,27 @@ async def test_extract_knowledge_from_text(
     mock_knowledge_extractor: convknowledge.KnowledgeExtractor,
 ):
     """Test extracting knowledge from a single text input."""
+
+    mock_knowledge_extractor.extract = AsyncMock(
+        side_effect=mock_knowledge_extractor.extract
+    )
+
     result = await extract_knowledge_from_text(mock_knowledge_extractor, "test text", 3)
     assert isinstance(result, Success)
     assert result.value.topics[0] == "test text"
 
+    assert mock_knowledge_extractor.extract.call_count == 1
+
     failure_result = await extract_knowledge_from_text(
         mock_knowledge_extractor, "error", 3
     )
     assert isinstance(failure_result, Failure)
     assert failure_result.message == "Extraction failed"
 
+    assert (
+        mock_knowledge_extractor.extract.call_count == 5
+    )  # 1 for success | 1 + 3 retries for failures
+
 
 @pytest.mark.asyncio
 async def test_extract_knowledge_from_text_batch(
diff --git a/typeagent/knowpro/knowledge.py b/typeagent/knowpro/knowledge.py
index 900a76c8..af2ba130 100644
--- a/typeagent/knowpro/knowledge.py
+++ b/typeagent/knowpro/knowledge.py
@@ -6,6 +6,7 @@
 from . import convknowledge
 from . import kplib
 from .interfaces import IKnowledgeExtractor
+from typechat import Result, Failure, Success
 
 
 def create_knowledge_extractor(
@@ -25,8 +26,18 @@ async def extract_knowledge_from_text(
     max_retries: int,
 ) -> Result[kplib.KnowledgeResponse]:
     """Extract knowledge from a single text input with retries."""
-    # TODO: Add a retry mechanism to handle transient errors.
-    return await knowledge_extractor.extract(text)
+
+    if max_retries < 0:
+        max_retries = 0
+
+    attempt = 0
+    while attempt <= max_retries:
+        result = await knowledge_extractor.extract(text)
+        if isinstance(result, Success) or attempt == max_retries:
+            return result
+        attempt += 1
+
+    return Failure("No result returned after retries")
 
 
 async def extract_knowledge_from_text_batch(