Make timeouts configurable in semantic LLM operations

YoungVor · YoungVor · commit cfce81bb5d96 · 2025-11-03T12:32:17.000-08:00
diff --git a/src/fenic/_backends/local/semantic_operators/analyze_sentiment.py b/src/fenic/_backends/local/semantic_operators/analyze_sentiment.py
@@ -134,6 +134,7 @@ def __init__(
         model: LanguageModel,
         temperature: float,
         model_alias: Optional[ResolvedModelAlias] = None,
+        request_timeout: Optional[float] = None,
     ):
         super().__init__(
             input,
@@ -145,6 +146,7 @@ def __init__(
                     temperature=temperature,
                     response_format=SENTIMENT_ANALYSIS_FORMAT,
                     model_profile=model_alias.profile if model_alias else None,
+                    request_timeout=request_timeout,
                 ),
             ),
             EXAMPLES,
diff --git a/src/fenic/_backends/local/semantic_operators/base.py b/src/fenic/_backends/local/semantic_operators/base.py
@@ -78,6 +78,7 @@ def send_requests(
             response_format=self.inference_config.response_format,
             top_logprobs=self.inference_config.top_logprobs,
             model_profile=self.inference_config.model_profile,
+            request_timeout=self.inference_config.request_timeout,
         )
 
         completions = [
diff --git a/src/fenic/_backends/local/semantic_operators/classify.py b/src/fenic/_backends/local/semantic_operators/classify.py
@@ -43,6 +43,7 @@ def __init__(
         temperature: float,
         examples: Optional[ClassifyExampleCollection] = None,
         model_alias: Optional[ResolvedModelAlias] = None,
+        request_timeout: Optional[float] = None,
     ):
         self.classes = classes
         self.valid_labels = {class_def.label for class_def in classes}
@@ -59,6 +60,7 @@ def __init__(
                     temperature=temperature,
                     response_format=ResolvedResponseFormat.from_pydantic_model(self.output_model, generate_struct_type=False),
                     model_profile=model_alias.profile if model_alias else None,
+                    request_timeout=request_timeout,
                 ),
             ),
             examples,
diff --git a/src/fenic/_backends/local/semantic_operators/extract.py b/src/fenic/_backends/local/semantic_operators/extract.py
@@ -49,6 +49,7 @@ def __init__(
         max_output_tokens: int,
         temperature: float,
         model_alias: Optional[ResolvedModelAlias] = None,
+        request_timeout: Optional[float] = None,
     ):
         self.resolved_format = response_format
         super().__init__(
@@ -60,6 +61,7 @@ def __init__(
                     temperature=temperature,
                     response_format=response_format,
                     model_profile=model_alias.profile if model_alias else None,
+                    request_timeout=request_timeout,
                 ),
                 model=model,
             ),
diff --git a/src/fenic/_backends/local/semantic_operators/map.py b/src/fenic/_backends/local/semantic_operators/map.py
@@ -52,6 +52,7 @@ def __init__(
         model_alias: Optional[ResolvedModelAlias] = None,
         response_format: Optional[ResolvedResponseFormat] = None,
         examples: Optional[MapExampleCollection] = None,
+        request_timeout: Optional[float] = None,
     ):
         super().__init__(
             input,
@@ -63,6 +64,7 @@ def __init__(
                     response_format=response_format,
                     temperature=temperature,
                     model_profile=model_alias.profile if model_alias else None,
+                    request_timeout=request_timeout,
                 ),
             ),
             jinja_template=jinja2.Template(jinja_template),
diff --git a/src/fenic/_backends/local/semantic_operators/parse_pdf.py b/src/fenic/_backends/local/semantic_operators/parse_pdf.py
@@ -49,6 +49,7 @@ def __init__(
         describe_images: bool = False,
         model_alias: Optional[ResolvedModelAlias] = None,
         max_output_tokens: Optional[int] = None,
+        request_timeout: Optional[float] = None,
     ):
         self.page_separator = page_separator
         self.describe_images = describe_images
@@ -67,6 +68,7 @@ def __init__(
                     max_output_tokens=max_output_tokens,
                     temperature=1.0,  # Use a higher temperature so gemini flash models can handle complex table formatting.  For more info see the conversation here: https://discuss.ai.google.dev/t/gemini-2-0-flash-has-a-weird-bug/65119/26
                     model_profile=model_alias.profile if model_alias else None,
+                    request_timeout=request_timeout,
                 ),
             ),
             examples=None,  # PDF parsing doesn't use examples
diff --git a/src/fenic/_backends/local/semantic_operators/predicate.py b/src/fenic/_backends/local/semantic_operators/predicate.py
@@ -46,6 +46,7 @@ def __init__(
         temperature: float,
         examples: Optional[PredicateExampleCollection] = None,
         model_alias: Optional[ResolvedModelAlias] = None,
+        request_timeout: Optional[float] = None,
     ):
         super().__init__(
             input,
@@ -56,6 +57,7 @@ def __init__(
                   response_format=PREDICATE_FORMAT,
                   temperature=temperature,
                   model_profile=model_alias.profile if model_alias else None,
+                  request_timeout=request_timeout,
                 ),
                 model=model,
             ),
diff --git a/src/fenic/_backends/local/semantic_operators/summarize.py b/src/fenic/_backends/local/semantic_operators/summarize.py
@@ -45,6 +45,7 @@ def __init__(
         temperature: float,
         model: LanguageModel,
         model_alias: Optional[ResolvedModelAlias] = None,
+        request_timeout: Optional[float] = None,
     ):
         self.format = format
 
@@ -56,6 +57,7 @@ def __init__(
                     max_output_tokens=self.get_max_tokens(),
                     temperature=temperature,
                     model_profile=model_alias.profile if model_alias else None,
+                    request_timeout=request_timeout,
                 ),
                 model=model,
             ),
diff --git a/src/fenic/_backends/local/transpiler/expr_converter.py b/src/fenic/_backends/local/transpiler/expr_converter.py
@@ -547,6 +547,7 @@ def sem_map_fn(batch: pl.Series) -> pl.Series:
                 temperature=logical.temperature,
                 response_format=logical.response_format,
                 model_alias=logical.model_alias,
+                request_timeout=logical.request_timeout,
             ).execute()
 
         column_exprs = [self._convert_expr(expr) for expr in logical.exprs]
@@ -646,6 +647,7 @@ def sem_ext_fn(batch: pl.Series) -> pl.Series:
                 max_output_tokens=logical.max_tokens,
                 temperature=logical.temperature,
                 model_alias=logical.model_alias,
+                request_timeout=logical.request_timeout,
             ).execute()
 
         return self._convert_expr(logical.expr).map_batches(
@@ -666,6 +668,7 @@ def sem_pred_fn(batch: pl.Series) -> pl.Series:
                 examples=logical.examples,
                 temperature=logical.temperature,
                 model_alias=logical.model_alias,
+                request_timeout=logical.request_timeout,
             ).execute()
 
         column_exprs = [self._convert_expr(expr) for expr in logical.exprs]
@@ -688,6 +691,7 @@ def sem_classify_fn(batch: pl.Series) -> pl.Series:
                 temperature=logical.temperature,
                 examples=logical.examples,
                 model_alias=logical.model_alias,
+                request_timeout=logical.request_timeout,
             ).execute()
 
         return self._convert_expr(logical.expr).map_batches(
@@ -703,6 +707,7 @@ def sem_sentiment_fn(batch: pl.Series) -> pl.Series:
                 model=self.session_state.get_language_model(logical.model_alias),
                 temperature=logical.temperature,
                 model_alias=logical.model_alias,
+                request_timeout=logical.request_timeout,
             ).execute()
 
         return self._convert_expr(logical.expr).map_batches(
@@ -719,7 +724,7 @@ def sem_summarize_fn(batch: pl.Series) -> pl.Series:
                 format=logical.format,
                 temperature=logical.temperature,
                 model=self.session_state.get_language_model(logical.model_alias),
-
+                request_timeout=logical.request_timeout,
             ).execute()
 
         return self._convert_expr(logical.expr).map_batches(
@@ -736,6 +741,7 @@ def parse_pdf_fn(batch: pl.Series) -> pl.Series:
                 describe_images=logical.describe_images,
                 model_alias=logical.model_alias,
                 max_output_tokens=logical.max_output_tokens,
+                request_timeout=logical.request_timeout,
             ).execute()
 
         return self._convert_expr(logical.expr).map_batches(
diff --git a/src/fenic/_inference/anthropic/anthropic_provider.py b/src/fenic/_inference/anthropic/anthropic_provider.py
@@ -3,7 +3,9 @@
 import logging
 
 import anthropic
+import httpx
 
+from fenic._inference.request_utils import MAX_CLIENT_TIMEOUT
 from fenic.core._inference.model_provider import ModelProviderClass
 
 logger = logging.getLogger(__name__)
@@ -18,11 +20,11 @@ def name(self) -> str:
     
     def create_client(self):
         """Create an Anthropic sync client instance."""
-        return anthropic.Client()
+        return anthropic.Client(http_client=httpx.Client(timeout=MAX_CLIENT_TIMEOUT))
 
     def create_aio_client(self):
         """Create an Anthropic async client instance."""
-        return anthropic.AsyncAnthropic()
+        return anthropic.AsyncAnthropic(http_client=httpx.AsyncClient(timeout=MAX_CLIENT_TIMEOUT))
     
     async def validate_api_key(self) -> None:
         """Validate Anthropic API key by making a minimal completion request."""
diff --git a/src/fenic/_inference/google/google_provider.py b/src/fenic/_inference/google/google_provider.py
@@ -7,9 +7,13 @@
 
 from google import genai
 from google.genai._common import ExperimentalWarning as _GoogleExperimentalWarning
+from google.genai.types import HttpOptions
 
+from fenic._inference.request_utils import MAX_CLIENT_TIMEOUT
 from fenic.core._inference.model_provider import ModelProviderClass
 
+MAX_CLIENT_TIMEOUT_MS = MAX_CLIENT_TIMEOUT * 1000
+
 logger = logging.getLogger(__name__)
 
 
@@ -45,9 +49,9 @@ def name(self) -> str:
     def create_client(self):
         """Create a Google Developer client instance."""
         if "GEMINI_API_KEY" in os.environ:
-            return genai.Client(api_key=os.environ["GEMINI_API_KEY"])
+            return genai.Client(api_key=os.environ["GEMINI_API_KEY"], http_options=HttpOptions(timeout=MAX_CLIENT_TIMEOUT_MS))
         else:
-            return genai.Client()
+            return genai.Client(http_options=HttpOptions(timeout=MAX_CLIENT_TIMEOUT_MS))
 
 
 class GoogleVertexModelProvider(GoogleModelProvider):
@@ -62,7 +66,7 @@ def create_client(self):
 
         Passing `vertexai=True` automatically routes traffic through Vertex-AI if the environment is configured for it.
         """
-        return genai.Client(vertexai=True)
+        return genai.Client(vertexai=True, http_options=HttpOptions(timeout=MAX_CLIENT_TIMEOUT_MS))
 
 
 def _suppress_google_experimental_warnings() -> None:
diff --git a/src/fenic/_inference/language_model.py b/src/fenic/_inference/language_model.py
@@ -28,6 +28,7 @@ class InferenceConfiguration:
     top_logprobs: Optional[int] = None
     response_format: Optional[ResolvedResponseFormat] = None  # Resolved JSON schema
     model_profile: Optional[str] = None
+    request_timeout: Optional[float] = None  # Timeout in seconds for a single LLM request
 
 class LanguageModel:
     def __init__(self, client: ModelClient[FenicCompletionsRequest, FenicCompletionsResponse]):
@@ -49,6 +50,7 @@ def get_completions(
         top_logprobs: Optional[int] = None,
         model_profile: Optional[str] = None,
         operation_name: Optional[str] = None,
+        request_timeout: Optional[float] = None,
     ) -> list[Optional[FenicCompletionsResponse]]:
         # Create batch requests
         requests = []
@@ -69,6 +71,7 @@ def get_completions(
                 structured_output=response_format,
                 temperature=temperature_param,
                 model_profile=model_profile,
+                request_timeout=request_timeout,
             )
             requests.append(request)
 
diff --git a/src/fenic/_inference/model_client.py b/src/fenic/_inference/model_client.py
@@ -578,13 +578,16 @@ async def _process_single_request(self, queue_item: QueueItem[RequestT]):
         """
         try:
             try:
-                # TODO: make the timeout configurable, or dynamic based on request size.
+                # Use request_timeout from request if available, otherwise default to 120 seconds
+                timeout = 120.0
+                if hasattr(queue_item.request, 'request_timeout') and queue_item.request.request_timeout is not None:
+                    timeout = queue_item.request.request_timeout
                 maybe_response = await asyncio.wait_for(
                     self.make_single_request(queue_item.request),
-                    timeout=120.0,
+                    timeout=timeout,
                 )
             except asyncio.TimeoutError:
-                logger.warning(f"Request for model {self.model} in batch {queue_item.batch_id} timed out. Retrying.")
+                logger.warning(f"Request for model {self.model} in batch {queue_item.batch_id} timed out after {timeout} seconds. Retrying.")
                 await self.retry_queue.put(queue_item)
                 return
 
diff --git a/src/fenic/_inference/openai/openai_provider.py b/src/fenic/_inference/openai/openai_provider.py
@@ -2,8 +2,10 @@
 
 import logging
 
+import httpx
 from openai import AsyncOpenAI, OpenAI
 
+from fenic._inference.request_utils import MAX_CLIENT_TIMEOUT
 from fenic.core._inference.model_provider import ModelProviderClass
 
 logger = logging.getLogger(__name__)
@@ -18,11 +20,11 @@ def name(self) -> str:
 
     def create_client(self):
         """Create an OpenAI client instance."""
-        return OpenAI()
+        return OpenAI(http_client=httpx.AsyncClient(timeout=MAX_CLIENT_TIMEOUT))
 
     def create_aio_client(self):
         """Create an OpenAI async client instance."""
-        return AsyncOpenAI()
+        return AsyncOpenAI(http_client=httpx.AsyncClient(timeout=MAX_CLIENT_TIMEOUT))
 
     async def validate_api_key(self) -> None:
         """Validate OpenAI API key by listing models."""
diff --git a/src/fenic/_inference/openrouter/openrouter_provider.py b/src/fenic/_inference/openrouter/openrouter_provider.py
@@ -4,9 +4,11 @@
 from functools import cached_property
 from typing import Any, Dict, Optional
 
+import httpx
 import requests
 from openai import AsyncOpenAI, OpenAI
 
+from fenic._inference.request_utils import MAX_CLIENT_TIMEOUT
 from fenic.core._inference.model_catalog import (
     CompletionModelParameters,
     ModelProvider,
@@ -55,6 +57,7 @@ def client(self):
         return OpenAI(
             default_headers=self._headers,
             base_url=OPENROUTER_BASE_URL,
+            http_client=httpx.AsyncClient(timeout=MAX_CLIENT_TIMEOUT),
         )
 
     @cached_property
@@ -63,6 +66,7 @@ def aio_client(self):
         return AsyncOpenAI(
             default_headers=self._headers,
             base_url=OPENROUTER_BASE_URL,
+            http_client=httpx.AsyncClient(timeout=MAX_CLIENT_TIMEOUT),
         )
 
     def create_client(self):
diff --git a/src/fenic/_inference/request_utils.py b/src/fenic/_inference/request_utils.py
@@ -3,12 +3,32 @@
 import base64
 import hashlib
 import logging
+from typing import Annotated, Optional
 
 import fitz  # PyMuPDF
+from pydantic import BeforeValidator
 
 from fenic._inference.types import FenicCompletionsRequest, LMRequestFile
+from fenic.core.error import ValidationError
 
 logger = logging.getLogger(__name__)
+MAX_CLIENT_TIMEOUT = 600 # 10 minutes
+
+def validate_timeout(value: Optional[float]) -> Optional[float]:
+    """Validate timeout value using Pydantic constraints."""
+    if value is not None:
+        if value <= 0:
+            raise ValidationError("The `request_timeout` argument must be a positive number.")
+        if value > MAX_CLIENT_TIMEOUT:
+            raise ValidationError(f"The `request_timeout` argument can't be greater than the system's max timeout of {MAX_CLIENT_TIMEOUT} seconds.")
+    return value
+
+
+# Type alias for validated timeout parameter
+TimeoutParam = Annotated[
+    Optional[float],
+    BeforeValidator(validate_timeout),
+]
 
 def parse_openrouter_rate_limit_headers(
     headers: dict | None,
diff --git a/src/fenic/_inference/types.py b/src/fenic/_inference/types.py
@@ -69,6 +69,7 @@ class FenicCompletionsRequest:
     structured_output: Optional[ResolvedResponseFormat]  # Resolved JSON schema
     temperature: Optional[float]
     model_profile: Optional[str] = None
+    request_timeout: Optional[float] = None  # Timeout in seconds for a single LLM request
 
 @dataclass
 class FenicEmbeddingsRequest:
diff --git a/src/fenic/api/dataframe/semantic_extensions.py b/src/fenic/api/dataframe/semantic_extensions.py
diff --git a/src/fenic/api/functions/semantic.py b/src/fenic/api/functions/semantic.py
diff --git a/src/fenic/core/_logical_plan/expressions/semantic.py b/src/fenic/core/_logical_plan/expressions/semantic.py
diff --git a/src/fenic/core/_logical_plan/plans/join.py b/src/fenic/core/_logical_plan/plans/join.py
diff --git a/src/fenic/core/_logical_plan/plans/transform.py b/src/fenic/core/_logical_plan/plans/transform.py
diff --git a/tests/_backends/local/dataframe/test_semantic_join.py b/tests/_backends/local/dataframe/test_semantic_join.py
diff --git a/tests/_backends/local/functions/test_semantic_map.py b/tests/_backends/local/functions/test_semantic_map.py
diff --git a/tools/long_running_semantic_requests.py b/tools/long_running_semantic_requests.py

Original file line number	Diff line number	Diff line change
`@@ -78,6 +78,7 @@ def send_requests(`
`78`	`78`	`response_format=self.inference_config.response_format,`
`79`	`79`	`top_logprobs=self.inference_config.top_logprobs,`
`80`	`80`	`model_profile=self.inference_config.model_profile,`
	`81`	`+ request_timeout=self.inference_config.request_timeout,`
`81`	`82`	`)`
`82`	`83`
`83`	`84`	`completions = [`