src.cost_tracking.ai_usage_middleware module

AI Usage Middleware for Token Tracking.

Intercepts all AI provider calls to track token usage per project.

class src.cost_tracking.ai_usage_middleware.AIUsageMiddleware

Bases: object

Middleware that wraps AI provider calls to track token usage.

This intercepts all calls to AI providers (OpenAI, Anthropic, etc.) and tracks token consumption per project.

__init__()

Return type:

None

set_project_context(agent_id, project_id, task_id=None)

Set the current project context for an agent.

This should be called when an agent starts working on a project/task.

Return type:

None

Parameters:

• agent_id (str)

• project_id (str)

• task_id (str | None)

clear_project_context(agent_id)

Clear project context when agent finishes.

Return type:

None

Parameters:

agent_id (str)

get_current_project(agent_id)

Get current project for an agent.

Return type:

Optional[str]

Parameters:

agent_id (str)

track_ai_usage(func)

Decorate AI provider methods to track token usage.

This wraps AI provider methods to capture token usage.

Return type:

Callable[..., Any]

Parameters:

func (Callable[[...], Any])

wrap_ai_provider(provider_instance)

Wrap an AI provider instance to track all its method calls.

This modifies the provider instance to track token usage on all methods that make API calls.

Return type:

Any

Parameters:

provider_instance (Any)

src.cost_tracking.ai_usage_middleware.track_project_tokens(project_id, agent_id=None)

Context manager to track AI tokens for a specific project.

Return type:

Any

Parameters:

• project_id (str)

• agent_id (str | None)

Usage:

with track_project_tokens(“project_123”, “agent_1”):

# All AI calls in this block will be tracked to project_123 await ai_engine.analyze(…)