api_client.py

import os
import json
import requests
from datetime import datetime, timedelta
from typing import Dict, List, Optional
import time

class AnthropicAPIClient:
    def __init__(self, api_key: Optional[str] = None):
        self.api_key = api_key or os.getenv('ANTHROPIC_API_KEY')
        if not self.api_key:
            raise ValueError("Anthropic API key not found. Set ANTHROPIC_API_KEY environment variable.")

        self.base_url = "https://api.anthropic.com"
        self.headers = {
            "x-api-key": self.api_key,
            "anthropic-version": "2023-06-01",
            "content-type": "application/json"
        }

        self.model_pricing = {
            "claude-3-opus-20240229": {"input": 0.015, "output": 0.075},
            "claude-3-sonnet-20240229": {"input": 0.003, "output": 0.015},
            "claude-3-haiku-20240307": {"input": 0.00025, "output": 0.00125},
            "claude-2.1": {"input": 0.008, "output": 0.024},
            "claude-2.0": {"input": 0.008, "output": 0.024},
            "claude-instant-1.2": {"input": 0.0008, "output": 0.0024}
        }

    def calculate_cost(self, model: str, input_tokens: int, output_tokens: int) -> float:
        """Calculate the cost based on model and token usage"""
        if model not in self.model_pricing:
            model = "claude-3-sonnet-20240229"

        pricing = self.model_pricing[model]
        input_cost = (input_tokens / 1000) * pricing["input"]
        output_cost = (output_tokens / 1000) * pricing["output"]

        return round(input_cost + output_cost, 6)

    def simulate_usage_data(self, days_back: int = 30) -> List[Dict]:
        """Generate simulated usage data for testing"""
        import random

        usage_data = []
        models = list(self.model_pricing.keys())

        for i in range(days_back):
            date = datetime.now() - timedelta(days=i)
            num_conversations = random.randint(5, 20)

            for j in range(num_conversations):
                model = random.choice(models)
                input_tokens = random.randint(100, 5000)
                output_tokens = random.randint(200, 8000)

                usage_data.append({
                    "timestamp": date.isoformat(),
                    "model": model,
                    "input_tokens": input_tokens,
                    "output_tokens": output_tokens,
                    "cost": self.calculate_cost(model, input_tokens, output_tokens),
                    "conversation_id": f"conv_{date.strftime('%Y%m%d')}_{j}",
                    "purpose": f"Test conversation {j}"
                })

        return usage_data

    def test_connection(self) -> bool:
        """Test the API connection"""
        try:
            url = f"{self.base_url}/v1/messages"
            data = {
                "model": "claude-3-haiku-20240307",
                "max_tokens": 10,
                "messages": [{"role": "user", "content": "Hi"}]
            }

            response = requests.post(url, headers=self.headers, json=data, timeout=10)
            return response.status_code == 200
        except Exception as e:
            print(f"Connection test failed: {e}")
            return False

    def fetch_usage_data(self, start_date: Optional[datetime] = None,
                        end_date: Optional[datetime] = None) -> List[Dict]:
        """
        Fetch usage data from Anthropic API
        Note: This is a placeholder as Anthropic doesn't currently provide usage endpoints
        In production, you would track usage as you make API calls
        """
        if not start_date:
            start_date = datetime.now() - timedelta(days=30)
        if not end_date:
            end_date = datetime.now()

        print("Note: Using simulated data. Real-time usage tracking requires monitoring actual API calls.")
        return self.simulate_usage_data(30)

    def track_api_call(self, model: str, input_tokens: int, output_tokens: int,
                      conversation_id: Optional[str] = None, purpose: Optional[str] = None) -> Dict:
        """Track an API call for usage monitoring"""
        cost = self.calculate_cost(model, input_tokens, output_tokens)

        return {
            "timestamp": datetime.utcnow().isoformat(),
            "model": model,
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "cost": cost,
            "conversation_id": conversation_id or f"conv_{int(time.time())}",
            "purpose": purpose or "API call"
        }

    def get_usage_summary(self, usage_data: List[Dict]) -> Dict:
        """Get a summary of usage data"""
        if not usage_data:
            return {
                "total_cost": 0,
                "total_input_tokens": 0,
                "total_output_tokens": 0,
                "conversation_count": 0,
                "models_used": {}
            }

        total_cost = sum(item.get("cost", 0) for item in usage_data)
        total_input = sum(item.get("input_tokens", 0) for item in usage_data)
        total_output = sum(item.get("output_tokens", 0) for item in usage_data)

        models_used = {}
        for item in usage_data:
            model = item.get("model", "unknown")
            if model not in models_used:
                models_used[model] = {"count": 0, "cost": 0, "tokens": 0}
            models_used[model]["count"] += 1
            models_used[model]["cost"] += item.get("cost", 0)
            models_used[model]["tokens"] += item.get("input_tokens", 0) + item.get("output_tokens", 0)

        return {
            "total_cost": round(total_cost, 2),
            "total_input_tokens": total_input,
            "total_output_tokens": total_output,
            "total_tokens": total_input + total_output,
            "conversation_count": len(usage_data),
            "models_used": models_used
        }

if __name__ == "__main__":
    client = AnthropicAPIClient()

    if client.test_connection():
        print("API connection successful!")

    usage_data = client.fetch_usage_data()
    summary = client.get_usage_summary(usage_data)

    print(f"\nUsage Summary (Last 30 days):")
    print(f"Total Cost: ${summary['total_cost']}")
    print(f"Total Tokens: {summary['total_tokens']:,}")
    print(f"Conversations: {summary['conversation_count']}")
    print(f"Models Used: {list(summary['models_used'].keys())}")