black

HumanSignal · Oct 25, 2024 · a722106 · a722106
1 parent 136bc94
commit a722106
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 7 deletions.
diff --git a/adala/runtimes/_litellm.py b/adala/runtimes/_litellm.py
@@ -565,7 +565,8 @@ def get_cost_estimate(
     ) -> CostEstimate:
         try:
             user_prompts = [
-                partial_str_format(prompt, **substitution) for substitution in substitutions
+                partial_str_format(prompt, **substitution)
+                for substitution in substitutions
             ]
             cumulative_prompt_cost = 0
             cumulative_completion_cost = 0

diff --git a/tests/test_cost_estimation.py b/tests/test_cost_estimation.py
@@ -2,12 +2,8 @@
 import pytest
 from adala.runtimes._litellm import AsyncLiteLLMChatRuntime
 from adala.runtimes.base import CostEstimate
-from adala.agents import Agent
-from adala.skills import ClassificationSkill
 import numpy as np
 import os
-from fastapi.testclient import TestClient
-from server.app import app, CostEstimateRequest
 
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 
@@ -60,7 +56,7 @@ def test_estimate_cost_endpoint(client):
                 }
             },
         },
-        "prompt": '''
+        "prompt": """
         test {text}
 
         Use the following JSON format:
@@ -72,7 +68,7 @@ def test_estimate_cost_endpoint(client):
                 }
             ]
         }
-        ''',
+        """,
         "substitutions": [{"text": "test"}],
     }
     resp = client.post(