Merge pull request #120 from agentsea/adding-auth-to-solve

adding auth to the new agentdesktop instance in solve
agentsea · Nov 19, 2024 · 6054189 · 6054189
2 parents 9c8b7b4 + 3f309e2
commit 6054189
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 5 deletions.
diff --git a/surfkit/runtime/agent/kube.py b/surfkit/runtime/agent/kube.py
@@ -834,6 +834,13 @@ def run(
             labels=labels if labels else {},
         )
 
+    # def _get_headers_with_auth(self, token) -> dict:
+    #     """Helper to return headers with optional Authorization"""
+    #     headers = {
+    #         "Authorization": f"Bearer {token}"
+    #     }
+    #     return headers
+
     def solve_task(
         self,
         name: str,
@@ -849,6 +856,7 @@ def solve_task(
                 method="POST",
                 port=9090,
                 data=task.model_dump(),
+                # headers=self._get_headers_with_auth(task.task.auth_token)
             )
             logger.debug(f"Task posted with response: {status_code}, {response_text}")
 

diff --git a/surfkit/server/routes.py b/surfkit/server/routes.py
@@ -47,7 +47,9 @@ async def solve_task(
         background_tasks: BackgroundTasks,
         task_model: V1SolveTask,
     ):
-        logger.info(f"solving task: {task_model.model_dump()}")
+        logger.info(
+            f"solving task: {task_model.model_dump()} with user {current_user.email}"
+        )
         try:
             # TODO: we need to find a way to do this earlier but get status back
             mllm_router.check_model()
@@ -60,11 +62,11 @@ async def solve_task(
                 detail=f"failed to conect to LLM providers: {e} -- did you provide a valid key?",
             )
 
-        background_tasks.add_task(_solve_task, task_model)
+        background_tasks.add_task(_solve_task, task_model, current_user)
         logger.info("created background task...")
         return
 
-    def _solve_task(task_model: V1SolveTask):
+    def _solve_task(task_model: V1SolveTask, current_user: V1UserProfile):
         owner_id = task_model.task.owner_id
         if not owner_id:
             owner_id = "local"
@@ -81,9 +83,13 @@ def _solve_task(task_model: V1SolveTask):
             for Device in Agent.supported_devices():
                 if Device.type() == task_model.task.device.type:
                     logger.debug(f"found device: {task_model.task.device.model_dump()}")
-
+                    api_key = (
+                        current_user.token if current_user.token else task.auth_token
+                    )
+                    if api_key is None:
+                        logger.info("No Api key/token on Task or in Auth")
                     config = Device.connect_config_type()(
-                        **task_model.task.device.config  # type: ignore
+                        api_key=api_key, **task_model.task.device.config  # type: ignore
                     )
                     device = Device.connect(config=config)