Merge pull request #202 from enoch3712/190-enhancement-timeout-para-f…

…or-llm-model

timout added with test

extract_thinker/llm.py

@@ -5,6 +5,7 @@
 
 class LLM:
     TEMPERATURE = 0  # Always zero for deterministic outputs (IDP)
+    TIMEOUT = 3000  # Timeout in milliseconds
 
     def __init__(self,
                  model: str,
@@ -28,6 +29,7 @@ def request(self, messages: List[Dict[str, str]], response_model: str) -> Any:
                 messages=messages,
                 response_model=response_model,
                 temperature=self.TEMPERATURE,
+                timeout=self.TIMEOUT,
             )
         else:
             response = self.client.chat.completions.create(
@@ -36,7 +38,8 @@ def request(self, messages: List[Dict[str, str]], response_model: str) -> Any:
                 temperature=self.TEMPERATURE,
                 response_model=response_model,
                 max_retries=1,
-                max_tokens=self.token_limit
+                max_tokens=self.token_limit,
+                timeout=self.TIMEOUT,
             )
 
         return response
@@ -54,4 +57,8 @@ def raw_completion(self, messages: List[Dict[str, str]]) -> str:
                 messages=messages,
                 max_tokens=self.token_limit
             )
-        return raw_response.choices[0].message.content
+        return raw_response.choices[0].message.content
+
+    def set_timeout(self, timeout_ms: int) -> None:
+        """Set the timeout value for LLM requests in milliseconds."""
+        self.TIMEOUT = timeout_ms

tests/test_extractor.py

@@ -367,3 +367,26 @@ def test_concatenation_handler():
         result_1.pages[0].content, 
         result_2.pages[0].content
     ), "Page contents are not semantically similar enough (threshold: 90%)"
+
+def test_llm_timeout():
+    # Arrange
+    test_file_path = os.path.join(cwd, "tests", "files", "invoice.pdf")
+
+    extractor = Extractor()
+    extractor.load_document_loader(DocumentLoaderPyPdf())
+
+    # Create LLM with very short timeout
+    llm = LLM("gpt-4o-mini")
+    llm.set_timeout(1)  # Set timeout to 1ms (extremely short to force timeout)
+    extractor.load_llm(llm)
+
+    # Act & Assert
+    with pytest.raises(Exception) as exc_info:
+        extractor.extract(test_file_path, InvoiceContract)
+
+    # Reset timeout to normal value
+    llm.set_timeout(3000)
+
+    # Verify normal operation works after reset
+    result = extractor.extract(test_file_path, InvoiceContract)
+    assert result is not None