chore(agents-api): lint

Author: Ahmad-mtos

agents-api/agents_api/clients/litellm.py

@@ -9,7 +9,7 @@
 from litellm import get_supported_openai_params
 from litellm.utils import CustomStreamWrapper, ModelResponse, get_valid_models
 
-from ..common.utils.usage import track_usage, track_embedding_usage
+from ..common.utils.usage import track_embedding_usage, track_usage
 from ..env import (
     embedding_dimensions,
     embedding_model_id,
@@ -79,7 +79,7 @@ async def acompletion(
     )
 
     response = patch_litellm_response(model_response)
-    
+
     # Track usage in database if we have a user ID (which should be the developer ID)
     user = settings.get("user")
     if user and isinstance(response, ModelResponse):
@@ -90,7 +90,7 @@ async def acompletion(
                 model=model,
                 messages=messages,
                 response=response,
-                custom_api_used= custom_api_key is not None,
+                custom_api_used=custom_api_key is not None,
                 metadata={"tags": kwargs.get("tags", [])},
             )
         except Exception as e:
@@ -134,7 +134,7 @@ async def aembedding(
         drop_params=True,
         **settings,
     )
-    
+
     # Track embedding usage if we have a user ID
     user = settings.get("user")
     if user:

agents-api/agents_api/common/utils/usage.py

@@ -2,16 +2,15 @@
 Utilities for tracking token usage and costs for LLM API calls.
 """
 
-from typing import Any, Optional
+from typing import Any
 from uuid import UUID
 
 from beartype import beartype
-from litellm import cost_per_token
-from litellm.utils import ModelResponse, _select_tokenizer as select_tokenizer
-from litellm.utils import token_counter
+from litellm.utils import ModelResponse, token_counter
 
 from ...queries.usage.create_usage_record import create_usage_record
 
+
 @beartype
 async def track_usage(
     *,
@@ -36,27 +35,28 @@ async def track_usage(
     Returns:
         None
     """
-    
+
     # Try to get token counts from response.usage
     if response.usage:
         prompt_tokens = response.usage.prompt_tokens
         completion_tokens = response.usage.completion_tokens
     else:
         # Calculate tokens manually if usage is not available
         prompt_tokens = token_counter(model=model, messages=messages)
-        
+
         # Calculate completion tokens from the response
-        completion_content = []
-        for choice in response.choices:
-            if hasattr(choice, "message") and choice.message:
-                if hasattr(choice.message, "content") and choice.message.content:
-                    completion_content.append({"content": choice.message.content})
-        
-        completion_tokens = token_counter(
-            model=model, 
-            messages=completion_content
-        ) if completion_content else 0
-    
+        completion_content = [
+            {"content": choice.message.content}
+            for choice in response.choices
+            if hasattr(choice, "message")
+            and choice.message
+            and hasattr(choice.message, "content")
+            and choice.message.content
+        ]
+
+        completion_tokens = (
+            token_counter(model=model, messages=completion_content) if completion_content else 0
+        )
 
     # Map the model name to the actual model name
     actual_model = model
@@ -99,13 +99,15 @@ async def track_embedding_usage(
     Returns:
         None
     """
-    
+
     # Try to get token count from response.usage
     if hasattr(response, "usage") and response.usage:
         prompt_tokens = response.usage.prompt_tokens
     else:
         # Calculate tokens manually if usage is not available
-        prompt_tokens = sum(token_counter(model=model, text=input_text) for input_text in inputs)
+        prompt_tokens = sum(
+            token_counter(model=model, text=input_text) for input_text in inputs
+        )
 
     # Map the model name to the actual model name
     actual_model = model
@@ -118,4 +120,4 @@ async def track_embedding_usage(
         completion_tokens=0,  # Embeddings don't have completion tokens
         custom_api_used=custom_api_used,
         metadata=metadata,
-    )
+    )

agents-api/agents_api/queries/__init__.py

@@ -17,5 +17,5 @@
 from . import sessions as sessions
 from . import tasks as tasks
 from . import tools as tools
+from . import usage as usage
 from . import users as users
-from . import usage as usage

agents-api/agents_api/queries/usage/__init__.py

@@ -17,4 +17,4 @@
 
 __all__ = [
     "create_usage_record",
-] 
+]

agents-api/agents_api/queries/usage/create_usage_record.py

@@ -71,30 +71,51 @@ async def create_usage_record(
     # Calculate cost based on token usage
     # For custom API keys, we still track usage but mark it as such
     total_cost = 0.0
-    
+
     if not custom_api_used:
         # Calculate cost using litellm's cost_per_token function
         try:
-            prompt_cost, completion_cost = cost_per_token(model, prompt_tokens=prompt_tokens, completion_tokens=completion_tokens)
+            prompt_cost, completion_cost = cost_per_token(
+                model, prompt_tokens=prompt_tokens, completion_tokens=completion_tokens
+            )
             total_cost = prompt_cost + completion_cost
-        except Exception as e:
+        except Exception:
             estimated = True
             fallback_pricing = {
-            # Meta Llama models
-            'meta-llama/llama-4-scout': {'api_request': 0.08/1000, 'api_response': 0.45/1000},
-            'meta-llama/llama-4-maverick': {'api_request': 0.19/1000, 'api_response': 0.85/1000},
-            'meta-llama/llama-4-maverick:free': {'api_request': 0.0/1000, 'api_response': 0.0/1000},
-            
-            # Qwen model
-            'qwen/qwen-2.5-72b-instruct': {'api_request': 0.7/1000, 'api_response': 0.7/1000},
-            
-            # Sao10k model
-            'sao10k/l3.3-euryale-70b': {'api_request': 0.7/1000, 'api_response': 0.8/1000},
-            'sao10k/l3.1-euryale-70b': {'api_request': 0.7/1000, 'api_response': 0.8/1000}
+                # Meta Llama models
+                "meta-llama/llama-4-scout": {
+                    "api_request": 0.08 / 1000,
+                    "api_response": 0.45 / 1000,
+                },
+                "meta-llama/llama-4-maverick": {
+                    "api_request": 0.19 / 1000,
+                    "api_response": 0.85 / 1000,
+                },
+                "meta-llama/llama-4-maverick:free": {
+                    "api_request": 0.0 / 1000,
+                    "api_response": 0.0 / 1000,
+                },
+                # Qwen model
+                "qwen/qwen-2.5-72b-instruct": {
+                    "api_request": 0.7 / 1000,
+                    "api_response": 0.7 / 1000,
+                },
+                # Sao10k model
+                "sao10k/l3.3-euryale-70b": {
+                    "api_request": 0.7 / 1000,
+                    "api_response": 0.8 / 1000,
+                },
+                "sao10k/l3.1-euryale-70b": {
+                    "api_request": 0.7 / 1000,
+                    "api_response": 0.8 / 1000,
+                },
             }
-        
+
             if model in fallback_pricing:
-                total_cost = fallback_pricing[model]['api_request'] * prompt_tokens + fallback_pricing[model]['api_response'] * completion_tokens
+                total_cost = (
+                    fallback_pricing[model]["api_request"] * prompt_tokens
+                    + fallback_pricing[model]["api_response"] * completion_tokens
+                )
             else:
                 print(f"No fallback pricing found for model {model}")
 
@@ -112,4 +133,4 @@ async def create_usage_record(
     return (
         usage_query,
         params,
-    ) 
+    )

agents-api/tests/test_usage_tracking.py

@@ -2,21 +2,19 @@
 Tests for usage tracking functionality.
 """
 
-import io
-from contextlib import redirect_stdout
-from decimal import Decimal
 from datetime import datetime
+from decimal import Decimal
 from unittest.mock import patch
 
+from agents_api.clients.pg import create_db_pool
+from agents_api.common.utils.usage import track_embedding_usage, track_usage
+from agents_api.queries.usage.create_usage_record import create_usage_record
+from litellm import cost_per_token
+from litellm.utils import Message, ModelResponse, Usage, token_counter
 from ward import test
 
-from agents_api.common.utils.usage import track_usage, track_embedding_usage
-from agents_api.queries.usage.create_usage_record import create_usage_record
-from litellm.utils import ModelResponse, Usage, Choices, Message
-from agents_api.clients.pg import create_db_pool
 from .fixtures import pg_dsn, test_developer_id
-from litellm import cost_per_token
-from litellm.utils import token_counter
+
 
 @test("query: create_usage_record creates a record with correct parameters")
 async def _(dsn=pg_dsn, developer_id=test_developer_id) -> None:
@@ -31,10 +29,10 @@ async def _(dsn=pg_dsn, developer_id=test_developer_id) -> None:
     assert len(response) == 1
     record = response[0]
     assert record["developer_id"] == developer_id
-    assert record["model"] == 'gpt-4o-mini'
+    assert record["model"] == "gpt-4o-mini"
     assert record["prompt_tokens"] == 100
     assert record["completion_tokens"] == 100
-    assert record["cost"] == Decimal('0.000075')
+    assert record["cost"] == Decimal("0.000075")
     assert record["estimated"] is False
     assert record["custom_api_used"] is False
     assert record["metadata"] == {}
@@ -60,7 +58,7 @@ async def _(dsn=pg_dsn, developer_id=test_developer_id) -> None:
         "meta-llama/llama-4-maverick:free",
         "qwen/qwen-2.5-72b-instruct",
         "sao10k/l3.3-euryale-70b",
-        "sao10k/l3.1-euryale-70b"
+        "sao10k/l3.1-euryale-70b",
     ]
     for model in models:
         response = await create_usage_record(
@@ -86,9 +84,11 @@ async def _(dsn=pg_dsn, developer_id=test_developer_id) -> None:
         connection_pool=pool,
     )
 
-    input_cost, completion_cost = cost_per_token("gpt-4o-mini", prompt_tokens=2041, completion_tokens=34198)
+    input_cost, completion_cost = cost_per_token(
+        "gpt-4o-mini", prompt_tokens=2041, completion_tokens=34198
+    )
     cost = input_cost + completion_cost
-    cost = Decimal(str(cost)).quantize(Decimal('0.000001'))
+    cost = Decimal(str(cost)).quantize(Decimal("0.000001"))
 
     assert len(response) == 1
     record = response[0]
@@ -125,13 +125,14 @@ async def _(dsn=pg_dsn, developer_id=test_developer_id) -> None:
 
     assert len(response) == 1
     record = response[0]
-    assert record["cost"] == Decimal('0.000000')
+    assert record["cost"] == Decimal("0.000000")
     assert record["estimated"] is True
 
+
 @test("query: create_usage_record with fallback pricing with model not in fallback pricing")
 async def _(dsn=pg_dsn, developer_id=test_developer_id) -> None:
     pool = await create_db_pool(dsn=dsn)
-    
+
     with patch("builtins.print") as mock_print:
         unknown_model = "unknown-model-name"
         response = await create_usage_record(
@@ -146,7 +147,7 @@ async def _(dsn=pg_dsn, developer_id=test_developer_id) -> None:
 
     assert len(response) == 1
     record = response[0]
-    assert record["cost"] == Decimal('0.000000')
+    assert record["cost"] == Decimal("0.000000")
     assert record["estimated"] is True
     assert expected_call == actual_call
 
@@ -171,23 +172,28 @@ async def _(developer_id=test_developer_id) -> None:
         assert call_args["prompt_tokens"] == 100
         assert call_args["completion_tokens"] == 100
 
-        
+
 @test("utils: track_usage without response.usage")
 async def _(developer_id=test_developer_id) -> None:
     with patch("agents_api.common.utils.usage.create_usage_record") as mock_create_usage_record:
         response = ModelResponse(
             usage=None,
-            choices=[{
-                "finish_reason": "stop",
-                "index": 0,
-                "message": Message(content="Hello, world!", role="assistant")
-            }]
+            choices=[
+                {
+                    "finish_reason": "stop",
+                    "index": 0,
+                    "message": Message(content="Hello, world!", role="assistant"),
+                }
+            ],
         )
         response.usage = None
         messages = [{"role": "user", "content": "Hello, world!"}]
 
         prompt_tokens = token_counter(model="gpt-4o-mini", messages=messages)
-        completion_tokens = token_counter(model="gpt-4o-mini", messages=[{"content": choice.message.content} for choice in response.choices])
+        completion_tokens = token_counter(
+            model="gpt-4o-mini",
+            messages=[{"content": choice.message.content} for choice in response.choices],
+        )
 
         await track_usage(
             developer_id=developer_id,
@@ -210,16 +216,16 @@ async def _(developer_id=test_developer_id) -> None:
                 completion_tokens=0,
             ),
         )
-        
+
         inputs = ["This is a test input for embedding"]
-        
+
         await track_embedding_usage(
             developer_id=developer_id,
             model="text-embedding-3-large",
             inputs=inputs,
             response=response,
         )
-        
+
         call_args = mock_create_usage_record.call_args[1]
         assert call_args["prompt_tokens"] == 150
         assert call_args["completion_tokens"] == 0
@@ -231,20 +237,23 @@ async def _(developer_id=test_developer_id) -> None:
     with patch("agents_api.common.utils.usage.create_usage_record") as mock_create_usage_record:
         response = ModelResponse()
         response.usage = None
-        
+
         inputs = ["First test input", "Second test input"]
-        
+
         # Calculate expected tokens manually
-        expected_tokens = sum(token_counter(model="text-embedding-3-large", text=input_text) for input_text in inputs)
-        
+        expected_tokens = sum(
+            token_counter(model="text-embedding-3-large", text=input_text)
+            for input_text in inputs
+        )
+
         await track_embedding_usage(
             developer_id=developer_id,
             model="text-embedding-3-large",
             inputs=inputs,
             response=response,
         )
-        
+
         call_args = mock_create_usage_record.call_args[1]
         assert call_args["prompt_tokens"] == expected_tokens
         assert call_args["completion_tokens"] == 0
-        assert call_args["model"] == "text-embedding-3-large"
+        assert call_args["model"] == "text-embedding-3-large"