fix: many miscellaneous fixes to integrations and agents API (#796)

> [!IMPORTANT] > Refines Remote Browser functionality by removing `wait_for_load`, updating models, and enhancing `PlaywrightActions` with error handling and screenshot capabilities. > > - **Behavior**: > - Removed `wait_for_load` action from `RemoteBrowserArguments` and `RemoteBrowserArgumentsUpdate`. > - Updated `task_to_spec()` in `tasks.py` to exclude `task_id`, `id`, and `agent_id`. > - Modified `create_task()` and `create_or_update_task()` to return query results directly. > - **Models**: > - Updated `RemoteBrowserSetup` to remove `RemoteBrowserSetupUpdate` and its fields. > - Refined `RemoteBrowserOutput` to include `output`, `error`, `base64_image`, and `system` fields. > - **Functions**: > - Enhanced `PlaywrightActions` with `_with_error_and_screenshot` decorator for error handling and screenshots. > - Added methods `_get_screen_size()`, `_set_screen_size()`, `_wait_for_load()`, `_execute_javascript()`, `_set_window_vars()`, `_get_mouse_coordinates()`, and `_get_element_coordinates()`. > - **Misc**: > - Removed `spec_to_task` import from `get_task.py` and `list_tasks.py`. > > <sup>This description was created by </sup>[<img alt="Ellipsis" src="https://img.shields.io/badge/Ellipsis-blue?color=175173">](https://www.ellipsis.dev?ref=julep-ai%2Fjulep&utm_source=github&utm_medium=referral)<sup> for c8ef196. It will automatically update as commits are pushed.</sup>  --------- Signed-off-by: Diwank Singh Tomer <[email protected]> Co-authored-by: creatorrr <[email protected]> Co-authored-by: HamadaSalhab <[email protected]> Co-authored-by: HamadaSalhab <[email protected]>
julep-ai · Nov 1, 2024 · 8ba2b5f · 8ba2b5f
1 parent 94ca66f
commit 8ba2b5f
Show file tree

Hide file tree

Showing 16 changed files with 284 additions and 229 deletions.
diff --git a/.gitignore b/.gitignore
@@ -1,3 +1,5 @@
+# BECAUSE DIWANK IS OBSESSED
+*gaga*
 .DS_Store
 *.swp
 ngrok*

diff --git a/agents-api/agents_api/autogen/Tools.py b/agents-api/agents_api/autogen/Tools.py
@@ -12,7 +12,6 @@
     BaseModel,
     ConfigDict,
     Field,
-    RootModel,
     StrictBool,
 )
 
@@ -1048,7 +1047,6 @@ class RemoteBrowserArguments(BaseModel):
         "cursor_position",
         "navigate",
         "refresh",
-        "wait_for_load",
     ]
     """
     The action to perform
@@ -1089,7 +1087,6 @@ class RemoteBrowserArgumentsUpdate(BaseModel):
             "cursor_position",
             "navigate",
             "refresh",
-            "wait_for_load",
         ]
         | None
     ) = None
@@ -1129,7 +1126,7 @@ class RemoteBrowserIntegrationDefUpdate(BaseIntegrationDefUpdate):
         populate_by_name=True,
     )
     provider: Literal["remote_browser"] = "remote_browser"
-    setup: RemoteBrowserSetupUpdate | None = None
+    setup: RemoteBrowserSetup | None = None
     method: Literal["perform_action"] = "perform_action"
     arguments: RemoteBrowserArgumentsUpdate | None = None
 
@@ -1139,28 +1136,6 @@ class RemoteBrowserSetup(BaseModel):
     The setup parameters for the remote browser
     """
 
-    model_config = ConfigDict(
-        populate_by_name=True,
-    )
-    connect_url: Annotated[str | None, Field(...)] = None
-    """
-    The connection URL for the remote browser
-    """
-    width: int | None = None
-    """
-    The width of the browser
-    """
-    height: int | None = None
-    """
-    The height of the browser
-    """
-
-
-class RemoteBrowserSetupUpdate(BaseModel):
-    """
-    The setup parameters for the remote browser
-    """
-
     model_config = ConfigDict(
         populate_by_name=True,
     )

diff --git a/agents-api/agents_api/common/protocol/tasks.py b/agents-api/agents_api/common/protocol/tasks.py
@@ -255,10 +255,7 @@ class StepOutcome(BaseModel):
 def task_to_spec(
     task: Task | CreateTaskRequest | UpdateTaskRequest | PatchTaskRequest, **model_opts
 ) -> TaskSpecDef | PartialTaskSpecDef:
-    task_data = task.model_dump(**model_opts)
-
-    if "task_id" in task_data:
-        task_data.pop("task_id")
+    task_data = task.model_dump(**model_opts, exclude={"task_id", "id", "agent_id"})
 
     if "tools" in task_data:
         del task_data["tools"]

diff --git a/agents-api/agents_api/models/docs/mmr.py b/agents-api/agents_api/models/docs/mmr.py
@@ -29,6 +29,9 @@ def _cosine_similarity(x: Matrix, y: Matrix) -> np.ndarray:
     if len(x) == 0 or len(y) == 0:
         return np.array([])
 
+    x = [xx for xx in x if xx is not None]
+    y = [yy for yy in y if yy is not None]
+
     x = np.array(x)
     y = np.array(y)
     if x.shape[1] != y.shape[1]:

diff --git a/agents-api/agents_api/models/task/create_task.py b/agents-api/agents_api/models/task/create_task.py
@@ -13,8 +13,9 @@
 
 from ...autogen.openapi_model import (
     CreateTaskRequest,
+    ResourceCreatedResponse,
 )
-from ...common.protocol.tasks import spec_to_task, task_to_spec
+from ...common.protocol.tasks import task_to_spec
 from ...common.utils.cozo import cozo_process_mutate_data
 from ...metrics.counters import increase_counter
 from ..utils import (
@@ -37,7 +38,16 @@
         TypeError: partialclass(HTTPException, status_code=400),
     }
 )
-@wrap_in_class(spec_to_task, one=True, _kind="inserted")
+@wrap_in_class(
+    ResourceCreatedResponse,
+    one=True,
+    transform=lambda d: {
+        "id": d["task_id"],
+        "jobs": [],
+        "created_at": d["created_at"],
+        **d,
+    },
+)
 @cozo_query
 @increase_counter("create_task")
 @beartype

diff --git a/agents-api/agents_api/models/task/get_task.py b/agents-api/agents_api/models/task/get_task.py
@@ -6,6 +6,7 @@
 from pycozo.client import QueryException
 from pydantic import ValidationError
 
+from ...common.protocol.tasks import spec_to_task
 from ..utils import (
     cozo_query,
     partialclass,
@@ -14,7 +15,6 @@
     verify_developer_owns_resource_query,
     wrap_in_class,
 )
-from .create_task import spec_to_task
 
 ModelT = TypeVar("ModelT", bound=Any)
 T = TypeVar("T")

diff --git a/agents-api/agents_api/models/task/list_tasks.py b/agents-api/agents_api/models/task/list_tasks.py
@@ -6,6 +6,7 @@
 from pycozo.client import QueryException
 from pydantic import ValidationError
 
+from ...common.protocol.tasks import spec_to_task
 from ..utils import (
     cozo_query,
     partialclass,
@@ -14,7 +15,6 @@
     verify_developer_owns_resource_query,
     wrap_in_class,
 )
-from .create_task import spec_to_task
 
 ModelT = TypeVar("ModelT", bound=Any)
 T = TypeVar("T")

diff --git a/agents-api/agents_api/routers/tasks/create_or_update_task.py b/agents-api/agents_api/routers/tasks/create_or_update_task.py
@@ -40,11 +40,9 @@ async def create_or_update_task(
     except ValidationError:
         pass
 
-    task = create_or_update_task_query(
+    return create_or_update_task_query(
         developer_id=x_developer_id,
         agent_id=agent_id,
         task_id=task_id,
         data=data,
     )
-
-    return ResourceUpdatedResponse(id=task.id, updated_at=task.updated_at, jobs=[])
diff --git a/agents-api/agents_api/routers/tasks/create_task.py b/agents-api/agents_api/routers/tasks/create_task.py
@@ -35,10 +35,8 @@ async def create_task(
     except ValidationError:
         pass
 
-    task = create_task_query(
+    return create_task_query(
         developer_id=x_developer_id,
         agent_id=agent_id,
         data=data,
     )
-
-    return ResourceCreatedResponse(id=task.id, created_at=task.created_at, jobs=[])
diff --git a/agents-api/tests/test_docs_routes.py b/agents-api/tests/test_docs_routes.py
@@ -1,4 +1,4 @@
-import asyncio
+import time
 
 from ward import skip, test
 
@@ -137,7 +137,7 @@ def _(make_request=make_request, agent=test_agent):
 # TODO: Fix this test. It fails sometimes and sometimes not.
 @test("route: search agent docs")
 async def _(make_request=make_request, agent=test_agent, doc=test_doc):
-    await asyncio.sleep(0.5)
+    time.sleep(0.5)
     search_params = dict(
         text=doc.content[0],
         limit=1,
@@ -161,7 +161,7 @@ async def _(make_request=make_request, agent=test_agent, doc=test_doc):
 @skip("Fails randomly on CI")
 @test("route: search user docs")
 async def _(make_request=make_request, user=test_user, doc=test_user_doc):
-    await asyncio.sleep(0.5)
+    time.sleep(0.5)
     search_params = dict(
         text=doc.content[0],
         limit=1,
@@ -184,7 +184,7 @@ async def _(make_request=make_request, user=test_user, doc=test_user_doc):
 
 @test("route: search agent docs hybrid with mmr")
 async def _(make_request=make_request, agent=test_agent, doc=test_doc):
-    await asyncio.sleep(0.5)
+    time.sleep(0.5)
 
     EMBEDDING_SIZE = 1024
     search_params = dict(

diff --git a/integrations-service/integrations/autogen/Tools.py b/integrations-service/integrations/autogen/Tools.py
@@ -12,7 +12,6 @@
     BaseModel,
     ConfigDict,
     Field,
-    RootModel,
     StrictBool,
 )
 
@@ -1048,7 +1047,6 @@ class RemoteBrowserArguments(BaseModel):
         "cursor_position",
         "navigate",
         "refresh",
-        "wait_for_load",
     ]
     """
     The action to perform
@@ -1089,7 +1087,6 @@ class RemoteBrowserArgumentsUpdate(BaseModel):
             "cursor_position",
             "navigate",
             "refresh",
-            "wait_for_load",
         ]
         | None
     ) = None
@@ -1129,7 +1126,7 @@ class RemoteBrowserIntegrationDefUpdate(BaseIntegrationDefUpdate):
         populate_by_name=True,
     )
     provider: Literal["remote_browser"] = "remote_browser"
-    setup: RemoteBrowserSetupUpdate | None = None
+    setup: RemoteBrowserSetup | None = None
     method: Literal["perform_action"] = "perform_action"
     arguments: RemoteBrowserArgumentsUpdate | None = None
 
@@ -1139,28 +1136,6 @@ class RemoteBrowserSetup(BaseModel):
     The setup parameters for the remote browser
     """
 
-    model_config = ConfigDict(
-        populate_by_name=True,
-    )
-    connect_url: Annotated[str | None, Field(...)] = None
-    """
-    The connection URL for the remote browser
-    """
-    width: int | None = None
-    """
-    The width of the browser
-    """
-    height: int | None = None
-    """
-    The height of the browser
-    """
-
-
-class RemoteBrowserSetupUpdate(BaseModel):
-    """
-    The setup parameters for the remote browser
-    """
-
     model_config = ConfigDict(
         populate_by_name=True,
     )

diff --git a/integrations-service/integrations/models/execution.py b/integrations-service/integrations/models/execution.py
@@ -29,7 +29,7 @@
 from .brave import BraveSearchOutput
 from .browserbase import (
     BrowserbaseCompleteSessionOutput,
-    # BrowserbaseContextOutput,
+    BrowserbaseContextOutput,
     BrowserbaseCreateSessionOutput,
     BrowserbaseExtensionOutput,
     BrowserbaseGetSessionConnectUrlOutput,
@@ -50,7 +50,6 @@
     BraveSearchSetup,
     BrowserbaseSetup,
     RemoteBrowserSetup,
-    # WikipediaSetup (not used)
 ]
 
 ExecutionArguments = Union[
@@ -59,12 +58,12 @@
     EmailArguments,
     WikipediaSearchArguments,
     BraveSearchArguments,
-    BrowserbaseCompleteSessionArguments,
-    BrowserbaseContextArguments,
     BrowserbaseCreateSessionArguments,
     BrowserbaseGetSessionArguments,
-    BrowserbaseGetSessionLiveUrlsArguments,
     BrowserbaseGetSessionConnectUrlArguments,
+    BrowserbaseGetSessionLiveUrlsArguments,
+    BrowserbaseCompleteSessionArguments,
+    BrowserbaseContextArguments,
     BrowserbaseExtensionArguments,
     BrowserbaseListSessionsArguments,
     RemoteBrowserArguments,
@@ -77,14 +76,14 @@
     WikipediaSearchOutput,
     BraveSearchOutput,
     BrowserbaseCreateSessionOutput,
-    BrowserbaseCompleteSessionOutput,
-    # BrowserbaseContextOutput,
-    BrowserbaseExtensionOutput,
     BrowserbaseGetSessionOutput,
     BrowserbaseGetSessionConnectUrlOutput,
     BrowserbaseGetSessionLiveUrlsOutput,
-    RemoteBrowserOutput,
+    BrowserbaseCompleteSessionOutput,
+    BrowserbaseContextOutput,
+    BrowserbaseExtensionOutput,
     BrowserbaseListSessionsOutput,
+    RemoteBrowserOutput,
 ]
 
 

diff --git a/integrations-service/integrations/models/remote_browser.py b/integrations-service/integrations/models/remote_browser.py
@@ -1,9 +1,12 @@
-from typing import Any
-
 from pydantic import Field
 
 from .base_models import BaseOutput
 
 
 class RemoteBrowserOutput(BaseOutput):
-    result: Any = Field(..., description="The result of the action")
+    output: str | None = Field(None, description="The output of the action")
+    error: str | None = Field(None, description="The error of the action")
+    base64_image: str | None = Field(
+        None, description="The base64 encoded image of the action"
+    )
+    system: str | None = Field(None, description="The system output of the action")