Add usage reference fields

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,
     )

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,

memory-store/README.md

@@ -8,3 +8,7 @@
 
 For creating a migration:
 `migrate -database "postgres://postgres:[email protected]: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`.

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;

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;