replypier/backend/app/utils/agno_utils.py

"""
Copyright 2024-2026 ReplyPier

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import re
from agno.agent import Agent
from agno.models.openai import OpenAIChat
from openai import AsyncOpenAI, OpenAI
from app.core.logger import get_logger
from app.core.config import settings
from app.core.outbound_http import guarded_async_http_client, guarded_http_client
from app.core.outbound_policy import check_outbound_target, public_only
from app.utils.agno_patches import apply_agno_patches
from typing import Dict, Any, Optional, List
from fastapi import HTTPException

logger = get_logger(__name__)

# Every code path that builds an agno model goes through this module, so this
# is the single choke point where the 1.7.6 runtime patches get applied.
apply_agno_patches()


class PinnedOpenAIChat(OpenAIChat):
    """OpenAIChat for a user-supplied base URL under OUTBOUND_URL_POLICY=
    public_only: every connection, redirect included, goes to a vetted public
    address. agno 1.7.6 takes a single http_client for both its sync and
    async OpenAI clients, so each getter supplies its own guarded one, fresh
    per call like agno's default."""

    def get_client(self):
        # follow_redirects matches the OpenAI SDK's own default client.
        return OpenAI(**{**self._get_client_params(),
                         "http_client": guarded_http_client(follow_redirects=True)})

    def get_async_client(self):
        return AsyncOpenAI(**{**self._get_client_params(),
                              "http_client": guarded_async_http_client(follow_redirects=True)})


def create_model(model_type: str, api_key: str, model_name: str, max_tokens: int = 1000, response_format: Optional[Dict[str, Any]] = None, base_url: Optional[str] = None) -> Any:
    """
    Create and return the specified model based on model_type.

    Args:
        model_type: The type of model (OPENAI, GROQ, etc.)
        api_key: The API key
        model_name: The name/ID of the model
        max_tokens: Maximum tokens for model output
        response_format: Optional response format specification
        base_url: Custom API endpoint, used by OPENAI_COMPATIBLE

    Returns:
        The initialized model object

    Raises:
        HTTPException: If the model type is not supported
    """
    model_type = model_type.upper()

    try:
        if model_type == 'OPENAI' or model_type == 'REPLYPIER': # own model for enterprise customers
            if response_format:
                return OpenAIChat(api_key=api_key, id=model_name, max_tokens=max_tokens, response_format=response_format)
            else:
                return OpenAIChat(api_key=api_key, id=model_name, max_tokens=max_tokens)
        elif model_type == 'OPENAI_COMPATIBLE':
            if not base_url:
                raise ValueError("base_url is required for OPENAI_COMPATIBLE")
            # Re-checked on every build, not only on save: DNS can change.
            check_outbound_target(base_url, "Base URL")
            base_url = re.sub(r'/chat/completions/?$', '', base_url.strip(), flags=re.IGNORECASE).rstrip('/')
            model_cls = PinnedOpenAIChat if public_only() else OpenAIChat
            if response_format:
                return model_cls(api_key=api_key, id=model_name, base_url=base_url, max_tokens=max_tokens, response_format=response_format)
            else:
                return model_cls(api_key=api_key, id=model_name, base_url=base_url, max_tokens=max_tokens)
        elif model_type == 'ANTHROPIC':
            from agno.models.anthropic import Claude
            return Claude(api_key=api_key, id=model_name, max_tokens=max_tokens)
        elif model_type == 'DEEPSEEK':
            from agno.models.deepseek import DeepSeek
            return DeepSeek(api_key=api_key, id=model_name, max_tokens=max_tokens)
        elif model_type == 'GOOGLE':
            from agno.models.google import Gemini
            # Gemini's agno model exposes max_output_tokens, not max_tokens.
            return Gemini(api_key=api_key, id=model_name, max_output_tokens=max_tokens)
        elif model_type == 'GOOGLEVERTEX':
            from agno.models.vertexai import Gemini
            return Gemini(api_key=api_key, id=model_name, max_tokens=max_tokens)
        elif model_type == 'GROQ':
            from agno.models.groq import Groq
            
            # Create a patched Groq model that handles the response_format + tools conflict
            class PatchedGroq(Groq):
                def __post_init__(self):
                    """Set supports_native_structured_outputs after initialization"""
                    super().__post_init__() if hasattr(super(), '__post_init__') else None
                    # Override to indicate Groq supports structured outputs via response_format
                    object.__setattr__(self, 'supports_native_structured_outputs', True)
                
                def get_request_params(self, response_format=None, tools=None, tool_choice=None):
                    """
                    Override get_request_params to handle the conflict between response_format and tools.
                    Groq API doesn't allow JSON response_format when tools are present.
                    """
                    # Get the base parameters from parent class
                    base_params = {
                        "frequency_penalty": self.frequency_penalty,
                        "logit_bias": self.logit_bias,
                        "logprobs": self.logprobs,
                        "max_tokens": self.max_tokens,
                        "presence_penalty": self.presence_penalty,
                        "seed": self.seed,
                        "stop": self.stop,
                        "temperature": self.temperature,
                        "top_logprobs": self.top_logprobs,
                        "top_p": self.top_p,
                        "user": self.user,
                        "extra_headers": self.extra_headers,
                        "extra_query": self.extra_query,
                    }
                    
                    # Filter out None values
                    request_params = {k: v for k, v in base_params.items() if v is not None}
                    
                    # Add tools first
                    if tools is not None:
                        request_params["tools"] = tools
                        if tool_choice is not None:
                            request_params["tool_choice"] = tool_choice
                        # Don't add response_format when tools are present to avoid conflict
                        logger.debug("Groq: Skipping response_format due to tools being present")
                    else:
                        # Only add response_format if no tools are present
                        if response_format is not None:
                            request_params["response_format"] = response_format
                    
                    # Add additional request params if provided
                    if self.request_params:
                        request_params.update(self.request_params)
                    
                    return request_params
            
            # Groq model doesn't support response_format parameter in constructor
            return PatchedGroq(api_key=api_key, id=model_name, max_tokens=max_tokens)
        elif model_type == 'MISTRAL':
            from agno.models.mistral import MistralChat
            return MistralChat(api_key=api_key, id=model_name, max_tokens=max_tokens)
        elif model_type == 'HUGGINGFACE':
            from agno.models.huggingface import HuggingFace
            return HuggingFace(api_key=api_key, id=model_name, max_tokens=max_tokens)
        elif model_type == 'OLLAMA':
            from agno.models.ollama import Ollama
            return Ollama(id=model_name)
        elif model_type == 'XAI':
            from agno.models.xai import xAI
            return xAI(api_key=api_key, id=model_name, max_tokens=max_tokens)
        else:
            raise ValueError(f"Unsupported model type: {model_type}")
    except ImportError as e:
        logger.error(f"Import error when creating model type {model_type}: {str(e)}")
        raise HTTPException(status_code=400, detail=f"Model type {model_type} is not available in this installation")
    except Exception as e:
        logger.error(f"Error creating model type {model_type}: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Failed to initialize model: {str(e)}")

async def test_model_api_key(api_key: str, model_type: str, model_name: str, base_url: Optional[str] = None) -> bool:
    """
    Test if the API key is valid for the given model type.

    Args:
        api_key: The API key to test
        model_type: The type of model
        model_name: The name of the model
        base_url: Custom API endpoint, used by OPENAI_COMPATIBLE

    Returns:
        bool: True if the API key is valid, False otherwise
    """
    try:
        # Create a simple model and agent for testing
        model = create_model(model_type, api_key, model_name, base_url=base_url)
        test_agent = Agent(
            name="Test Agent",
            model=model,
            instructions="You are a test agent. Just respond with 'Test successful.'",
            debug_mode=False
        )
        
        # Run a simple test query
        await test_agent.arun(message="This is a test message.")
        return True
    except Exception as e:
        logger.error(f"API key test failed: {str(e)}")
        return False