UN-4137 [FIX] Stream LLM completions under the hood so long generations complete instead of timing out

unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/anthropic.json

@@ -45,7 +45,13 @@
       "multipleOf": 1,
       "title": "Timeout",
       "default": 900,
-      "description": "Timeout in seconds"
+      "description": "Timeout in seconds. Replies are streamed, so this bounds the wait for the first token and any silence between chunks, not the length of the whole reply."
+    },
+    "enable_streaming": {
+      "type": "boolean",
+      "title": "Enable Streaming",
+      "default": true,
+      "description": "Stream the model's reply and assemble it in the platform. Keeps the connection alive on long replies so a provider read timeout is not hit. Turn it off only for an endpoint that cannot stream."
     },
     "enable_thinking": {
       "type": "boolean",

unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/anyscale.json

@@ -49,6 +49,12 @@
       "title": "Max Retries",
       "default": 5,
       "description": "Maximum number of retries to attempt when a request fails."
+    },
+    "enable_streaming": {
+      "type": "boolean",
+      "title": "Enable Streaming",
+      "default": true,
+      "description": "Stream the model's reply and assemble it in the platform. Keeps the connection alive on long replies so a provider read timeout is not hit. Turn it off only for an endpoint that cannot stream."
     }
   }
 }

unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/azure.json

@@ -69,6 +69,12 @@
       "default": 900,
       "description": "Timeout in seconds"
     },
+    "enable_streaming": {
+      "type": "boolean",
+      "title": "Enable Streaming",
+      "default": true,
+      "description": "Stream the model's reply and assemble it in the platform. Keeps the connection alive on long replies so a provider read timeout is not hit. Turn it off only for an endpoint that cannot stream."
+    },
     "enable_reasoning": {
       "type": "boolean",
       "title": "Enable Reasoning",

unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/azure_ai.json

@@ -54,6 +54,12 @@
       "title": "Timeout",
       "default": 900,
       "description": "Request timeout in seconds."
+    },
+    "enable_streaming": {
+      "type": "boolean",
+      "title": "Enable Streaming",
+      "default": true,
+      "description": "Stream the model's reply and assemble it in the platform. Keeps the connection alive on long replies so a provider read timeout is not hit. Turn it off only for an endpoint that cannot stream."
     }
   }
 }

unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/bedrock.json

@@ -111,6 +111,12 @@
       "title": "Timeout",
       "default": 900,
       "description": "Timeout in seconds"
+    },
+    "enable_streaming": {
+      "type": "boolean",
+      "title": "Enable Streaming",
+      "default": true,
+      "description": "Stream the model's reply and assemble it in the platform. Keeps the connection alive on long replies so a provider read timeout is not hit. Turn it off only for an endpoint that cannot stream."
     }
   },
   "dependencies": {

unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/custom_openai.json

@@ -56,6 +56,12 @@
       "default": 900,
       "description": "Timeout in seconds."
     },
+    "enable_streaming": {
+      "type": "boolean",
+      "title": "Enable Streaming",
+      "default": true,
+      "description": "Stream the model's reply and assemble it in the platform. Keeps the connection alive on long replies so a provider read timeout is not hit. Turn it off only for an endpoint that cannot stream."
+    },
     "enable_reasoning": {
       "type": "boolean",
       "title": "Enable Reasoning",

unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json

@@ -46,6 +46,12 @@
       "default": 600,
       "description": "Timeout in seconds"
     },
+    "enable_streaming": {
+      "type": "boolean",
+      "title": "Enable Streaming",
+      "default": true,
+      "description": "Stream the model's reply and assemble it in the platform. Keeps the connection alive on long replies so a provider read timeout is not hit. Turn it off only for an endpoint that cannot stream."
+    },
     "max_retries": {
       "type": "number",
       "minimum": 0,

unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/minimax.json

@@ -70,6 +70,12 @@
       "default": 900,
       "description": "Timeout in seconds."
     },
+    "enable_streaming": {
+      "type": "boolean",
+      "title": "Enable Streaming",
+      "default": true,
+      "description": "Stream the model's reply and assemble it in the platform. Keeps the connection alive on long replies so a provider read timeout is not hit. Turn it off only for an endpoint that cannot stream."
+    },
     "enable_thinking": {
       "type": "boolean",
       "title": "Enable Thinking",

unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/mistral.json

@@ -49,6 +49,12 @@
       "default": 900,
       "description": "Timeout in seconds"
     },
+    "enable_streaming": {
+      "type": "boolean",
+      "title": "Enable Streaming",
+      "default": true,
+      "description": "Stream the model's reply and assemble it in the platform. Keeps the connection alive on long replies so a provider read timeout is not hit. Turn it off only for an endpoint that cannot stream."
+    },
     "enable_reasoning": {
       "type": "boolean",
       "title": "Enable Reasoning",

unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/nvidia_build.json

@@ -55,6 +55,12 @@
       "default": 900,
       "description": "Timeout in seconds."
     },
+    "enable_streaming": {
+      "type": "boolean",
+      "title": "Enable Streaming",
+      "default": true,
+      "description": "Stream the model's reply and assemble it in the platform. Keeps the connection alive on long replies so a provider read timeout is not hit. Turn it off only for an endpoint that cannot stream."
+    },
     "enable_reasoning": {
       "type": "boolean",
       "title": "Enable Reasoning",

unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/ollama.json

@@ -56,6 +56,12 @@
       "default": 3,
       "description": "The maximum number of times to retry a request if it fails."
     },
+    "enable_streaming": {
+      "type": "boolean",
+      "title": "Enable Streaming",
+      "default": true,
+      "description": "Stream the model's reply and assemble it in the platform. Keeps the connection alive on long replies so a provider read timeout is not hit. Turn it off only for an endpoint that cannot stream."
+    },
     "request_timeout": {
       "type": "number",
       "minimum": 0,

unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/openai.json

@@ -62,6 +62,12 @@
       "default": 900,
       "description": "Timeout in seconds"
     },
+    "enable_streaming": {
+      "type": "boolean",
+      "title": "Enable Streaming",
+      "default": true,
+      "description": "Stream the model's reply and assemble it in the platform. Keeps the connection alive on long replies so a provider read timeout is not hit. Turn it off only for an endpoint that cannot stream."
+    },
     "enable_reasoning": {
       "type": "boolean",
       "title": "Enable Reasoning",

unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/openrouter.json

@@ -55,6 +55,12 @@
       "default": 900,
       "description": "Timeout in seconds."
     },
+    "enable_streaming": {
+      "type": "boolean",
+      "title": "Enable Streaming",
+      "default": true,
+      "description": "Stream the model's reply and assemble it in the platform. Keeps the connection alive on long replies so a provider read timeout is not hit. Turn it off only for an endpoint that cannot stream."
+    },
     "enable_reasoning": {
       "type": "boolean",
       "title": "Enable Reasoning",

unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/vertex_ai.json

@@ -46,6 +46,12 @@
       "default": 5,
       "description": "Maximum number of retries to attempt when a request fails."
     },
+    "enable_streaming": {
+      "type": "boolean",
+      "title": "Enable Streaming",
+      "default": true,
+      "description": "Stream the model's reply and assemble it in the platform. Keeps the connection alive on long replies so a provider read timeout is not hit. Turn it off only for an endpoint that cannot stream."
+    },
     "max_tokens": {
       "type": "number",
       "minimum": 0,

unstract/sdk1/src/unstract/sdk1/llm.py

@@ -25,6 +25,7 @@
 from unstract.sdk1.utils.retry_utils import (
     acall_with_retry,
     call_with_retry,
+    collect_with_retry,
     is_retryable_litellm_error,
     iter_with_retry,
     pop_litellm_retry_kwargs,
@@ -302,6 +303,12 @@ def __init__(  # noqa: C901
                 or enable_prompt_caching
                 or is_prompt_caching_enabled()
             )
+            # Per-adapter streaming switch. Read from the raw metadata because
+            # the provider validators drop keys they do not model, so it never
+            # reaches litellm as a completion param.
+            self._enable_streaming = self._resolve_enable_streaming(
+                self._adapter_metadata
+            )
 
             # REF: https://docs.litellm.ai/docs/completion/input#translated-openai-params
             # supported = get_supported_openai_params(model=self.kwargs["model"],
@@ -496,6 +503,86 @@ def _build_messages(
             {"role": "user", "content": user_content},
         ]
 
+    # Streaming under the hood is on by default and a per-adapter choice
+    # (``enable_streaming`` in the adapter metadata, surfaced as "Enable
+    # Streaming" in every LLM adapter form). Adapters stored before the field
+    # existed carry no value and stream too; the switch exists to opt an
+    # endpoint out when it cannot stream.
+    #
+    # Why streaming: a non-streaming endpoint keeps the socket silent until the
+    # last token, so a reply that takes longer than the read timeout, or than
+    # an intermediary tolerates on an idle connection, never arrives (observed
+    # on Anthropic as ``litellm.Timeout`` after 900 s and 1800 s on generations
+    # the Anthropic console finishes in 16 minutes; Anthropic's own SDK refuses
+    # non-streaming requests that may exceed 10 minutes). With streaming,
+    # ``complete()`` collects the chunks and rebuilds the full response, so
+    # callers see the same shape as before.
+    _STREAM_BY_DEFAULT = True
+
+    @classmethod
+    def _resolve_enable_streaming(cls, adapter_metadata: Mapping[str, object]) -> bool:
+        """Adapter flag if stored, else the platform default (on)."""
+        flag = adapter_metadata.get("enable_streaming")
+        if flag is None:
+            return cls._STREAM_BY_DEFAULT
+        return bool(flag)
+
+    def _streams_under_the_hood(self) -> bool:
+        return self._enable_streaming
+
+    @staticmethod
+    def _chunk_has_content(chunk: object) -> bool:
+        """Whether a stream chunk carries generated text or reasoning.
+
+        Bookkeeping chunks (``message_start``, usage-only tails) do not count:
+        a failure before the first content chunk is still a failed request.
+        """
+        get = getattr(chunk, "get", None)
+        if not callable(get):
+            return False
+        choices = get("choices")
+        if not choices:
+            return False
+        delta = choices[0].get("delta") or {}
+        return bool(delta.get("content") or delta.get("reasoning_content"))
+
+    def _complete_via_stream(
+        self,
+        messages: list[dict[str, object]],
+        completion_kwargs: dict[str, object],
+        max_retries: int,
+    ) -> dict[str, object]:
+        """Run a streamed completion and rebuild the non-streaming response.
+
+        Retries follow ``collect_with_retry``: a failure before the first
+        content chunk is retried like a non-streaming call; a failure after
+        content started is raised as-is so a long generation is never
+        replayed. ``litellm.stream_chunk_builder`` reassembles content,
+        thinking blocks, ``finish_reason``, usage (including cache tokens)
+        and the provider response headers.
+        """
+        chunks = collect_with_retry(
+            lambda: litellm.completion(
+                messages=messages,
+                stream=True,
+                stream_options={"include_usage": True},
+                **completion_kwargs,
+            ),
+            max_retries=max_retries,
+            retry_predicate=is_retryable_litellm_error,
+            is_content=self._chunk_has_content,
+            description=self._get_adapter_info(),
+        )
+        response = litellm.stream_chunk_builder(chunks, messages=messages)
+        if response is None:
+            raise LLMError(
+                message=(
+                    f"Error from LLM adapter '{self._get_adapter_info()}': "
+                    "provider returned an empty stream"
+                )
+            )
+        return cast("dict[str, object]", response)
+
     @capture_metrics
     def complete(
         self,
@@ -541,12 +628,23 @@ def complete(
             max_retries = pop_litellm_retry_kwargs(
                 completion_kwargs, self._get_adapter_info()
             )
-            response: dict[str, object] = call_with_retry(
-                lambda: litellm.completion(messages=messages, **completion_kwargs),
-                max_retries=max_retries,
-                retry_predicate=is_retryable_litellm_error,
-                description=self._get_adapter_info(),
-            )
+            # A mocked completion (UNSTRACT_LLM_MOCK_RESPONSE, used by the e2e
+            # rig) never touches the network, so streaming buys nothing there,
+            # and litellm's mock reports fixed usage (10/20/30) only on the
+            # non-streaming path. The rig counts calls with that contract.
+            mocked = "mock_response" in completion_kwargs
+            response: dict[str, object]
+            if self._streams_under_the_hood() and not mocked:
+                response = self._complete_via_stream(
+                    messages, completion_kwargs, max_retries
+                )
+            else:
+                response = call_with_retry(
+                    lambda: litellm.completion(messages=messages, **completion_kwargs),
+                    max_retries=max_retries,
+                    retry_predicate=is_retryable_litellm_error,
+                    description=self._get_adapter_info(),
+                )
 
             response_text = response["choices"][0]["message"]["content"]
             finish_reason = response["choices"][0].get("finish_reason")

unstract/sdk1/src/unstract/sdk1/utils/retry_utils.py

@@ -274,6 +274,57 @@
             time.sleep(delay)
 
 
+def collect_with_retry[T](
+    fn: Callable[[], Iterable[T]],
+    *,
+    max_retries: int,
+    retry_predicate: Callable[[Exception], bool],
+    is_content: Callable[[T], bool],
+    description: str = "",
+    logger_instance: logging.Logger | None = None,
+) -> list[T]:
+    """Drain fn() into a list with retry. Only retries before content.
+
+    Streamed completions emit bookkeeping items (e.g. Anthropic
+    ``message_start``) before any content. A failure at that point is a
+    failed *request*, retried exactly as a non-streaming call would be. A
+    failure after ``is_content`` matched an item is a failed *generation*
+    and is raised immediately: replaying it would re-run a generation that
+    may already have consumed minutes. Items from a failed attempt are
+    discarded so a retry never duplicates them.
+    """
+    _validate_max_retries(max_retries)
+    log = logger_instance or logger
+    for attempt in range(max_retries + 1):
+        items: list[T] = []
+        has_content = False
+        gen: Iterable[T] | None = None
+        try:
+            # ``fn()`` itself can fail: litellm's streaming ``completion()``
+            # sends the request when called, so a 429/5xx/connection error
+            # is raised here, before any chunk is yielded. That is a failed
+            # request and must retry like the non-streaming path.
+            gen = fn()
+            for item in gen:
+                items.append(item)
+                has_content = has_content or is_content(item)
+            return items
+        except Exception as e:
+            # Release the in-flight HTTP/socket resources before retrying.
+            close = getattr(gen, "close", None)
+            if callable(close):
+                close()
+            if has_content:
+                raise
+            delay = _get_retry_delay(
+                e, attempt, max_retries, retry_predicate, description, log
+            )
+            if delay is None:
+                raise
+            time.sleep(delay)
+    raise RuntimeError("unreachable")  # for type-checker: loop always returns or raises
+
+
 def is_retryable_error(error: Exception) -> bool:
     """Check if a requests-library HTTP error should trigger a retry.