local-deep-research/examples/optimization/update_llm_config.py

#!/usr/bin/env python
"""
Update LLM configuration in the database for benchmarks.

This script updates the LLM configuration in the database to ensure
consistent behavior when running benchmarks with different LLM models.

Usage:
    # Install dependencies with PDM
    cd /path/to/local-deep-research
    pdm install

    # Run the script with PDM
    pdm run python examples/optimization/update_llm_config.py --model "google/gemini-2.0-flash" --provider "openai_endpoint" --endpoint "https://openrouter.ai/api/v1" --api-key "your-api-key"

    # Or to reset to default configuration
    pdm run python examples/optimization/update_llm_config.py --reset
"""

import argparse
import sys
from pathlib import Path
from typing import Optional

from loguru import logger

# Add the src directory to the Python path
project_root = str(Path(__file__).parent.parent.parent.resolve())
sys.path.insert(0, str(Path(project_root) / "src"))


def update_llm_configuration(
    model_name: Optional[str] = None,
    provider: Optional[str] = None,
    endpoint_url: Optional[str] = None,
    api_key: Optional[str] = None,
    temperature: Optional[float] = None,
    reset: bool = False,
) -> bool:
    """
    Update LLM configuration in the database.

    Args:
        model_name: LLM model name to set
        provider: LLM provider to set
        endpoint_url: Endpoint URL for OpenRouter or similar services
        api_key: API key for the provider
        temperature: Temperature setting for the LLM
        reset: If True, reset to default configuration

    Returns:
        True if successful, False otherwise
    """
    # Import database utility functions
    try:
        from local_deep_research.utilities.db_utils import (
            get_db_setting,
            update_db_setting,
        )
    except ImportError:
        logger.exception(
            "Could not import database utilities. Make sure you're in the correct directory."
        )
        return False

    # Default configuration
    default_config = {
        "llm.model": "gemma3:12b",
        "llm.provider": "ollama",
        "llm.temperature": 0.7,
        "llm.max_tokens": 30000,
    }

    try:
        if reset:
            # Reset to default configuration
            logger.info("Resetting LLM configuration to defaults")

            for key, value in default_config.items():
                update_db_setting(key, value)
                logger.info(f"Reset {key} to {value}")

            # Clear API keys
            update_db_setting("llm.openai_endpoint.api_key", "")
            update_db_setting("llm.openai_endpoint.url", "")

            logger.info("LLM configuration reset to defaults")
            return True

        # Update model and provider if provided
        if model_name:
            update_db_setting("llm.model", model_name)
            logger.info(f"Updated llm.model to {model_name}")

        if provider:
            update_db_setting("llm.provider", provider)
            logger.info(f"Updated llm.provider to {provider}")

        if temperature is not None:
            update_db_setting("llm.temperature", temperature)
            logger.info(f"Updated llm.temperature to {temperature}")

        # Handle provider-specific settings
        if provider == "openai_endpoint":
            if endpoint_url:
                update_db_setting("llm.openai_endpoint.url", endpoint_url)
                logger.info(
                    f"Updated llm.openai_endpoint.url to {endpoint_url}"
                )

            if api_key:
                update_db_setting("llm.openai_endpoint.api_key", api_key)
                logger.info(
                    "Updated llm.openai_endpoint.api_key (value hidden)"
                )

        elif provider == "openai":
            if api_key:
                update_db_setting("llm.openai.api_key", api_key)
                logger.info("Updated llm.openai.api_key (value hidden)")

        elif provider == "anthropic":
            if api_key:
                update_db_setting("llm.anthropic.api_key", api_key)
                logger.info("Updated llm.anthropic.api_key (value hidden)")

        # Verify settings were updated
        current_model = get_db_setting("llm.model")
        current_provider = get_db_setting("llm.provider")

        logger.info(
            f"Current LLM configuration: model={current_model}, provider={current_provider}"
        )

        if provider == "openai_endpoint":
            endpoint = get_db_setting("llm.openai_endpoint.url")
            has_key = bool(get_db_setting("llm.openai_endpoint.api_key"))
            logger.info(f"OpenAI Endpoint URL: {endpoint}")
            logger.info(f"Has API key: {has_key}")

        return True

    except Exception:
        logger.exception("Error updating LLM configuration")
        return False


def main():
    parser = argparse.ArgumentParser(
        description="Update LLM configuration in the database"
    )

    # Configuration options
    parser.add_argument("--model", help="LLM model name")
    parser.add_argument(
        "--provider",
        help="LLM provider (e.g., 'anthropic', 'openai', 'openai_endpoint')",
    )
    parser.add_argument(
        "--endpoint", help="Endpoint URL for OpenRouter or similar services"
    )
    parser.add_argument("--api-key", help="API key for the provider")
    parser.add_argument(
        "--temperature", type=float, help="Temperature setting for the LLM"
    )

    # Reset option
    parser.add_argument(
        "--reset", action="store_true", help="Reset to default configuration"
    )

    args = parser.parse_args()

    # Check if any argument is provided
    if not any(
        [
            args.model,
            args.provider,
            args.endpoint,
            args.api_key,
            args.temperature,
            args.reset,
        ]
    ):
        parser.print_help()
        return 1

    # Update LLM configuration
    success = update_llm_configuration(
        model_name=args.model,
        provider=args.provider,
        endpoint_url=args.endpoint,
        api_key=args.api_key,
        temperature=args.temperature,
        reset=args.reset,
    )

    return 0 if success else 1


if __name__ == "__main__":
    sys.exit(main())