How to Track Agent API Spending with Budgets

import os, requests, json, time
from datetime import datetime
from collections import defaultdict

API_KEY = os.environ['SCAVIO_API_KEY']
SH = {'x-api-key': API_KEY, 'Content-Type': 'application/json'}

class BudgetTracker:
    def __init__(self, cost_per_call=0.005):
        self.cost_per_call = cost_per_call
        self.agents = {}  # agent_id -> {budget, spent, calls}
        self.log = []

    def register_agent(self, agent_id, daily_budget):
        self.agents[agent_id] = {
            'budget': daily_budget,
            'spent': 0.0,
            'calls': 0,
            'created': datetime.now().strftime('%Y-%m-%d')
        }
        print(f'Registered: {agent_id} (budget: ${daily_budget}/day)')

    def can_spend(self, agent_id):
        agent = self.agents.get(agent_id)
        if not agent:
            return False
        remaining = agent['budget'] - agent['spent']
        return remaining >= self.cost_per_call

    def search(self, agent_id, query, **kwargs):
        if not self.can_spend(agent_id):
            remaining = self.agents[agent_id]['budget'] - self.agents[agent_id]['spent']
            print(f'BUDGET EXCEEDED: {agent_id} (${remaining:.3f} remaining)')
            return {'error': 'budget_exceeded', 'remaining': remaining}
        data = requests.post('https://api.scavio.dev/api/v1/search',
            headers=SH, json={'query': query, 'country_code': 'us', **kwargs}).json()
        self.agents[agent_id]['spent'] += self.cost_per_call
        self.agents[agent_id]['calls'] += 1
        self.log.append({'agent': agent_id, 'query': query, 'cost': self.cost_per_call,
                         'time': datetime.now().isoformat()})
        return data

tracker = BudgetTracker()
tracker.register_agent('research-agent', daily_budget=0.50)
tracker.register_agent('seo-agent', daily_budget=1.00)
tracker.register_agent('lead-agent', daily_budget=0.25)

# Simulate agent usage
queries = {
    'research-agent': ['best python framework 2026', 'fastapi vs django', 'python testing tools'],
    'seo-agent': ['serp api comparison', 'search api pricing', 'mcp search tool', 'ai agent search'],
    'lead-agent': ['saas companies hiring 2026', 'startup funding round'],
}

for agent_id, agent_queries in queries.items():
    print(f'\n--- {agent_id} ---')
    for q in agent_queries:
        result = tracker.search(agent_id, q)
        if 'error' in result:
            print(f'  BLOCKED: {q}')
        else:
            count = len(result.get('organic_results', []))
            print(f'  "{q}" -> {count} results')

# Check remaining budgets
print(f'\n=== Budget Status ===')
for agent_id, data in tracker.agents.items():
    remaining = data['budget'] - data['spent']
    pct = data['spent'] / data['budget'] * 100
    bar = '#' * int(pct / 5)
    print(f'  {agent_id:20} | ${data["spent"]:.3f}/${data["budget"]:.2f} ({pct:.0f}%) | {bar}')

def spending_report(tracker):
    print(f'\n=== Agent Spending Report - {datetime.now().strftime("%Y-%m-%d")} ===')
    total_spent = sum(a['spent'] for a in tracker.agents.values())
    total_budget = sum(a['budget'] for a in tracker.agents.values())
    total_calls = sum(a['calls'] for a in tracker.agents.values())
    print(f'  Total spent: ${total_spent:.3f} / ${total_budget:.2f} budget')
    print(f'  Total calls: {total_calls}')
    print(f'  Avg cost/call: ${total_spent/total_calls:.4f}' if total_calls else '')
    print(f'\n  Per-Agent Breakdown:')
    for agent_id, data in sorted(tracker.agents.items(), key=lambda x: x[1]['spent'], reverse=True):
        efficiency = data['spent'] / data['calls'] if data['calls'] else 0
        print(f'    {agent_id:20} | {data["calls"]:3} calls | ${data["spent"]:.3f} spent | ${efficiency:.4f}/call')
    # Monthly projection
    monthly = total_spent * 30
    print(f'\n  Monthly projection: ${monthly:.2f}')
    print(f'  Annual projection: ${monthly * 12:.2f}')
    # Optimization suggestions
    print(f'\n  Optimization Suggestions:')
    for agent_id, data in tracker.agents.items():
        usage = data['spent'] / data['budget'] * 100
        if usage < 30:
            print(f'    {agent_id}: Budget underused ({usage:.0f}%). Consider reducing to ${data["spent"] * 2:.2f}/day.')
        elif usage > 90:
            print(f'    {agent_id}: Near limit ({usage:.0f}%). Consider increasing budget.')

spending_report(tracker)