Merge branch 'main' into pierre-clean-docstring

Author: Pierre

.vscode/launch.json

@@ -1,6 +1,13 @@
 {
   "version": "0.2.0",
   "configurations": [
+    {
+      "name": "Python Debugger: Current File",
+      "type": "debugpy",
+      "request": "launch",
+      "program": "${file}",
+      "console": "integratedTerminal"
+    },
     {
       "name": "Python: Debug Tests",
       "type": "debugpy",

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "workflowai"
-version = "0.6.0.dev6"
+version = "0.6.0.dev8"
 description = ""
 authors = ["Guillaume Aquilina <[email protected]>"]
 readme = "README.md"

tests/e2e/no_schema_test.py

@@ -18,5 +18,13 @@ async def summarize(_: SummarizeTaskInput) -> SummarizeTaskOutput: ...
 
 
 async def test_summarize():
-    summarized = await summarize(SummarizeTaskInput(text="Hello, world!"))
+    summarized = await summarize(
+        SummarizeTaskInput(
+            text="""The first computer programmer was Ada Lovelace. She wrote the first algorithm
+intended to be processed by a machine in the 1840s. Her work was on Charles Babbage's
+proposed mechanical computer, the Analytical Engine. She is celebrated annually on Ada
+Lovelace Day, which promotes women in science and technology.""",
+        ),
+        use_cache="never",
+    )
     assert summarized.summary_points

tests/integration/conftest.py

@@ -1,5 +1,6 @@
 import json
-from typing import Any, Optional
+from collections.abc import Callable
+from typing import Any, Optional, Union
 from unittest.mock import patch
 
 import pytest
@@ -51,34 +52,70 @@ def mock_response(
             status_code=status_code,
         )
 
-    def mock_stream(self, task_id: str = "city-to-capital"):
+    def mock_stream(
+        self,
+        task_id: str = "city-to-capital",
+        outputs: Optional[list[dict[str, Any]]] = None,
+        run_id: str = "1",
+        metadata: Optional[dict[str, Any]] = None,
+    ):
+        outputs = outputs or [
+            {"capital": ""},
+            {"capital": "Tok"},
+            {"capital": "Tokyo"},
+        ]
+        if metadata is None:
+            metadata = {"cost_usd": 0.01, "duration_seconds": 10.1}
+
+        payloads = [{"id": run_id, "task_output": o} for o in outputs]
+
+        final_payload = {**payloads[-1], **metadata}
+        payloads.append(final_payload)
+        streams = [f"data: {json.dumps(p)}\n\n".encode() for p in payloads]
+
         self.httpx_mock.add_response(
             url=f"https://run.workflowai.dev/v1/_/agents/{task_id}/schemas/1/run",
-            stream=IteratorStream(
-                [
-                    b'data: {"id":"1","task_output":{"capital":""}}\n\n',
-                    b'data: {"id":"1","task_output":{"capital":"Tok"}}\n\ndata: {"id":"1","task_output":{"capital":"Tokyo"}}\n\n',  # noqa: E501
-                    b'data: {"id":"1","task_output":{"capital":"Tokyo"},"cost_usd":0.01,"duration_seconds":10.1}\n\n',
-                ],
-            ),
+            stream=IteratorStream(streams),
         )
 
+    def check_register(
+        self,
+        task_id: str = "city-to-capital",
+        input_schema: Optional[Union[dict[str, Any], Callable[[dict[str, Any]], None]]] = None,
+        output_schema: Optional[Union[dict[str, Any], Callable[[dict[str, Any]], None]]] = None,
+    ):
+        request = self.httpx_mock.get_request(url=self.REGISTER_URL)
+        assert request is not None
+        assert request.headers["Authorization"] == "Bearer test"
+        assert request.headers["Content-Type"] == "application/json"
+        assert request.headers["x-workflowai-source"] == "sdk"
+        assert request.headers["x-workflowai-language"] == "python"
+
+        body = json.loads(request.content)
+        assert body["id"] == task_id
+        if callable(input_schema):
+            input_schema(body["input_schema"])
+        else:
+            assert body["input_schema"] == input_schema or {"city": {"type": "string"}}
+        if callable(output_schema):
+            output_schema(body["output_schema"])
+        else:
+            assert body["output_schema"] == output_schema or {"capital": {"type": "string"}}
+
     def check_request(
         self,
         version: Any = "production",
         task_id: str = "city-to-capital",
         task_input: Optional[dict[str, Any]] = None,
         **matchers: Any,
     ):
+        if not matchers:
+            matchers = {"url": f"https://run.workflowai.dev/v1/_/agents/{task_id}/schemas/1/run"}
         request = self.httpx_mock.get_request(**matchers)
         assert request is not None
-        assert request.url == f"https://run.workflowai.dev/v1/_/agents/{task_id}/schemas/1/run"
         body = json.loads(request.content)
-        assert body == {
-            "task_input": task_input or {"city": "Hello"},
-            "version": version,
-            "stream": False,
-        }
+        assert body["task_input"] == task_input or {"city": "Hello"}
+        assert body["version"] == version
         assert request.headers["Authorization"] == "Bearer test"
         assert request.headers["Content-Type"] == "application/json"
         assert request.headers["x-workflowai-source"] == "sdk"

tests/integration/run_test.py

@@ -52,6 +52,21 @@ async def city_to_capital(task_input: CityToCapitalTaskInput) -> Run[CityToCapit
     test_client.check_request(version="staging")
 
 
+async def test_run_task_run_with_model(test_client: IntTestClient) -> None:
+    @workflowai.agent(schema_id=1, version="staging")
+    async def city_to_capital(task_input: CityToCapitalTaskInput) -> Run[CityToCapitalTaskOutput]: ...
+
+    test_client.mock_response()
+
+    task_input = CityToCapitalTaskInput(city="Hello")
+    with_run = await city_to_capital(task_input, model="gpt-4o-latest")
+
+    assert with_run.id == "123"
+    assert with_run.output.capital == "Tokyo"
+
+    test_client.check_request(version={"model": "gpt-4o-latest"})
+
+
 async def test_stream_task_run(test_client: IntTestClient) -> None:
     @workflowai.agent(schema_id=1)
     def city_to_capital(task_input: CityToCapitalTaskInput) -> AsyncIterator[CityToCapitalTaskOutput]: ...
@@ -112,27 +127,23 @@ async def city_to_capital(task_input: CityToCapitalTaskInput) -> CityToCapitalTa
         "input_schema": {
             "properties": {
                 "city": {
-                    "title": "City",
                     "type": "string",
                 },
             },
             "required": [
                 "city",
             ],
-            "title": "CityToCapitalTaskInput",
             "type": "object",
         },
         "output_schema": {
             "properties": {
                 "capital": {
-                    "title": "Capital",
                     "type": "string",
                 },
             },
             "required": [
                 "capital",
             ],
-            "title": "CityToCapitalTaskOutput",
             "type": "object",
         },
     }