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

local_deep_research/config/llm_config.py

@@ -1,4 +1,5 @@
 import os
+from functools import cache
 
 from langchain_anthropic import ChatAnthropic
 from langchain_community.llms import VLLM
@@ -24,7 +25,163 @@ VALID_PROVIDERS = [
 ]
 
 
-def get_llm(model_name=None, temperature=None, provider=None, openai_endpoint_url=None):
+def is_openai_available():
+    """Check if OpenAI is available"""
+    try:
+        api_key = get_db_setting("llm.openai.api_key")
+        return bool(api_key)
+    except Exception:
+        return False
+
+
+def is_anthropic_available():
+    """Check if Anthropic is available"""
+    try:
+        api_key = get_db_setting("llm.anthropic.api_key")
+        return bool(api_key)
+    except Exception:
+        return False
+
+
+def is_openai_endpoint_available():
+    """Check if OpenAI endpoint is available"""
+    try:
+        api_key = get_db_setting("llm.openai_endpoint.api_key")
+        return bool(api_key)
+    except Exception:
+        return False
+
+
+def is_ollama_available():
+    """Check if Ollama is running"""
+    try:
+        import requests
+
+        raw_base_url = get_db_setting(
+            "llm.ollama.url", "http://localhost:11434"
+        )
+        base_url = (
+            normalize_url(raw_base_url)
+            if raw_base_url
+            else "http://localhost:11434"
+        )
+        logger.info(f"Checking Ollama availability at {base_url}/api/tags")
+
+        try:
+            response = requests.get(f"{base_url}/api/tags", timeout=3.0)
+            if response.status_code == 200:
+                logger.info(
+                    f"Ollama is available. Status code: {response.status_code}"
+                )
+                # Log first 100 chars of response to debug
+                logger.info(f"Response preview: {str(response.text)[:100]}")
+                return True
+            else:
+                logger.warning(
+                    f"Ollama API returned status code: {response.status_code}"
+                )
+                return False
+        except requests.exceptions.RequestException as req_error:
+            logger.error(
+                f"Request error when checking Ollama: {str(req_error)}"
+            )
+            return False
+        except Exception:
+            logger.exception("Unexpected error when checking Ollama")
+            return False
+    except Exception:
+        logger.exception("Error in is_ollama_available")
+        return False
+
+
+def is_vllm_available():
+    """Check if VLLM capability is available"""
+    try:
+        import torch  # noqa: F401
+        import transformers  # noqa: F401
+
+        return True
+    except ImportError:
+        return False
+
+
+def is_lmstudio_available():
+    """Check if LM Studio is available"""
+    try:
+        import requests
+
+        lmstudio_url = get_db_setting(
+            "llm.lmstudio.url", "http://localhost:1234"
+        )
+        # LM Studio typically uses OpenAI-compatible endpoints
+        response = requests.get(f"{lmstudio_url}/v1/models", timeout=1.0)
+        return response.status_code == 200
+    except Exception:
+        return False
+
+
+def is_llamacpp_available():
+    """Check if LlamaCpp is available and configured"""
+    try:
+        from langchain_community.llms import LlamaCpp  # noqa: F401
+
+        model_path = get_db_setting("llm.llamacpp_model_path")
+        return bool(model_path) and os.path.exists(model_path)
+    except Exception:
+        return False
+
+
+@cache
+def get_available_providers():
+    """Return available model providers"""
+    providers = {}
+
+    if is_ollama_available():
+        providers["ollama"] = "Ollama (local models)"
+
+    if is_openai_available():
+        providers["openai"] = "OpenAI API"
+
+    if is_anthropic_available():
+        providers["anthropic"] = "Anthropic API"
+
+    if is_openai_endpoint_available():
+        providers["openai_endpoint"] = "OpenAI-compatible Endpoint"
+
+    if is_lmstudio_available():
+        providers["lmstudio"] = "LM Studio (local models)"
+
+    if is_llamacpp_available():
+        providers["llamacpp"] = "LlamaCpp (local models)"
+
+    # Check for VLLM capability
+    try:
+        import torch  # noqa: F401
+        import transformers  # noqa: F401
+
+        providers["vllm"] = "VLLM (local models)"
+    except ImportError:
+        pass
+
+    # Default fallback
+    if not providers:
+        providers["none"] = "No model providers available"
+
+    return providers
+
+
+def get_selected_llm_provider():
+    return get_db_setting("llm.provider", "ollama").lower()
+
+
+def get_llm(
+    model_name=None,
+    temperature=None,
+    provider=None,
+    openai_endpoint_url=None,
+    research_id=None,
+    research_context=None,
+):
     """
     Get LLM instance based on model name and provider.
 
@@ -34,6 +191,8 @@ def get_llm(model_name=None, temperature=None, provider=None, openai_endpoint_ur
         provider: Provider to use (if None, uses database setting)
         openai_endpoint_url: Custom endpoint URL to use (if None, uses database
             setting)
+        research_id: Optional research ID for token tracking
+        research_context: Optional research context for enhanced token tracking
 
     Returns:
         A LangChain LLM instance with automatic think-tag removal
@@ -47,6 +206,16 @@ def get_llm(model_name=None, temperature=None, provider=None, openai_endpoint_ur
     if provider is None:
         provider = get_db_setting("llm.provider", "ollama")
 
+    # Check if we're in testing mode and should use fallback
+    if os.environ.get("LDR_USE_FALLBACK_LLM", ""):
+        logger.info("LDR_USE_FALLBACK_LLM is set, using fallback model")
+        return wrap_llm_without_think_tags(
+            get_fallback_model(temperature),
+            research_id=research_id,
+            provider="fallback",
+            research_context=research_context,
+        )
+
     # Clean model name: remove quotes and extra whitespace
     if model_name:
         model_name = model_name.strip().strip("\"'").strip()
@@ -79,7 +248,8 @@ def get_llm(model_name=None, temperature=None, provider=None, openai_endpoint_ur
     if get_db_setting("llm.supports_max_tokens", True):
         # Use 80% of context window to leave room for prompts
         max_tokens = min(
-            int(get_db_setting("llm.max_tokens", 30000)), int(context_window_size * 0.8)
+            int(get_db_setting("llm.max_tokens", 30000)),
+            int(context_window_size * 0.8),
         )
         common_params["max_tokens"] = max_tokens
 
@@ -95,16 +265,28 @@ def get_llm(model_name=None, temperature=None, provider=None, openai_endpoint_ur
         llm = ChatAnthropic(
             model=model_name, anthropic_api_key=api_key, **common_params
         )
-        return wrap_llm_without_think_tags(llm)
+        return wrap_llm_without_think_tags(
+            llm,
+            research_id=research_id,
+            provider=provider,
+            research_context=research_context,
+        )
 
     elif provider == "openai":
         api_key = get_db_setting("llm.openai.api_key")
         if not api_key:
-            logger.warning("OPENAI_API_KEY not found. Falling back to default model.")
+            logger.warning(
+                "OPENAI_API_KEY not found. Falling back to default model."
+            )
             return get_fallback_model(temperature)
 
         llm = ChatOpenAI(model=model_name, api_key=api_key, **common_params)
-        return wrap_llm_without_think_tags(llm)
+        return wrap_llm_without_think_tags(
+            llm,
+            research_id=research_id,
+            provider=provider,
+            research_context=research_context,
+        )
 
     elif provider == "openai_endpoint":
         api_key = get_db_setting("llm.openai_endpoint.api_key")
@@ -126,7 +308,12 @@ def get_llm(model_name=None, temperature=None, provider=None, openai_endpoint_ur
             openai_api_base=openai_endpoint_url,
             **common_params,
         )
-        return wrap_llm_without_think_tags(llm)
+        return wrap_llm_without_think_tags(
+            llm,
+            research_id=research_id,
+            provider=provider,
+            research_context=research_context,
+        )
 
     elif provider == "vllm":
         try:
@@ -138,7 +325,12 @@ def get_llm(model_name=None, temperature=None, provider=None, openai_endpoint_ur
                 top_p=0.95,
                 temperature=temperature,
             )
-            return wrap_llm_without_think_tags(llm)
+            return wrap_llm_without_think_tags(
+                llm,
+                research_id=research_id,
+                provider=provider,
+                research_context=research_context,
+            )
         except Exception:
             logger.exception("Error loading VLLM model")
             return get_fallback_model(temperature)
@@ -146,7 +338,9 @@ def get_llm(model_name=None, temperature=None, provider=None, openai_endpoint_ur
     elif provider == "ollama":
         try:
             # Use the configurable Ollama base URL
-            raw_base_url = get_db_setting("llm.ollama.url", "http://localhost:11434")
+            raw_base_url = get_db_setting(
+                "llm.ollama.url", "http://localhost:11434"
+            )
             base_url = (
                 normalize_url(raw_base_url)
                 if raw_base_url
@@ -164,7 +358,9 @@ def get_llm(model_name=None, temperature=None, provider=None, openai_endpoint_ur
             import requests
 
             try:
-                logger.info(f"Checking if model '{model_name}' exists in Ollama")
+                logger.info(
+                    f"Checking if model '{model_name}' exists in Ollama"
+                )
                 response = requests.get(f"{base_url}/api/tags", timeout=3.0)
                 if response.status_code == 200:
                     # Handle both newer and older Ollama API formats
@@ -189,21 +385,30 @@ def get_llm(model_name=None, temperature=None, provider=None, openai_endpoint_ur
                         )
                         return get_fallback_model(temperature)
             except Exception:
-                logger.exception(f"Error checking for model '{model_name}' in Ollama")
+                logger.exception(
+                    f"Error checking for model '{model_name}' in Ollama"
+                )
                 # Continue anyway, let ChatOllama handle potential errors
 
             logger.info(
                 f"Creating ChatOllama with model={model_name}, base_url={base_url}"
             )
             try:
-                llm = ChatOllama(model=model_name, base_url=base_url, **common_params)
+                llm = ChatOllama(
+                    model=model_name, base_url=base_url, **common_params
+                )
                 # Test invoke to validate model works
                 logger.info("Testing Ollama model with simple invocation")
                 test_result = llm.invoke("Hello")
                 logger.info(
                     f"Ollama test successful. Response type: {type(test_result)}"
                 )
-                return wrap_llm_without_think_tags(llm)
+                return wrap_llm_without_think_tags(
+                    llm,
+                    research_id=research_id,
+                    provider=provider,
+                    research_context=research_context,
+                )
             except Exception:
                 logger.exception("Error creating or testing ChatOllama")
                 return get_fallback_model(temperature)
@@ -213,7 +418,9 @@ def get_llm(model_name=None, temperature=None, provider=None, openai_endpoint_ur
 
     elif provider == "lmstudio":
         # LM Studio supports OpenAI API format, so we can use ChatOpenAI directly
-        lmstudio_url = get_db_setting("llm.lmstudio.url", "http://localhost:1234")
+        lmstudio_url = get_db_setting(
+            "llm.lmstudio.url", "http://localhost:1234"
+        )
 
         llm = ChatOpenAI(
             model=model_name,
@@ -222,7 +429,12 @@ def get_llm(model_name=None, temperature=None, provider=None, openai_endpoint_ur
             temperature=temperature,
             max_tokens=max_tokens,  # Use calculated max_tokens based on context size
         )
-        return wrap_llm_without_think_tags(llm)
+        return wrap_llm_without_think_tags(
+            llm,
+            research_id=research_id,
+            provider=provider,
+            research_context=research_context,
+        )
 
     # Update the llamacpp section in get_llm function
     elif provider == "llamacpp":
@@ -230,7 +442,9 @@ def get_llm(model_name=None, temperature=None, provider=None, openai_endpoint_ur
         from langchain_community.llms import LlamaCpp
 
         # Get LlamaCpp connection mode from settings
-        connection_mode = get_db_setting("llm.llamacpp_connection_mode", "local")
+        connection_mode = get_db_setting(
+            "llm.llamacpp_connection_mode", "local"
+        )
 
         if connection_mode == "http":
             # Use HTTP client mode
@@ -270,10 +484,20 @@ def get_llm(model_name=None, temperature=None, provider=None, openai_endpoint_ur
                 verbose=True,
             )
 
-        return wrap_llm_without_think_tags(llm)
+        return wrap_llm_without_think_tags(
+            llm,
+            research_id=research_id,
+            provider=provider,
+            research_context=research_context,
+        )
 
     else:
-        return wrap_llm_without_think_tags(get_fallback_model(temperature))
+        return wrap_llm_without_think_tags(
+            get_fallback_model(temperature),
+            research_id=research_id,
+            provider=provider,
+            research_context=research_context,
+        )
 
 
 def get_fallback_model(temperature=None):
@@ -285,8 +509,36 @@ def get_fallback_model(temperature=None):
     )
 
 
-def wrap_llm_without_think_tags(llm):
-    """Create a wrapper class that processes LLM outputs with remove_think_tags"""
+def wrap_llm_without_think_tags(
+    llm, research_id=None, provider=None, research_context=None
+):
+    """Create a wrapper class that processes LLM outputs with remove_think_tags and token counting"""
+
+    # Import token counting functionality if research_id is provided
+    callbacks = []
+    if research_id is not None:
+        from ..metrics import TokenCounter
+
+        token_counter = TokenCounter()
+        token_callback = token_counter.create_callback(
+            research_id, research_context
+        )
+        # Set provider and model info on the callback
+        if provider:
+            token_callback.preset_provider = provider
+        # Try to extract model name from the LLM instance
+        if hasattr(llm, "model_name"):
+            token_callback.preset_model = llm.model_name
+        elif hasattr(llm, "model"):
+            token_callback.preset_model = llm.model
+        callbacks.append(token_callback)
+
+    # Add callbacks to the LLM if it supports them
+    if callbacks and hasattr(llm, "callbacks"):
+        if llm.callbacks is None:
+            llm.callbacks = callbacks
+        else:
+            llm.callbacks.extend(callbacks)
 
     class ProcessingLLMWrapper:
         def __init__(self, base_llm):
@@ -308,153 +560,3 @@ def wrap_llm_without_think_tags(llm):
             return getattr(self.base_llm, name)
 
     return ProcessingLLMWrapper(llm)
-
-
-def get_available_provider_types():
-    """Return available model providers"""
-    providers = {}
-
-    if is_ollama_available():
-        providers["ollama"] = "Ollama (local models)"
-
-    if is_openai_available():
-        providers["openai"] = "OpenAI API"
-
-    if is_anthropic_available():
-        providers["anthropic"] = "Anthropic API"
-
-    if is_openai_endpoint_available():
-        providers["openai_endpoint"] = "OpenAI-compatible Endpoint"
-
-    if is_lmstudio_available():
-        providers["lmstudio"] = "LM Studio (local models)"
-
-    if is_llamacpp_available():
-        providers["llamacpp"] = "LlamaCpp (local models)"
-
-    # Check for VLLM capability
-    try:
-        import torch  # noqa: F401
-        import transformers  # noqa: F401
-
-        providers["vllm"] = "VLLM (local models)"
-    except ImportError:
-        pass
-
-    # Default fallback
-    if not providers:
-        providers["none"] = "No model providers available"
-
-    return providers
-
-
-def is_openai_available():
-    """Check if OpenAI is available"""
-    try:
-        api_key = get_db_setting("llm.openai.api_key")
-        return bool(api_key)
-    except Exception:
-        return False
-
-
-def is_anthropic_available():
-    """Check if Anthropic is available"""
-    try:
-        api_key = get_db_setting("llm.anthropic.api_key")
-        return bool(api_key)
-    except Exception:
-        return False
-
-
-def is_openai_endpoint_available():
-    """Check if OpenAI endpoint is available"""
-    try:
-        api_key = get_db_setting("llm.openai_endpoint.api_key")
-        return bool(api_key)
-    except Exception:
-        return False
-
-
-def is_ollama_available():
-    """Check if Ollama is running"""
-    try:
-        import requests
-
-        raw_base_url = get_db_setting("llm.ollama.url", "http://localhost:11434")
-        base_url = (
-            normalize_url(raw_base_url) if raw_base_url else "http://localhost:11434"
-        )
-        logger.info(f"Checking Ollama availability at {base_url}/api/tags")
-
-        try:
-            response = requests.get(f"{base_url}/api/tags", timeout=3.0)
-            if response.status_code == 200:
-                logger.info(f"Ollama is available. Status code: {response.status_code}")
-                # Log first 100 chars of response to debug
-                logger.info(f"Response preview: {str(response.text)[:100]}")
-                return True
-            else:
-                logger.warning(
-                    f"Ollama API returned status code: {response.status_code}"
-                )
-                return False
-        except requests.exceptions.RequestException as req_error:
-            logger.error(f"Request error when checking Ollama: {str(req_error)}")
-            return False
-        except Exception:
-            logger.exception("Unexpected error when checking Ollama")
-            return False
-    except Exception:
-        logger.exception("Error in is_ollama_available")
-        return False
-
-
-def is_vllm_available():
-    """Check if VLLM capability is available"""
-    try:
-        import torch  # noqa: F401
-        import transformers  # noqa: F401
-
-        return True
-    except ImportError:
-        return False
-
-
-def is_lmstudio_available():
-    """Check if LM Studio is available"""
-    try:
-        import requests
-
-        lmstudio_url = get_db_setting("llm.lmstudio.url", "http://localhost:1234")
-        # LM Studio typically uses OpenAI-compatible endpoints
-        response = requests.get(f"{lmstudio_url}/v1/models", timeout=1.0)
-        return response.status_code == 200
-    except Exception:
-        return False
-
-
-def is_llamacpp_available():
-    """Check if LlamaCpp is available and configured"""
-    try:
-        from langchain_community.llms import LlamaCpp  # noqa: F401
-
-        model_path = get_db_setting("llm.llamacpp_model_path")
-        return bool(model_path) and os.path.exists(model_path)
-    except Exception:
-        return False
-
-
-def get_available_providers():
-    """Get dictionary of available providers"""
-    return get_available_provider_types()
-
-
-AVAILABLE_PROVIDERS = get_available_providers()
-selected_provider = get_db_setting("llm.provider", "ollama").lower()
-
-# Log which providers are available
-logger.info(f"Available providers: {list(AVAILABLE_PROVIDERS.keys())}")
-
-# Check if selected provider is available
-if selected_provider not in AVAILABLE_PROVIDERS and selected_provider != "none":
-    logger.warning(f"Selected provider {selected_provider} is not available.")

local_deep_research/config/search_config.py

@@ -2,13 +2,20 @@
 from loguru import logger
 
 from ..utilities.db_utils import get_db_setting
-from ..web_search_engines.search_engine_factory import get_search as factory_get_search
+from ..web_search_engines.search_engine_factory import (
+    get_search as factory_get_search,
+)
 from .llm_config import get_llm
 
 # Whether to check the quality search results using the LLM.
 QUALITY_CHECK_DDG_URLS = True
-# Whether to only retrieve snippets instead of full search results.
-SEARCH_SNIPPETS_ONLY = get_db_setting("search.snippets_only", True)
+
+
+def get_search_snippets_only_setting():
+    """
+    Lazily retrieve the 'search.snippets_only' setting.
+    """
+    return get_db_setting("search.snippets_only", True)
 
 
 # Expose get_search function
@@ -36,9 +43,11 @@ def get_search(search_tool=None, llm_instance=None):
         "region": get_db_setting("search.region", "wt-wt"),
         "time_period": get_db_setting("search.time_period", "all"),
         "safe_search": get_db_setting("search.safe_search", True),
-        "search_snippets_only": SEARCH_SNIPPETS_ONLY,
+        "search_snippets_only": get_search_snippets_only_setting(),
         "search_language": get_db_setting("search.search_language", "English"),
-        "max_filtered_results": get_db_setting("search.max_filtered_results", 5),
+        "max_filtered_results": get_db_setting(
+            "search.max_filtered_results", 5
+        ),
     }
 
     # Log NULL parameters for debugging

local_deep_research/defaults/__init__.py

@@ -6,7 +6,6 @@ configuration files and resources used throughout the application.
 """
 
 import logging
-import os
 from pathlib import Path
 
 logger = logging.getLogger(__name__)

local_deep_research/metrics/__init__.py

@@ -0,0 +1,13 @@
+"""Metrics module for tracking LLM usage and token counts."""
+
+from .database import get_metrics_db
+from .db_models import ModelUsage, TokenUsage
+from .token_counter import TokenCounter, TokenCountingCallback
+
+__all__ = [
+    "TokenCounter",
+    "TokenCountingCallback",
+    "TokenUsage",
+    "ModelUsage",
+    "get_metrics_db",
+]

local_deep_research/metrics/database.py

@@ -0,0 +1,58 @@
+"""Database utilities for metrics module with SQLAlchemy."""
+
+from contextlib import contextmanager
+from typing import Generator
+
+from loguru import logger
+from sqlalchemy import create_engine
+from sqlalchemy.orm import Session, sessionmaker
+
+from ..utilities.db_utils import DB_PATH
+from .db_models import Base
+
+
+class MetricsDatabase:
+    """Database manager for metrics using SQLAlchemy."""
+
+    def __init__(self):
+        # Use the same database as the rest of the app
+        self.engine = create_engine(
+            f"sqlite:///{DB_PATH}", connect_args={"check_same_thread": False}
+        )
+        self.SessionLocal = sessionmaker(
+            bind=self.engine, autocommit=False, autoflush=False
+        )
+        self._init_database()
+
+    def _init_database(self):
+        """Initialize database tables for metrics."""
+        try:
+            Base.metadata.create_all(self.engine)
+            logger.info("Metrics tables initialized successfully")
+        except Exception as e:
+            logger.exception(f"Error initializing metrics tables: {e}")
+
+    @contextmanager
+    def get_session(self) -> Generator[Session, None, None]:
+        """Get a database session with automatic cleanup."""
+        session = self.SessionLocal()
+        try:
+            yield session
+            session.commit()
+        except Exception:
+            session.rollback()
+            raise
+        finally:
+            session.close()
+
+
+# Singleton instance
+_metrics_db = None
+
+
+def get_metrics_db() -> MetricsDatabase:
+    """Get the singleton metrics database instance."""
+    global _metrics_db
+    if _metrics_db is None:
+        _metrics_db = MetricsDatabase()
+    return _metrics_db