diff --git a/agents-api/agents_api/common/utils/usage.py b/agents-api/agents_api/common/utils/usage.py index f528fcdfb..f2a62a38f 100644 --- a/agents-api/agents_api/common/utils/usage.py +++ b/agents-api/agents_api/common/utils/usage.py @@ -18,6 +18,11 @@ async def track_usage( model: str, messages: list[dict], response: ModelResponse, + execution_id: UUID | None = None, + transition_id: UUID | None = None, + session_id: UUID | None = None, + entry_id: UUID | None = None, + provider: str | None = None, custom_api_used: bool = False, metadata: dict[str, Any] = {}, connection_pool: Any = None, # This is for testing purposes @@ -30,6 +35,11 @@ async def track_usage( model (str): The model used for the API call. messages (list[dict]): The messages sent to the model. response (ModelResponse): The response from the LLM API call. + execution_id (UUID | None): Execution that triggered the call. + transition_id (UUID | None): Transition that invoked the call. + session_id (UUID | None): Session the call belongs to. + entry_id (UUID | None): Entry that prompted the call. + provider (str | None): LLM provider used. custom_api_used (bool): Whether a custom API key was used. metadata (dict): Additional metadata about the usage. @@ -68,6 +78,11 @@ async def track_usage( model=actual_model, prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, + execution_id=execution_id, + transition_id=transition_id, + session_id=session_id, + entry_id=entry_id, + provider=provider, custom_api_used=custom_api_used, metadata={ "request_id": response.id if hasattr(response, "id") else None, @@ -84,6 +99,11 @@ async def track_embedding_usage( model: str, inputs: list[str], response: Any, + execution_id: UUID | None = None, + transition_id: UUID | None = None, + session_id: UUID | None = None, + entry_id: UUID | None = None, + provider: str | None = None, custom_api_used: bool = False, metadata: dict[str, Any] = {}, ) -> None: @@ -95,6 +115,11 @@ async def track_embedding_usage( model (str): The model used for the embedding. inputs (list[str]): The inputs sent for embedding. response (Any): The response from the embedding API call. + execution_id (UUID | None): Execution that triggered the call. + transition_id (UUID | None): Transition that invoked the call. + session_id (UUID | None): Session the call belongs to. + entry_id (UUID | None): Entry that prompted the call. + provider (str | None): LLM provider used. custom_api_used (bool): Whether a custom API key was used. metadata (dict): Additional metadata about the usage. @@ -120,6 +145,11 @@ async def track_embedding_usage( model=actual_model, prompt_tokens=prompt_tokens, completion_tokens=0, # Embeddings don't have completion tokens + execution_id=execution_id, + transition_id=transition_id, + session_id=session_id, + entry_id=entry_id, + provider=provider, custom_api_used=custom_api_used, metadata=metadata, ) diff --git a/agents-api/agents_api/queries/usage/create_usage_record.py b/agents-api/agents_api/queries/usage/create_usage_record.py index 3842c8da2..6ca1bed21 100644 --- a/agents-api/agents_api/queries/usage/create_usage_record.py +++ b/agents-api/agents_api/queries/usage/create_usage_record.py @@ -96,6 +96,11 @@ model, prompt_tokens, completion_tokens, + execution_id, + transition_id, + session_id, + entry_id, + provider, cost, estimated, custom_api_used, @@ -106,10 +111,15 @@ $2, -- model $3, -- prompt_tokens $4, -- completion_tokens - $5, -- cost - $6, -- estimated - $7, -- custom_api_used - $8 -- metadata + $5, -- execution_id + $6, -- transition_id + $7, -- session_id + $8, -- entry_id + $9, -- provider + $10, -- cost + $11, -- estimated + $12, -- custom_api_used + $13 -- metadata ) RETURNING *; """ @@ -125,6 +135,12 @@ async def create_usage_record( model: str, prompt_tokens: int, completion_tokens: int, + execution_id: UUID | None = None, + transition_id: UUID | None = None, + session_id: UUID | None = None, + entry_id: UUID | None = None, + provider: str | None = None, + # AIDEV-NOTE: new parameters capture context for each usage record custom_api_used: bool = False, estimated: bool = False, metadata: dict[str, Any] | None = None, @@ -137,6 +153,11 @@ async def create_usage_record( model (str): The model used for the API call. prompt_tokens (int): Number of tokens in the prompt. completion_tokens (int): Number of tokens in the completion. + execution_id (UUID | None): Execution that triggered the call. + transition_id (UUID | None): Transition that invoked the call. + session_id (UUID | None): Session the call belongs to. + entry_id (UUID | None): Entry that prompted the call. + provider (str | None): LLM provider used. custom_api_used (bool): Whether a custom API key was used. estimated (bool): Whether the token count is estimated. metadata (dict | None): Additional metadata about the usage. @@ -173,6 +194,11 @@ async def create_usage_record( model, prompt_tokens, completion_tokens, + execution_id, + transition_id, + session_id, + entry_id, + provider, total_cost, estimated, custom_api_used, diff --git a/memory-store/README.md b/memory-store/README.md index 3422c92a9..69398d4ac 100644 --- a/memory-store/README.md +++ b/memory-store/README.md @@ -8,3 +8,7 @@ For creating a migration: `migrate -database "postgres://postgres:postgres@0.0.0.0:5432/postgres?sslmode=disable" -path migrations create -ext sql -seq -dir migrations switch_to_hypercore` + +### Usage table +New reference columns were added in migration `000042_usage_reference_fields`: +`execution_id`, `transition_id`, `session_id`, `entry_id`, and `provider`. diff --git a/memory-store/migrations/000042_usage_reference_fields.down.sql b/memory-store/migrations/000042_usage_reference_fields.down.sql new file mode 100644 index 000000000..8261f473a --- /dev/null +++ b/memory-store/migrations/000042_usage_reference_fields.down.sql @@ -0,0 +1,17 @@ +BEGIN; + +-- Remove reference columns from usage table +ALTER TABLE usage + DROP COLUMN IF EXISTS execution_id, + DROP COLUMN IF EXISTS transition_id, + DROP COLUMN IF EXISTS session_id, + DROP COLUMN IF EXISTS entry_id, + DROP COLUMN IF EXISTS provider; + +DROP INDEX IF EXISTS idx_usage_execution_id; +DROP INDEX IF EXISTS idx_usage_transition_id; +DROP INDEX IF EXISTS idx_usage_session_id; +DROP INDEX IF EXISTS idx_usage_entry_id; +DROP INDEX IF EXISTS idx_usage_provider; + +COMMIT; diff --git a/memory-store/migrations/000042_usage_reference_fields.up.sql b/memory-store/migrations/000042_usage_reference_fields.up.sql new file mode 100644 index 000000000..58bd87d52 --- /dev/null +++ b/memory-store/migrations/000042_usage_reference_fields.up.sql @@ -0,0 +1,20 @@ +BEGIN; + +-- Add reference columns to usage table +ALTER TABLE usage + ADD COLUMN execution_id UUID, + ADD COLUMN transition_id UUID, + ADD COLUMN session_id UUID, + ADD COLUMN entry_id UUID, + ADD COLUMN provider TEXT; + +-- Index new columns for query efficiency +CREATE INDEX IF NOT EXISTS idx_usage_execution_id ON usage (execution_id); +CREATE INDEX IF NOT EXISTS idx_usage_transition_id ON usage (transition_id); +CREATE INDEX IF NOT EXISTS idx_usage_session_id ON usage (session_id); +CREATE INDEX IF NOT EXISTS idx_usage_entry_id ON usage (entry_id); +CREATE INDEX IF NOT EXISTS idx_usage_provider ON usage (provider); + +COMMENT ON COLUMN usage.provider IS 'Model provider for monitoring'; + +COMMIT;