LLM Factory

def create_llm_from_provider(
    provider_type: str,
    model_id: str,
    api_key: Optional[str] = None,
    base_url: Optional[str] = None,
    provider_config: Optional[Dict[str, Any]] = None,
    temperature: float = DEFAULT_TEMPERATURE,
    max_retries: int = DEFAULT_MAX_RETRIES,
) -> BaseChatModel:
    """Create a LangChain chat model from provider configuration.
    
    Args:
        provider_type: Type of provider (anthropic, openai, google, custom).
        model_id: Model identifier (e.g., 'claude-3-opus-20240229', 'gpt-4').
        api_key: API key for the provider (decrypted).
        base_url: Optional base URL for custom endpoints or proxies.
        provider_config: Additional provider-specific configuration.
        temperature: LLM temperature setting.
        max_retries: Maximum retry attempts for API calls.
    
    Returns:
        Configured BaseChatModel instance.
    
    Raises:
        ValueError: If provider type is unsupported or required dependencies missing.
    """
    provider_config = provider_config or {}
    
    if provider_type == "anthropic":
        if not _anthropic_available:
            raise ValueError(
                "Anthropic provider requires langchain-anthropic package. "
                "Install with: pip install langchain-anthropic"
            )
        if not api_key:
            raise ValueError("Anthropic provider requires an API key")
        
        kwargs = {
            "model": model_id,
            "api_key": api_key,
            "temperature": temperature,
            "max_retries": max_retries,
        }
        if base_url:
            kwargs["base_url"] = base_url
        
        logger.info(f"Creating Anthropic LLM with model: {model_id}")
        return ChatAnthropic(**kwargs)
    
    elif provider_type == "openai":
        if not _openai_available:
            raise ValueError(
                "OpenAI provider requires langchain-openai package. "
                "Install with: pip install langchain-openai"
            )
        if not api_key:
            raise ValueError("OpenAI provider requires an API key")
        
        kwargs = {
            "model": model_id,
            "api_key": api_key,
            "temperature": temperature,
            "max_retries": max_retries,
        }
        if base_url:
            kwargs["base_url"] = base_url
        
        # Handle organization ID if provided
        if provider_config.get("organization_id"):
            kwargs["organization"] = provider_config["organization_id"]
        
        logger.info(f"Creating OpenAI LLM with model: {model_id}")
        return ChatOpenAI(**kwargs)
    
    elif provider_type == "google":
        if not _google_available:
            raise ValueError(
                "Google provider requires langchain-google-genai package. "
                "Install with: pip install langchain-google-genai"
            )
        if not api_key:
            raise ValueError("Google provider requires an API key")
        
        kwargs = {
            "model": model_id,
            "google_api_key": api_key,
            "temperature": temperature,
            "max_retries": max_retries,
        }
        
        logger.info(f"Creating Google LLM with model: {model_id}")
        return ChatGoogleGenerativeAI(**kwargs)
    
    elif provider_type == "custom":
        # Custom provider uses Ollama-compatible or OpenAI-compatible API
        if not base_url:
            raise ValueError("Custom provider requires a base URL")
        
        auth_type = provider_config.get("auth_type", "none")
        
        # Check if it's OpenAI-compatible (has /v1/chat/completions endpoint)
        is_openai_compatible = provider_config.get("openai_compatible", False)
        
        if is_openai_compatible and _openai_available:
            # Use OpenAI client for OpenAI-compatible APIs
            kwargs = {
                "model": model_id,
                "base_url": base_url,
                "temperature": temperature,
                "max_retries": max_retries,
            }
            
            if auth_type == "bearer" and api_key:
                kwargs["api_key"] = api_key
            elif auth_type == "none":
                # Some local models don't need auth
                kwargs["api_key"] = "not-needed"
            
            logger.info(f"Creating OpenAI-compatible LLM at {base_url} with model: {model_id}")
            return ChatOpenAI(**kwargs)
        else:
            # Default to Ollama-compatible API
            kwargs = {
                "model": model_id,
                "base_url": base_url,
                "temperature": temperature,
                "max_retries": max_retries,
            }
            
            logger.info(f"Creating Ollama-compatible LLM at {base_url} with model: {model_id}")
            return ChatOllama(**kwargs)
    
    else:
        raise ValueError(f"Unsupported provider type: {provider_type}")

def get_default_llm(
    temperature: float = DEFAULT_TEMPERATURE,
    max_retries: int = DEFAULT_MAX_RETRIES,
) -> BaseChatModel:
    """Get the default Ollama LLM from environment configuration.
    
    This is the fallback when no provider is configured in the database.
    
    Args:
        temperature: LLM temperature setting.
        max_retries: Maximum retry attempts.
    
    Returns:
        ChatOllama instance configured from environment variables.
    """
    logger.info(
        f"Creating default Ollama LLM: {config.DEFAULT_OLLAMA_MODEL} "
        f"at {config.OLLAMA_API_URL}"
    )
    return ChatOllama(
        model=config.DEFAULT_OLLAMA_MODEL,
        temperature=temperature,
        base_url=config.OLLAMA_API_URL,
        max_retries=max_retries,
    )

def create_llm_from_state(
    state: Dict[str, Any],
    temperature: Optional[float] = None,
) -> BaseChatModel:
    """Create an LLM instance based on agent state configuration.
    
    This function extracts provider configuration from the agent state
    and creates the appropriate LLM instance.
    
    Args:
        state: Agent state dictionary containing optional provider config:
            - provider_type: Type of provider
            - model_id: Model identifier
            - api_key: Decrypted API key
            - base_url: Provider base URL
            - provider_config: Additional config
            - temperature: Optional temperature override
        temperature: Override temperature (uses state value or default if None).
    
    Returns:
        Configured BaseChatModel instance.
    """
    # Check if provider configuration is present in state
    provider_type = state.get("provider_type")
    model_id = state.get("model_id")
    
    if not provider_type or not model_id:
        # No provider configured, use default
        logger.debug("No provider configuration in state, using default Ollama")
        return get_default_llm(
            temperature=temperature or state.get("temperature", DEFAULT_TEMPERATURE)
        )
    
    # Use provider from state
    return create_llm_from_provider(
        provider_type=provider_type,
        model_id=model_id,
        api_key=state.get("api_key"),
        base_url=state.get("base_url"),
        provider_config=state.get("provider_config", {}),
        temperature=temperature or state.get("temperature", DEFAULT_TEMPERATURE),
    )

state = {
    "provider_type": "anthropic",
    "model_id": "claude-3-opus-20240229",
    "api_key": "sk-...",  # Decrypted
    "base_url": None,
    "provider_config": {},
    "temperature": 0.5,
}

llm = create_llm_from_provider(
    provider_type="anthropic",
    model_id="claude-3-opus-20240229",
    api_key="sk-ant-...",
    temperature=0.33,
    max_retries=2,
)

llm = create_llm_from_provider(
    provider_type="openai",
    model_id="gpt-4-turbo",
    api_key="sk-...",
    provider_config={
        "organization_id": "org-..."
    },
    temperature=0.33,
)

llm = create_llm_from_provider(
    provider_type="google",
    model_id="gemini-1.5-pro",
    api_key="AIza...",
    temperature=0.33,
)

llm = create_llm_from_provider(
    provider_type="custom",
    model_id="custom-model-name",
    base_url="https://api.custom-provider.com/v1",
    provider_config={
        "openai_compatible": True,
        "auth_type": "bearer"
    },
    api_key="custom-api-key",
    temperature=0.33,
)

llm = create_llm_from_provider(
    provider_type="custom",
    model_id="mixtral:latest",
    base_url="http://custom-ollama-server:11434",
    provider_config={
        "openai_compatible": False,
        "auth_type": "none"
    },
    temperature=0.33,
)

# Lazy imports for optional providers
_anthropic_available = False
_openai_available = False
_google_available = False

try:
    from langchain_anthropic import ChatAnthropic
    _anthropic_available = True
except ImportError:
    logger.debug("langchain-anthropic not installed, Anthropic provider unavailable")

try:
    from langchain_openai import ChatOpenAI
    _openai_available = True
except ImportError:
    logger.debug("langchain-openai not installed, OpenAI provider unavailable")

try:
    from langchain_google_genai import ChatGoogleGenerativeAI
    _google_available = True
except ImportError:
    logger.debug("langchain-google-genai not installed, Google provider unavailable")

# Global LLM instance cache (refreshed when provider config changes)
_cached_llm: Optional[BaseChatModel] = None
_cached_llm_config_hash: Optional[str] = None

def set_llm_from_config(
    provider_type: Optional[str] = None,
    model_id: Optional[str] = None,
    api_key: Optional[str] = None,
    base_url: Optional[str] = None,
    provider_config: Optional[Dict[str, Any]] = None,
    temperature: Optional[float] = None,
) -> None:
    """Set the LLM instance from provider configuration.
    
    This function should be called before invoking the graph to configure
    which LLM provider to use. The LLM is cached and reused until the
    configuration changes.
    """
    global _cached_llm, _cached_llm_config_hash
    
    # Create a hash of the config to detect changes
    config_tuple = (provider_type, model_id, base_url, temperature)
    config_hash = str(hash(config_tuple))
    
    # Only recreate if config changed
    if _cached_llm_config_hash == config_hash and _cached_llm is not None:
        return
    
    # Create new LLM from config
    if provider_type and model_id:
        from src.copilot.llm_factory import create_llm_from_provider
        _cached_llm = create_llm_from_provider(
            provider_type=provider_type,
            model_id=model_id,
            api_key=api_key,
            base_url=base_url,
            provider_config=provider_config or {},
            temperature=temperature or config.DEFAULT_LLM_TEMPERATURE,
        )
        logger.info(f"Configured LLM: {provider_type}/{model_id}")
    else:
        from src.copilot.llm_factory import get_default_llm
        _cached_llm = get_default_llm(
            temperature=temperature or config.DEFAULT_LLM_TEMPERATURE
        )
        logger.info("Using default Ollama LLM")
    
    _cached_llm_config_hash = config_hash

def get_configured_llm() -> BaseChatModel:
    """Get the currently configured LLM instance.
    
    Returns:
        The cached LLM instance, or creates a default one if not configured.
    """
    global _cached_llm
    if _cached_llm is None:
        from src.copilot.llm_factory import get_default_llm
        _cached_llm = get_default_llm()
    return _cached_llm

# Configure LLM before creating graph
set_llm_from_config(
    provider_type="anthropic",
    model_id="claude-3-opus-20240229",
    api_key=decrypted_api_key,
    temperature=0.5,
)

# Create and invoke graph
app = create_agent_graph()
result = app.invoke({"messages": [("user", "What is INC-2025-08-24-001?")]})

# Default LLM settings
DEFAULT_TEMPERATURE = 0.33
DEFAULT_MAX_RETRIES = 2

if provider_type == "anthropic":
    if not _anthropic_available:
        raise ValueError(
            "Anthropic provider requires langchain-anthropic package. "
            "Install with: pip install langchain-anthropic"
        )

if not api_key:
    raise ValueError("Anthropic provider requires an API key")

else:
    raise ValueError(f"Unsupported provider type: {provider_type}")