Merge pull request #8 from workflowai/guillaume/fix-stream-and-timeout

Fix stream errors and increase httpx timeout

Author: guillaq

tests/e2e/deploy_test.py

@@ -28,14 +28,18 @@ async def test_deploy_task(wai: workflowai.Client):
 
     # Run using the environment and the same input
     task_run2 = await wai.run(
-        task, task_input=CityToCapitalTaskInput(city="Osaka"), environment="dev",
+        task,
+        task_input=CityToCapitalTaskInput(city="Osaka"),
+        environment="dev",
     )
     # IDs will match since we are using cache
     assert task_run.id == task_run2.id
 
     # Run using the environment and a different input
     task_run3 = await wai.run(
-        task, task_input=CityToCapitalTaskInput(city="Toulouse"), environment="dev",
+        task,
+        task_input=CityToCapitalTaskInput(city="Toulouse"),
+        environment="dev",
     )
     assert task_run3.task_output.capital == "Paris"
     assert task_run3.id != task_run2.id

tests/e2e/stream_test.py

@@ -0,0 +1,43 @@
+from typing import Optional
+
+import pytest
+from pydantic import BaseModel
+
+import workflowai
+from workflowai.core.domain.task import Task
+
+
+class ImprovePromptTaskInput(BaseModel):
+    original_prompt: Optional[str] = None
+    prompt_input: Optional[str] = None
+    prompt_output: Optional[str] = None
+    user_evaluation: Optional[str] = None
+
+
+class ImprovePromptTaskOutput(BaseModel):
+    improved_prompt: Optional[str] = None
+    changelog: Optional[str] = None
+
+
+class ImprovePromptTask(Task[ImprovePromptTaskInput, ImprovePromptTaskOutput]):
+    id: str = "improve-prompt"
+    schema_id: int = 3
+    input_class: type[ImprovePromptTaskInput] = ImprovePromptTaskInput
+    output_class: type[ImprovePromptTaskOutput] = ImprovePromptTaskOutput
+
+
+@pytest.mark.skip("This hits the API")
+async def test_stream_task(wai: workflowai.Client):
+    task = ImprovePromptTask()
+
+    task_input = ImprovePromptTaskInput(
+        original_prompt="Say hello to the guest",
+        prompt_input='{"guest": "John", "language": "French"}',
+        prompt_output='{"greeting": "Hello John"}',
+        user_evaluation="Not in the right language",
+    )
+
+    streamed = await wai.run(task, task_input=task_input, stream=True, use_cache="never")
+    chunks = [chunk async for chunk in streamed]
+
+    assert len(chunks) > 1

workflowai/__init__.py

@@ -2,7 +2,7 @@
 
 from workflowai.core.client import Client as Client
 from workflowai.core.domain.cache_usage import CacheUsage as CacheUsage
-from workflowai.core.domain.errors import NotFoundError as NotFoundError
+from workflowai.core.domain.errors import WorkflowAIError as WorkflowAIError
 from workflowai.core.domain.llm_completion import LLMCompletion as LLMCompletion
 from workflowai.core.domain.task import Task as Task
 from workflowai.core.domain.task_evaluation import TaskEvaluation as TaskEvaluation

workflowai/core/client/api.py

@@ -1,7 +1,10 @@
 from typing import Any, AsyncIterator, Literal, Optional, TypeVar, Union, overload
 
 import httpx
-from pydantic import BaseModel, TypeAdapter
+from pydantic import BaseModel, TypeAdapter, ValidationError
+
+from workflowai.core.client.utils import split_chunks
+from workflowai.core.domain.errors import BaseError, ErrorResponse, WorkflowAIError
 
 # A type for return values
 _R = TypeVar("_R")
@@ -24,6 +27,7 @@ def _client(self) -> httpx.AsyncClient:
                 "Authorization": f"Bearer {self.api_key}",
                 **(self.source_headers or {}),
             },
+            timeout=120.0,
         )
         return client
 
@@ -84,21 +88,37 @@ async def delete(self, path: str) -> None:
             response = await client.delete(path)
             response.raise_for_status()
 
+    def _extract_error(self, data: Union[bytes, str], exception: Optional[Exception] = None) -> WorkflowAIError:
+        try:
+            res = ErrorResponse.model_validate_json(data)
+            return WorkflowAIError(res.error, task_run_id=res.task_run_id)
+        except ValidationError:
+            raise WorkflowAIError(
+                error=BaseError(
+                    message="Unknown error" if exception is None else str(exception),
+                    details={
+                        "raw": str(data),
+                    },
+                ),
+            ) from exception
+
     async def stream(
         self,
         method: Literal["GET", "POST"],
         path: str,
         data: BaseModel,
         returns: type[_M],
     ) -> AsyncIterator[_M]:
-        # TODO: error handling
-        async with self._client() as client:  # noqa: SIM117
-            async with client.stream(
-                method,
-                path,
-                content=data.model_dump_json(exclude_none=True),
-                headers={"Content-Type": "application/json"},
-            ) as response:
-                async for chunk in response.aiter_bytes():
-                    stripped = chunk.removeprefix(b"data: ").removesuffix(b"\n\n")
-                    yield returns.model_validate_json(stripped)
+        async with self._client() as client, client.stream(
+            method,
+            path,
+            content=data.model_dump_json(exclude_none=True),
+            headers={"Content-Type": "application/json"},
+        ) as response:
+            async for chunk in response.aiter_bytes():
+                payload = ""
+                try:
+                    for payload in split_chunks(chunk):
+                        yield returns.model_validate_json(payload)
+                except ValidationError as e:
+                    raise self._extract_error(payload, e) from None

workflowai/core/client/client.py

@@ -26,7 +26,7 @@
     TaskRunResponse,
 )
 from workflowai.core.domain.cache_usage import CacheUsage
-from workflowai.core.domain.errors import NotFoundError
+from workflowai.core.domain.errors import BaseError, WorkflowAIError
 from workflowai.core.domain.task import Task, TaskInput, TaskOutput
 from workflowai.core.domain.task_example import TaskExample
 from workflowai.core.domain.task_run import TaskRun
@@ -146,7 +146,13 @@ async def run(  # noqa: C901
                     return res.to_domain(task)
                 except HTTPStatusError as e:
                     if e.response.status_code == 404:
-                        raise NotFoundError("Task not found") from e
+                        raise WorkflowAIError(
+                            error=BaseError(
+                                status_code=404,
+                                code="not_found",
+                                message="Task not found",
+                            ),
+                        ) from e
                     retry_after = e.response.headers.get("Retry-After")
                     if retry_after:
                         try:
@@ -180,7 +186,7 @@ async def _stream():
                         yield task.output_class.model_construct(None, **chunk.task_output)
                 except HTTPStatusError as e:
                     if e.response.status_code == 404:
-                        raise NotFoundError("Task not found") from e
+                        raise WorkflowAIError(error=BaseError(message="Task not found")) from e
                     retry_after = e.response.headers.get("Retry-After")
 
                     if retry_after:
@@ -194,7 +200,7 @@ async def _stream():
                             except (TypeError, ValueError, OverflowError):
                                 delay = min(delay * 2, max_retry_delay / 1000)
                     elif e.response.status_code == 429 and delay < max_retry_delay / 1000:
-                            delay = min(delay * 2, max_retry_delay / 1000)
+                        delay = min(delay * 2, max_retry_delay / 1000)
                     await asyncio.sleep(delay)
                 retry_count += 1
 

workflowai/core/client/utils.py

@@ -0,0 +1,15 @@
+# Sometimes, 2 payloads are sent in a single message.
+# By adding the " at the end we more or less guarantee that
+# the delimiter is not withing a quoted string
+import re
+
+delimiter = re.compile(r'\}\n\ndata: \{"')
+
+
+def split_chunks(chunk: bytes):
+    start = 0
+    chunk_str = chunk.removeprefix(b"data: ").removesuffix(b"\n\n").decode()
+    for match in delimiter.finditer(chunk_str):
+        yield chunk_str[start : match.start() + 1]
+        start = match.end() - 2
+    yield chunk_str[start:]

workflowai/core/client/utils_test.py

@@ -0,0 +1,17 @@
+import pytest
+
+from workflowai.core.client.utils import split_chunks
+
+
+@pytest.mark.parametrize(
+    ("chunk", "expected"),
+    [
+        (b'data: {"foo": "bar"}\n\ndata: {"foo": "baz"}', ['{"foo": "bar"}', '{"foo": "baz"}']),
+        (
+            b'data: {"foo": "bar"}\n\ndata: {"foo": "baz"}\n\ndata: {"foo": "qux"}',
+            ['{"foo": "bar"}', '{"foo": "baz"}', '{"foo": "qux"}'],
+        ),
+    ],
+)
+def test_split_chunks(chunk: bytes, expected: list[bytes]):
+    assert list(split_chunks(chunk)) == expected

workflowai/core/domain/errors.py

@@ -1,12 +1,73 @@
-class WorkflowAIError(Exception):
-    def __init__(self, message: str):
-        self.message = message
-        super().__init__(message)
+from typing import Any, Literal, Optional, Union
+
+from pydantic import BaseModel
+
+ProviderErrorCode = Literal[
+    # Max number of tokens were exceeded in the prompt
+    "max_tokens_exceeded",
+    # The model failed to generate a response
+    "failed_generation",
+    # The model generated a response but it was not valid
+    "invalid_generation",
+    # The model returned an error that we currently do not handle
+    # The returned status code will match the provider status code and the entire
+    # provider response will be provided the error details.
+    #
+    # This error is intended as a fallback since we do not control what the providers
+    # return. We track this error on our end and the error should eventually
+    # be assigned a different status code
+    "unknown_provider_error",
+    # The provider returned a rate limit error
+    "rate_limit",
+    # The provider returned a server overloaded error
+    "server_overloaded",
+    # The requested provider does not support the model
+    "invalid_provider_config",
+    # The provider returned a 500
+    "provider_internal_error",
+    # The provider returned a 502 or 503
+    "provider_unavailable",
+    # The request timed out
+    "read_timeout",
+]
+
+ErrorCode = Union[
+    ProviderErrorCode,
+    Literal[
+        # The object was not found
+        "object_not_found",
+        # There are no configured providers supporting the requested model
+        # This error will never happen when using WorkflowAI keys
+        "no_provider_supporting_model",
+        # The requested provider does not support the model
+        "provider_does_not_support_model",
+        # The requested model does not support the requested generation mode
+        # (e-g a model that does not support images generation was sent an image)
+        "model_does_not_support_mode",
+        # Run properties are invalid, for example the model does not exist
+        "invalid_run_properties",
+        # An internal error occurred
+        "internal_error",
+        # The request was invalid
+        "bad_request",
+    ],
+    str,  # Using as a fallback to avoid validation error if an error code is added to the API
+]
 
-    def __str__(self):
-        return self.message
 
+class BaseError(BaseModel):
+    details: Optional[dict[str, Any]] = None
+    message: str
+    status_code: Optional[int] = None
+    code: Optional[ErrorCode] = None
 
-class NotFoundError(WorkflowAIError):
-    def __init__(self, message: str):
-        super().__init__(message)
+
+class ErrorResponse(BaseModel):
+    error: BaseError
+    task_run_id: Optional[str] = None
+
+
+class WorkflowAIError(Exception):
+    def __init__(self, error: BaseError, task_run_id: Optional[str] = None):
+        self.error = error
+        self.task_run_id = task_run_id