fix: run vs non run endpoint

Author: guillaq

tests/e2e/no_schema_test.py

@@ -13,7 +13,7 @@ class SummarizeTaskOutput(BaseModel):
     summary_points: Optional[list[str]] = None
 
 
-@workflowai.agent(id="summarize")
+@workflowai.agent(id="summarize", model="gemini-1.5-flash-latest")
 async def summarize(task_input: SummarizeTaskInput) -> SummarizeTaskOutput: ...
 
 

tests/integration/run_test.py

@@ -17,9 +17,9 @@ class CityToCapitalTaskOutput(BaseModel):
     capital: str
 
 
-workflowai.init(api_key="test", url="http://localhost:8000")
+workflowai.init(api_key="test", url="https://run.workflowai.dev")
 
-_REGISTER_URL = "http://localhost:8000/v1/_/agents"
+_REGISTER_URL = "https://api.workflowai.dev/v1/_/agents"
 
 
 def _mock_register(httpx_mock: HTTPXMock, schema_id: int = 1, task_id: str = "city-to-capital", variant_id: str = "1"):
@@ -33,14 +33,14 @@ def _mock_register(httpx_mock: HTTPXMock, schema_id: int = 1, task_id: str = "ci
 def _mock_response(httpx_mock: HTTPXMock, task_id: str = "city-to-capital", capital: str = "Tokyo"):
     httpx_mock.add_response(
         method="POST",
-        url=f"http://localhost:8000/v1/_/agents/{task_id}/schemas/1/run",
+        url=f"https://run.workflowai.dev/v1/_/agents/{task_id}/schemas/1/run",
         json={"id": "123", "task_output": {"capital": capital}},
     )
 
 
 def _mock_stream(httpx_mock: HTTPXMock, task_id: str = "city-to-capital"):
     httpx_mock.add_response(
-        url=f"http://localhost:8000/v1/_/agents/{task_id}/schemas/1/run",
+        url=f"https://run.workflowai.dev/v1/_/agents/{task_id}/schemas/1/run",
         stream=IteratorStream(
             [
                 b'data: {"id":"1","task_output":{"capital":""}}\n\n',
@@ -53,7 +53,7 @@ def _mock_stream(httpx_mock: HTTPXMock, task_id: str = "city-to-capital"):
 
 def _check_request(request: Optional[Request], version: Any = "production", task_id: str = "city-to-capital"):
     assert request is not None
-    assert request.url == f"http://localhost:8000/v1/_/agents/{task_id}/schemas/1/run"
+    assert request.url == f"https://run.workflowai.dev/v1/_/agents/{task_id}/schemas/1/run"
     body = json.loads(request.content)
     assert body == {
         "task_input": {"city": "Hello"},

workflowai/core/client/_api.py

@@ -20,11 +20,16 @@ def __init__(self, endpoint: str, api_key: str, source_headers: Optional[dict[st
         self.api_key = api_key
         self.source_headers = source_headers or {}
 
+    def _get_endpoint(self, run: bool = False):
+        if run:
+            return self.endpoint
+        return self.endpoint.replace("https://run.", "https://api.")
+
     @asynccontextmanager
-    async def _client(self):
+    async def _client(self, run: bool = False):
         source_headers = self.source_headers or {}
         async with httpx.AsyncClient(
-            base_url=self.endpoint,
+            base_url=self._get_endpoint(run),
             headers={
                 "Authorization": f"Bearer {self.api_key}",
                 **source_headers,
@@ -48,18 +53,19 @@ async def get(self, path: str, returns: type[_R], query: Union[dict[str, Any], N
             return TypeAdapter(returns).validate_python(response.json())
 
     @overload
-    async def post(self, path: str, data: BaseModel, returns: type[_R]) -> _R: ...
+    async def post(self, path: str, data: BaseModel, returns: type[_R], run: bool = False) -> _R: ...
 
     @overload
-    async def post(self, path: str, data: BaseModel) -> None: ...
+    async def post(self, path: str, data: BaseModel, returns: None = None, run: bool = False) -> None: ...
 
     async def post(
         self,
         path: str,
         data: BaseModel,
         returns: Optional[type[_R]] = None,
+        run: bool = False,
     ) -> Optional[_R]:
-        async with self._client() as client:
+        async with self._client(run) as client:
             response = await client.post(
                 path,
                 content=data.model_dump_json(exclude_none=True),
@@ -155,8 +161,9 @@ async def stream(
         path: str,
         data: BaseModel,
         returns: type[_M],
+        run: bool = False,
     ) -> AsyncIterator[_M]:
-        async with self._client() as client, client.stream(
+        async with self._client(run=run) as client, client.stream(
             method,
             path,
             content=data.model_dump_json(exclude_none=True),

workflowai/core/client/agent.py

@@ -127,7 +127,7 @@ async def run(
         last_error = None
         while prepared_run.should_retry():
             try:
-                res = await self.api.post(prepared_run.route, prepared_run.request, returns=RunResponse)
+                res = await self.api.post(prepared_run.route, prepared_run.request, returns=RunResponse, run=True)
                 return res.to_domain(self.agent_id, prepared_run.schema_id, validator)
             except WorkflowAIError as e:  # noqa: PERF203
                 last_error = e
@@ -174,6 +174,7 @@ async def stream(
                     path=prepared_run.route,
                     data=prepared_run.request,
                     returns=RunResponse,
+                    run=True,
                 ):
                     yield chunk.to_domain(self.agent_id, prepared_run.schema_id, validator)
                 return