local-deep-research 0.4.3__py3-none-any.whl → 0.5.0__py3-none-any.whl

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

local_deep_research/metrics/search_tracker.py

@@ -0,0 +1,380 @@
+"""
+Search call tracking system for metrics collection.
+Similar to token_counter.py but tracks search engine usage.
+"""
+
+from typing import Any, Dict, List, Optional
+
+from loguru import logger
+from sqlalchemy import case, func
+
+from .database import MetricsDatabase
+from .db_models import SearchCall
+from .query_utils import get_research_mode_condition, get_time_filter_condition
+
+
+class SearchTracker:
+    """Track search engine calls and performance metrics."""
+
+    def __init__(self, db: Optional[MetricsDatabase] = None):
+        """Initialize the search tracker."""
+        self.db = db or MetricsDatabase()
+        self.research_context = {}
+
+    def set_research_context(self, context: Dict[str, Any]) -> None:
+        """Set the current research context for search tracking."""
+        self.research_context = context or {}
+        logger.debug(f"Search tracker context updated: {self.research_context}")
+
+    def record_search(
+        self,
+        engine_name: str,
+        query: str,
+        results_count: int = 0,
+        response_time_ms: int = 0,
+        success: bool = True,
+        error_message: Optional[str] = None,
+    ) -> None:
+        """Record a completed search operation directly to database."""
+
+        # Extract research context
+        research_id = self.research_context.get("research_id")
+        research_query = self.research_context.get("research_query")
+        research_mode = self.research_context.get("research_mode", "unknown")
+        research_phase = self.research_context.get("research_phase", "search")
+        search_iteration = self.research_context.get("search_iteration", 0)
+
+        # Determine success status
+        success_status = "success" if success else "error"
+        error_type = None
+        if error_message:
+            error_type = (
+                type(error_message).__name__
+                if isinstance(error_message, Exception)
+                else "unknown_error"
+            )
+
+        # Record search call in database
+        try:
+            with self.db.get_session() as session:
+                # Create search call record
+                search_call = SearchCall(
+                    research_id=research_id,
+                    research_query=research_query,
+                    research_mode=research_mode,
+                    research_phase=research_phase,
+                    search_iteration=search_iteration,
+                    search_engine=engine_name,
+                    query=query,
+                    results_count=results_count,
+                    response_time_ms=response_time_ms,
+                    success_status=success_status,
+                    error_type=error_type,
+                    error_message=str(error_message) if error_message else None,
+                )
+                session.add(search_call)
+                session.commit()
+
+                logger.debug(
+                    f"Search call recorded: {engine_name} - "
+                    f"{results_count} results in {response_time_ms}ms"
+                )
+
+        except Exception as e:
+            logger.error(f"Failed to record search call: {e}")
+
+    def get_search_metrics(
+        self, period: str = "30d", research_mode: str = "all"
+    ) -> Dict[str, Any]:
+        """Get search engine usage metrics."""
+        with self.db.get_session() as session:
+            try:
+                # Build base query with filters
+                query = session.query(SearchCall).filter(
+                    SearchCall.search_engine.isnot(None)
+                )
+
+                # Apply time filter
+                time_condition = get_time_filter_condition(
+                    period, SearchCall.timestamp
+                )
+                if time_condition is not None:
+                    query = query.filter(time_condition)
+
+                # Apply research mode filter
+                mode_condition = get_research_mode_condition(
+                    research_mode, SearchCall.research_mode
+                )
+                if mode_condition is not None:
+                    query = query.filter(mode_condition)
+
+                # Get search engine statistics using ORM aggregation
+                search_stats = session.query(
+                    SearchCall.search_engine,
+                    func.count().label("call_count"),
+                    func.avg(SearchCall.response_time_ms).label(
+                        "avg_response_time"
+                    ),
+                    func.sum(SearchCall.results_count).label("total_results"),
+                    func.avg(SearchCall.results_count).label(
+                        "avg_results_per_call"
+                    ),
+                    func.sum(
+                        case(
+                            (SearchCall.success_status == "success", 1), else_=0
+                        )
+                    ).label("success_count"),
+                    func.sum(
+                        case((SearchCall.success_status == "error", 1), else_=0)
+                    ).label("error_count"),
+                ).filter(SearchCall.search_engine.isnot(None))
+
+                # Apply same filters to stats query
+                if time_condition is not None:
+                    search_stats = search_stats.filter(time_condition)
+                if mode_condition is not None:
+                    search_stats = search_stats.filter(mode_condition)
+
+                search_stats = (
+                    search_stats.group_by(SearchCall.search_engine)
+                    .order_by(func.count().desc())
+                    .all()
+                )
+
+                # Get recent search calls
+                recent_calls_query = session.query(SearchCall)
+                if time_condition is not None:
+                    recent_calls_query = recent_calls_query.filter(
+                        time_condition
+                    )
+                if mode_condition is not None:
+                    recent_calls_query = recent_calls_query.filter(
+                        mode_condition
+                    )
+
+                recent_calls = (
+                    recent_calls_query.order_by(SearchCall.timestamp.desc())
+                    .limit(20)
+                    .all()
+                )
+
+                return {
+                    "search_engine_stats": [
+                        {
+                            "engine": stat.search_engine,
+                            "call_count": stat.call_count,
+                            "avg_response_time": stat.avg_response_time or 0,
+                            "total_results": stat.total_results or 0,
+                            "avg_results_per_call": stat.avg_results_per_call
+                            or 0,
+                            "success_rate": (
+                                (stat.success_count / stat.call_count * 100)
+                                if stat.call_count > 0
+                                else 0
+                            ),
+                            "error_count": stat.error_count or 0,
+                        }
+                        for stat in search_stats
+                    ],
+                    "recent_calls": [
+                        {
+                            "engine": call.search_engine,
+                            "query": (
+                                call.query[:100] + "..."
+                                if len(call.query or "") > 100
+                                else call.query
+                            ),
+                            "results_count": call.results_count,
+                            "response_time_ms": call.response_time_ms,
+                            "success_status": call.success_status,
+                            "timestamp": str(call.timestamp),
+                        }
+                        for call in recent_calls
+                    ],
+                }
+
+            except Exception as e:
+                logger.exception(f"Error getting search metrics: {e}")
+                return {"search_engine_stats": [], "recent_calls": []}
+
+    def get_research_search_metrics(self, research_id: int) -> Dict[str, Any]:
+        """Get search metrics for a specific research session."""
+        with self.db.get_session() as session:
+            try:
+                # Get all search calls for this research
+                search_calls = (
+                    session.query(SearchCall)
+                    .filter(SearchCall.research_id == research_id)
+                    .order_by(SearchCall.timestamp.asc())
+                    .all()
+                )
+
+                # Get search engine stats for this research
+                engine_stats = (
+                    session.query(
+                        SearchCall.search_engine,
+                        func.count().label("call_count"),
+                        func.avg(SearchCall.response_time_ms).label(
+                            "avg_response_time"
+                        ),
+                        func.sum(SearchCall.results_count).label(
+                            "total_results"
+                        ),
+                        func.sum(
+                            case(
+                                (SearchCall.success_status == "success", 1),
+                                else_=0,
+                            )
+                        ).label("success_count"),
+                    )
+                    .filter(SearchCall.research_id == research_id)
+                    .group_by(SearchCall.search_engine)
+                    .order_by(func.count().desc())
+                    .all()
+                )
+
+                # Calculate totals
+                total_searches = len(search_calls)
+                total_results = sum(
+                    call.results_count or 0 for call in search_calls
+                )
+                avg_response_time = (
+                    sum(call.response_time_ms or 0 for call in search_calls)
+                    / total_searches
+                    if total_searches > 0
+                    else 0
+                )
+                successful_searches = sum(
+                    1
+                    for call in search_calls
+                    if call.success_status == "success"
+                )
+                success_rate = (
+                    (successful_searches / total_searches * 100)
+                    if total_searches > 0
+                    else 0
+                )
+
+                return {
+                    "total_searches": total_searches,
+                    "total_results": total_results,
+                    "avg_response_time": round(avg_response_time),
+                    "success_rate": round(success_rate, 1),
+                    "search_calls": [
+                        {
+                            "engine": call.search_engine,
+                            "query": call.query,
+                            "results_count": call.results_count,
+                            "response_time_ms": call.response_time_ms,
+                            "success_status": call.success_status,
+                            "timestamp": str(call.timestamp),
+                        }
+                        for call in search_calls
+                    ],
+                    "engine_stats": [
+                        {
+                            "engine": stat.search_engine,
+                            "call_count": stat.call_count,
+                            "avg_response_time": stat.avg_response_time or 0,
+                            "total_results": stat.total_results or 0,
+                            "success_rate": (
+                                (stat.success_count / stat.call_count * 100)
+                                if stat.call_count > 0
+                                else 0
+                            ),
+                        }
+                        for stat in engine_stats
+                    ],
+                }
+
+            except Exception as e:
+                logger.exception(f"Error getting research search metrics: {e}")
+                return {
+                    "total_searches": 0,
+                    "total_results": 0,
+                    "avg_response_time": 0,
+                    "success_rate": 0,
+                    "search_calls": [],
+                    "engine_stats": [],
+                }
+
+    def get_search_time_series(
+        self, period: str = "30d", research_mode: str = "all"
+    ) -> List[Dict[str, Any]]:
+        """Get search activity time series data for charting.
+
+        Args:
+            period: Time period to filter by ('7d', '30d', '3m', '1y', 'all')
+            research_mode: Research mode to filter by ('quick', 'detailed', 'all')
+
+        Returns:
+            List of time series data points with search engine activity
+        """
+        with self.db.get_session() as session:
+            try:
+                # Build base query
+                query = session.query(SearchCall).filter(
+                    SearchCall.search_engine.isnot(None),
+                    SearchCall.timestamp.isnot(None),
+                )
+
+                # Apply time filter
+                time_condition = get_time_filter_condition(
+                    period, SearchCall.timestamp
+                )
+                if time_condition is not None:
+                    query = query.filter(time_condition)
+
+                # Apply research mode filter
+                mode_condition = get_research_mode_condition(
+                    research_mode, SearchCall.research_mode
+                )
+                if mode_condition is not None:
+                    query = query.filter(mode_condition)
+
+                # Get all search calls ordered by time
+                search_calls = query.order_by(SearchCall.timestamp.asc()).all()
+
+                # Create time series data
+                time_series = []
+                for call in search_calls:
+                    time_series.append(
+                        {
+                            "timestamp": (
+                                str(call.timestamp) if call.timestamp else None
+                            ),
+                            "search_engine": call.search_engine,
+                            "results_count": call.results_count or 0,
+                            "response_time_ms": call.response_time_ms or 0,
+                            "success_status": call.success_status,
+                            "query": (
+                                call.query[:50] + "..."
+                                if call.query and len(call.query) > 50
+                                else call.query
+                            ),
+                        }
+                    )
+
+                return time_series
+
+            except Exception as e:
+                logger.exception(f"Error getting search time series: {e}")
+                return []
+
+
+# Global search tracker instance
+_search_tracker = None
+
+
+def get_search_tracker() -> SearchTracker:
+    """Get the global search tracker instance."""
+    global _search_tracker
+    if _search_tracker is None:
+        _search_tracker = SearchTracker()
+    return _search_tracker
+
+
+def set_search_context(context: Dict[str, Any]) -> None:
+    """Set search context for the global tracker."""
+    tracker = get_search_tracker()
+    tracker.set_research_context(context)