From 1e0ab6c71eac8060d7b0a42e29f5443e69c510a5 Mon Sep 17 00:00:00 2001 From: whiterabbit1983 Date: Tue, 12 Nov 2024 12:23:23 +0000 Subject: [PATCH] refactor: Lint agents-api (CI) --- .../agents_api/activities/task_steps/prompt_step.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/agents-api/agents_api/activities/task_steps/prompt_step.py b/agents-api/agents_api/activities/task_steps/prompt_step.py index 50d7fb312..5bbc0a1ae 100644 --- a/agents-api/agents_api/activities/task_steps/prompt_step.py +++ b/agents-api/agents_api/activities/task_steps/prompt_step.py @@ -8,11 +8,12 @@ from langchain_core.tools.convert import tool as tool_decorator from litellm import ChatCompletionMessageToolCall, Function, Message from litellm.types.utils import Choices, ModelResponse +from pydantic import BaseModel from temporalio import activity from temporalio.exceptions import ApplicationError -from pydantic import BaseModel from ...autogen.Tools import ( + BaseIntegrationDef, BraveIntegrationDef, BrowserbaseCompleteSessionIntegrationDef, BrowserbaseContextIntegrationDef, @@ -29,7 +30,6 @@ Tool, WeatherIntegrationDef, WikipediaIntegrationDef, - BaseIntegrationDef, ) from ...clients import ( litellm, # We dont directly import `acompletion` so we can mock it @@ -65,7 +65,9 @@ def _get_integration_arguments(tool: Tool): "remote_browser": RemoteBrowserIntegrationDef, } - integration: BaseIntegrationDef | dict[str, BaseIntegrationDef] = providers_map.get(tool.integration.provider) + integration: BaseIntegrationDef | dict[str, BaseIntegrationDef] = providers_map.get( + tool.integration.provider + ) if isinstance(integration, dict): integration: BaseIntegrationDef = integration.get(tool.integration.method)