local-deep-research 0.4.4__py3-none-any.whl → 0.5.2__py3-none-any.whl

local_deep_research/metrics/pricing/cost_calculator.py

@@ -0,0 +1,237 @@
+"""
+Cost Calculator
+
+Calculates LLM usage costs based on token usage and pricing data.
+Integrates with pricing fetcher and cache systems.
+"""
+
+import logging
+from typing import Any, Dict, List, Optional
+
+from .pricing_cache import PricingCache
+from .pricing_fetcher import PricingFetcher
+
+logger = logging.getLogger(__name__)
+
+
+class CostCalculator:
+    """Calculates LLM usage costs."""
+
+    def __init__(self, cache_dir: Optional[str] = None):
+        self.cache = PricingCache(cache_dir)
+        self.pricing_fetcher = None
+
+    async def __aenter__(self):
+        self.pricing_fetcher = PricingFetcher()
+        await self.pricing_fetcher.__aenter__()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        if self.pricing_fetcher:
+            await self.pricing_fetcher.__aexit__(exc_type, exc_val, exc_tb)
+
+    async def get_model_pricing(
+        self, model_name: str, provider: str = None
+    ) -> Dict[str, float]:
+        """Get pricing for a model and provider (cached or fetched)."""
+        # Create cache key that includes provider
+        cache_key = f"{provider}:{model_name}" if provider else model_name
+
+        # Try cache first
+        cached_pricing = self.cache.get(f"model:{cache_key}")
+        if cached_pricing:
+            return cached_pricing
+
+        # Fetch from API
+        if self.pricing_fetcher:
+            pricing = await self.pricing_fetcher.get_model_pricing(
+                model_name, provider
+            )
+            if pricing:
+                self.cache.set(f"model:{cache_key}", pricing)
+                return pricing
+
+        # No pricing found
+        logger.warning(
+            f"No pricing found for {model_name} (provider: {provider})"
+        )
+        return None
+
+    async def calculate_cost(
+        self,
+        model_name: str,
+        prompt_tokens: int,
+        completion_tokens: int,
+        provider: str = None,
+    ) -> Dict[str, float]:
+        """
+        Calculate cost for a single LLM call.
+
+        Returns:
+            Dict with prompt_cost, completion_cost, total_cost
+        """
+        pricing = await self.get_model_pricing(model_name, provider)
+
+        # If no pricing found, return zero cost
+        if pricing is None:
+            return {
+                "prompt_cost": 0.0,
+                "completion_cost": 0.0,
+                "total_cost": 0.0,
+                "pricing_used": None,
+                "error": "No pricing data available for this model",
+            }
+
+        # Convert tokens to thousands for pricing calculation
+        prompt_cost = (prompt_tokens / 1000) * pricing["prompt"]
+        completion_cost = (completion_tokens / 1000) * pricing["completion"]
+        total_cost = prompt_cost + completion_cost
+
+        return {
+            "prompt_cost": round(prompt_cost, 6),
+            "completion_cost": round(completion_cost, 6),
+            "total_cost": round(total_cost, 6),
+            "pricing_used": pricing,
+        }
+
+    async def calculate_batch_costs(
+        self, usage_records: List[Dict[str, Any]]
+    ) -> List[Dict[str, Any]]:
+        """
+        Calculate costs for multiple usage records.
+
+        Expected record format:
+        {
+            "model_name": str,
+            "provider": str (optional),
+            "prompt_tokens": int,
+            "completion_tokens": int,
+            "research_id": int (optional),
+            "timestamp": datetime (optional)
+        }
+        """
+        results = []
+
+        for record in usage_records:
+            try:
+                cost_data = await self.calculate_cost(
+                    record["model_name"],
+                    record["prompt_tokens"],
+                    record["completion_tokens"],
+                    record.get("provider"),
+                )
+
+                result = {**record, **cost_data}
+                results.append(result)
+
+            except Exception as e:
+                logger.error(
+                    f"Failed to calculate cost for record {record}: {e}"
+                )
+                # Add record with zero cost on error
+                results.append(
+                    {
+                        **record,
+                        "prompt_cost": 0.0,
+                        "completion_cost": 0.0,
+                        "total_cost": 0.0,
+                        "error": str(e),
+                    }
+                )
+
+        return results
+
+    def calculate_cost_sync(
+        self, model_name: str, prompt_tokens: int, completion_tokens: int
+    ) -> Dict[str, float]:
+        """
+        Synchronous cost calculation using cached pricing only.
+        Fallback for when async is not available.
+        """
+        # Use cached pricing only
+        pricing = self.cache.get_model_pricing(model_name)
+        if not pricing:
+            # Use static fallback with exact matching only
+            fetcher = PricingFetcher()
+            # Try exact match
+            pricing = fetcher.static_pricing.get(model_name)
+            if not pricing:
+                # Try exact match without provider prefix
+                if "/" in model_name:
+                    model_only = model_name.split("/")[-1]
+                    pricing = fetcher.static_pricing.get(model_only)
+
+        # If no pricing found, return zero cost
+        if not pricing:
+            return {
+                "prompt_cost": 0.0,
+                "completion_cost": 0.0,
+                "total_cost": 0.0,
+                "pricing_used": None,
+                "error": "No pricing data available for this model",
+            }
+
+        prompt_cost = (prompt_tokens / 1000) * pricing["prompt"]
+        completion_cost = (completion_tokens / 1000) * pricing["completion"]
+        total_cost = prompt_cost + completion_cost
+
+        return {
+            "prompt_cost": round(prompt_cost, 6),
+            "completion_cost": round(completion_cost, 6),
+            "total_cost": round(total_cost, 6),
+            "pricing_used": pricing,
+        }
+
+    async def get_research_cost_summary(
+        self, usage_records: List[Dict[str, Any]]
+    ) -> Dict[str, Any]:
+        """
+        Get cost summary for research session(s).
+        """
+        costs = await self.calculate_batch_costs(usage_records)
+
+        total_cost = sum(c["total_cost"] for c in costs)
+        total_prompt_cost = sum(c["prompt_cost"] for c in costs)
+        total_completion_cost = sum(c["completion_cost"] for c in costs)
+
+        total_prompt_tokens = sum(r["prompt_tokens"] for r in usage_records)
+        total_completion_tokens = sum(
+            r["completion_tokens"] for r in usage_records
+        )
+        total_tokens = total_prompt_tokens + total_completion_tokens
+
+        # Model breakdown
+        model_costs = {}
+        for cost in costs:
+            model = cost["model_name"]
+            if model not in model_costs:
+                model_costs[model] = {
+                    "total_cost": 0.0,
+                    "prompt_tokens": 0,
+                    "completion_tokens": 0,
+                    "calls": 0,
+                }
+
+            model_costs[model]["total_cost"] += cost["total_cost"]
+            model_costs[model]["prompt_tokens"] += cost["prompt_tokens"]
+            model_costs[model]["completion_tokens"] += cost["completion_tokens"]
+            model_costs[model]["calls"] += 1
+
+        return {
+            "total_cost": round(total_cost, 6),
+            "prompt_cost": round(total_prompt_cost, 6),
+            "completion_cost": round(total_completion_cost, 6),
+            "total_tokens": total_tokens,
+            "prompt_tokens": total_prompt_tokens,
+            "completion_tokens": total_completion_tokens,
+            "total_calls": len(usage_records),
+            "model_breakdown": model_costs,
+            "avg_cost_per_call": (
+                round(total_cost / len(usage_records), 6)
+                if usage_records
+                else 0.0
+            ),
+            "cost_per_token": (
+                round(total_cost / total_tokens, 8) if total_tokens > 0 else 0.0
+            ),
+        }

local_deep_research/metrics/pricing/pricing_cache.py

@@ -0,0 +1,143 @@
+"""
+Pricing Cache System
+
+Caches pricing data to avoid repeated API calls and improve performance.
+Includes cache expiration and refresh mechanisms.
+"""
+
+import json
+import logging
+import time
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class PricingCache:
+    """Cache for LLM pricing data."""
+
+    def __init__(self, cache_dir: Optional[str] = None, cache_ttl: int = 3600):
+        """
+        Initialize pricing cache.
+
+        Args:
+            cache_dir: Directory to store cache files
+            cache_ttl: Cache time-to-live in seconds (default: 1 hour)
+        """
+        self.cache_ttl = cache_ttl
+
+        if cache_dir:
+            self.cache_dir = Path(cache_dir)
+        else:
+            # Default to data directory
+            self.cache_dir = Path.cwd() / "data" / "cache" / "pricing"
+
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+        self.cache_file = self.cache_dir / "pricing_cache.json"
+
+        self._cache = {}
+        self._load_cache()
+
+    def _load_cache(self):
+        """Load cache from disk."""
+        try:
+            if self.cache_file.exists():
+                with open(self.cache_file, "r") as f:
+                    data = json.load(f)
+                    self._cache = data.get("cache", {})
+                    logger.info(
+                        f"Loaded pricing cache with {len(self._cache)} entries"
+                    )
+        except Exception as e:
+            logger.warning(f"Failed to load pricing cache: {e}")
+            self._cache = {}
+
+    def _save_cache(self):
+        """Save cache to disk."""
+        try:
+            cache_data = {
+                "cache": self._cache,
+                "last_updated": datetime.now().isoformat(),
+            }
+            with open(self.cache_file, "w") as f:
+                json.dump(cache_data, f, indent=2)
+        except Exception as e:
+            logger.warning(f"Failed to save pricing cache: {e}")
+
+    def _is_expired(self, timestamp: float) -> bool:
+        """Check if cache entry is expired."""
+        return (time.time() - timestamp) > self.cache_ttl
+
+    def get(self, key: str) -> Optional[Any]:
+        """Get cached pricing data."""
+        if key not in self._cache:
+            return None
+
+        entry = self._cache[key]
+        if self._is_expired(entry["timestamp"]):
+            # Remove expired entry
+            del self._cache[key]
+            self._save_cache()
+            return None
+
+        return entry["data"]
+
+    def set(self, key: str, data: Any):
+        """Set cached pricing data."""
+        self._cache[key] = {"data": data, "timestamp": time.time()}
+        self._save_cache()
+
+    def get_model_pricing(self, model_name: str) -> Optional[Dict[str, float]]:
+        """Get cached pricing for a specific model."""
+        return self.get(f"model:{model_name}")
+
+    def set_model_pricing(self, model_name: str, pricing: Dict[str, float]):
+        """Cache pricing for a specific model."""
+        self.set(f"model:{model_name}", pricing)
+
+    def get_all_pricing(self) -> Optional[Dict[str, Dict[str, float]]]:
+        """Get cached pricing for all models."""
+        return self.get("all_models")
+
+    def set_all_pricing(self, pricing: Dict[str, Dict[str, float]]):
+        """Cache pricing for all models."""
+        self.set("all_models", pricing)
+
+    def clear(self):
+        """Clear all cached data."""
+        self._cache = {}
+        self._save_cache()
+        logger.info("Pricing cache cleared")
+
+    def clear_expired(self):
+        """Remove expired cache entries."""
+        expired_keys = []
+        for key, entry in self._cache.items():
+            if self._is_expired(entry["timestamp"]):
+                expired_keys.append(key)
+
+        for key in expired_keys:
+            del self._cache[key]
+
+        if expired_keys:
+            self._save_cache()
+            logger.info(f"Removed {len(expired_keys)} expired cache entries")
+
+    def get_cache_stats(self) -> Dict[str, Any]:
+        """Get cache statistics."""
+        total_entries = len(self._cache)
+        expired_count = 0
+
+        for entry in self._cache.values():
+            if self._is_expired(entry["timestamp"]):
+                expired_count += 1
+
+        return {
+            "total_entries": total_entries,
+            "expired_entries": expired_count,
+            "valid_entries": total_entries - expired_count,
+            "cache_file": str(self.cache_file),
+            "cache_ttl": self.cache_ttl,
+        }

local_deep_research/metrics/pricing/pricing_fetcher.py

@@ -0,0 +1,240 @@
+"""
+LLM Pricing Data Fetcher
+
+Fetches real-time pricing data from various LLM providers.
+Supports multiple providers and fallback to static pricing.
+"""
+
+from typing import Any, Dict, Optional
+
+import aiohttp
+from loguru import logger
+
+
+class PricingFetcher:
+    """Fetches LLM pricing data from various sources."""
+
+    def __init__(self):
+        self.session = None
+        self.static_pricing = self._load_static_pricing()
+
+    async def __aenter__(self):
+        self.session = aiohttp.ClientSession()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        if self.session:
+            await self.session.close()
+
+    def _load_static_pricing(self) -> Dict[str, Dict[str, float]]:
+        """Load static pricing as fallback (per 1K tokens in USD)."""
+        return {
+            # OpenAI Models
+            "gpt-4": {"prompt": 0.03, "completion": 0.06},
+            "gpt-4-turbo": {"prompt": 0.01, "completion": 0.03},
+            "gpt-4o": {"prompt": 0.005, "completion": 0.015},
+            "gpt-4o-mini": {"prompt": 0.00015, "completion": 0.0006},
+            "gpt-3.5-turbo": {"prompt": 0.001, "completion": 0.002},
+            # Anthropic Models
+            "claude-3-opus": {"prompt": 0.015, "completion": 0.075},
+            "claude-3-sonnet": {"prompt": 0.003, "completion": 0.015},
+            "claude-3-haiku": {"prompt": 0.00025, "completion": 0.00125},
+            "claude-3-5-sonnet": {"prompt": 0.003, "completion": 0.015},
+            # Google Models
+            "gemini-pro": {"prompt": 0.0005, "completion": 0.0015},
+            "gemini-pro-vision": {"prompt": 0.0005, "completion": 0.0015},
+            "gemini-1.5-pro": {"prompt": 0.0035, "completion": 0.0105},
+            "gemini-1.5-flash": {"prompt": 0.00035, "completion": 0.00105},
+            # Local/Open Source (free)
+            "ollama": {"prompt": 0.0, "completion": 0.0},
+            "llama": {"prompt": 0.0, "completion": 0.0},
+            "mistral": {"prompt": 0.0, "completion": 0.0},
+            "gemma": {"prompt": 0.0, "completion": 0.0},
+            "qwen": {"prompt": 0.0, "completion": 0.0},
+            "codellama": {"prompt": 0.0, "completion": 0.0},
+            "vicuna": {"prompt": 0.0, "completion": 0.0},
+            "alpaca": {"prompt": 0.0, "completion": 0.0},
+            "vllm": {"prompt": 0.0, "completion": 0.0},
+            "lmstudio": {"prompt": 0.0, "completion": 0.0},
+            "llamacpp": {"prompt": 0.0, "completion": 0.0},
+        }
+
+    async def fetch_openai_pricing(self) -> Optional[Dict[str, Any]]:
+        """Fetch OpenAI pricing from their API (if available)."""
+        try:
+            # Note: OpenAI doesn't have a public pricing API
+            # This would need to be web scraping or manual updates
+            logger.info("Using static OpenAI pricing (no public API available)")
+            return None
+        except Exception as e:
+            logger.warning(f"Failed to fetch OpenAI pricing: {e}")
+            return None
+
+    async def fetch_anthropic_pricing(self) -> Optional[Dict[str, Any]]:
+        """Fetch Anthropic pricing."""
+        try:
+            # Note: Anthropic doesn't have a public pricing API
+            # This would need to be web scraping or manual updates
+            logger.info(
+                "Using static Anthropic pricing (no public API available)"
+            )
+            return None
+        except Exception as e:
+            logger.warning(f"Failed to fetch Anthropic pricing: {e}")
+            return None
+
+    async def fetch_google_pricing(self) -> Optional[Dict[str, Any]]:
+        """Fetch Google/Gemini pricing."""
+        try:
+            # Note: Google doesn't have a dedicated pricing API for individual models
+            # This would need to be web scraping or manual updates
+            logger.info("Using static Google pricing (no public API available)")
+            return None
+        except Exception as e:
+            logger.warning(f"Failed to fetch Google pricing: {e}")
+            return None
+
+    async def fetch_huggingface_pricing(self) -> Optional[Dict[str, Any]]:
+        """Fetch HuggingFace Inference API pricing."""
+        try:
+            if not self.session:
+                return None
+
+            # HuggingFace has some pricing info but not a structured API
+            # This is more for hosted inference endpoints
+            url = "https://huggingface.co/pricing"
+            async with self.session.get(url) as response:
+                if response.status == 200:
+                    # Would need to parse HTML for pricing info
+                    logger.info(
+                        "HuggingFace pricing would require HTML parsing"
+                    )
+                    return None
+        except Exception as e:
+            logger.warning(f"Failed to fetch HuggingFace pricing: {e}")
+            return None
+
+    async def get_model_pricing(
+        self, model_name: str, provider: str = None
+    ) -> Optional[Dict[str, float]]:
+        """Get pricing for a specific model and provider."""
+        # Normalize inputs
+        model_name = model_name.lower() if model_name else ""
+        provider = provider.lower() if provider else ""
+
+        # Provider-first approach: Check if provider indicates local/free models
+        local_providers = ["ollama", "vllm", "lmstudio", "llamacpp"]
+        if provider in local_providers:
+            logger.debug(
+                f"Local provider '{provider}' detected - returning zero cost"
+            )
+            return {"prompt": 0.0, "completion": 0.0}
+
+        # Try to fetch live pricing first (most providers don't have APIs)
+        if (
+            provider == "openai"
+            or "gpt" in model_name
+            or "openai" in model_name
+        ):
+            await self.fetch_openai_pricing()
+        elif (
+            provider == "anthropic"
+            or "claude" in model_name
+            or "anthropic" in model_name
+        ):
+            await self.fetch_anthropic_pricing()
+        elif (
+            provider == "google"
+            or "gemini" in model_name
+            or "google" in model_name
+        ):
+            await self.fetch_google_pricing()
+
+        # Fallback to static pricing with provider priority
+        if provider:
+            # First try provider-specific lookup with exact matching
+            provider_models = self._get_models_by_provider(provider)
+            # Try exact match
+            if model_name in provider_models:
+                return provider_models[model_name]
+            # Try exact match without provider prefix
+            if "/" in model_name:
+                model_only = model_name.split("/")[-1]
+                if model_only in provider_models:
+                    return provider_models[model_only]
+
+        # Exact model name matching only
+        # First try exact match
+        if model_name in self.static_pricing:
+            return self.static_pricing[model_name]
+
+        # Try exact match without provider prefix (e.g., "openai/gpt-4o-mini" -> "gpt-4o-mini")
+        if "/" in model_name:
+            model_only = model_name.split("/")[-1]
+            if model_only in self.static_pricing:
+                return self.static_pricing[model_only]
+
+        # No pricing found - return None instead of default pricing
+        logger.warning(
+            f"No pricing found for model: {model_name}, provider: {provider}"
+        )
+        return None
+
+    def _get_models_by_provider(
+        self, provider: str
+    ) -> Dict[str, Dict[str, float]]:
+        """Get models for a specific provider."""
+        provider = provider.lower()
+        provider_models = {}
+
+        if provider == "openai":
+            provider_models = {
+                k: v
+                for k, v in self.static_pricing.items()
+                if k.startswith("gpt")
+            }
+        elif provider == "anthropic":
+            provider_models = {
+                k: v
+                for k, v in self.static_pricing.items()
+                if k.startswith("claude")
+            }
+        elif provider == "google":
+            provider_models = {
+                k: v
+                for k, v in self.static_pricing.items()
+                if k.startswith("gemini")
+            }
+        elif provider in ["ollama", "vllm", "lmstudio", "llamacpp"]:
+            # All local models are free
+            provider_models = {
+                k: v
+                for k, v in self.static_pricing.items()
+                if v["prompt"] == 0.0 and v["completion"] == 0.0
+            }
+
+        return provider_models
+
+    async def get_all_pricing(self) -> Dict[str, Dict[str, float]]:
+        """Get pricing for all known models."""
+        # In the future, this could aggregate from multiple live sources
+        return self.static_pricing.copy()
+
+    def get_provider_from_model(self, model_name: str) -> str:
+        """Determine the provider from model name."""
+        model_name = model_name.lower()
+
+        if "gpt" in model_name or "openai" in model_name:
+            return "openai"
+        elif "claude" in model_name or "anthropic" in model_name:
+            return "anthropic"
+        elif "gemini" in model_name or "google" in model_name:
+            return "google"
+        elif "llama" in model_name or "meta" in model_name:
+            return "meta"
+        elif "mistral" in model_name:
+            return "mistral"
+        elif "ollama" in model_name:
+            return "ollama"
+        else:
+            return "unknown"

local_deep_research/metrics/query_utils.py

@@ -0,0 +1,51 @@
+"""Common query utilities for metrics module."""
+
+from datetime import datetime, timedelta
+from typing import Any
+
+from sqlalchemy import Column
+
+
+def get_time_filter_condition(period: str, timestamp_column: Column) -> Any:
+    """Get SQLAlchemy condition for time filtering.
+
+    Args:
+        period: Time period ('7d', '30d', '3m', '1y', 'all')
+        timestamp_column: SQLAlchemy timestamp column to filter on
+
+    Returns:
+        SQLAlchemy condition object or None for 'all'
+    """
+    if period == "all":
+        return None
+    elif period == "7d":
+        cutoff = datetime.now() - timedelta(days=7)
+    elif period == "30d":
+        cutoff = datetime.now() - timedelta(days=30)
+    elif period == "3m":
+        cutoff = datetime.now() - timedelta(days=90)
+    elif period == "1y":
+        cutoff = datetime.now() - timedelta(days=365)
+    else:
+        # Default to 30 days for unknown periods
+        cutoff = datetime.now() - timedelta(days=30)
+
+    return timestamp_column >= cutoff
+
+
+def get_research_mode_condition(research_mode: str, mode_column: Column) -> Any:
+    """Get SQLAlchemy condition for research mode filtering.
+
+    Args:
+        research_mode: Research mode ('quick', 'detailed', 'all')
+        mode_column: SQLAlchemy column to filter on
+
+    Returns:
+        SQLAlchemy condition object or None for 'all'
+    """
+    if research_mode == "all":
+        return None
+    elif research_mode in ["quick", "detailed"]:
+        return mode_column == research_mode
+    else:
+        return None