clients.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
LLMEx Testing Tool - LLM Client Implementations
=============================================
Client implementations for different LLM APIs.
"""

import abc
import argparse
import json
import logging
import os
import time
from typing import Dict, List, Optional, Any
import requests

from config import config_manager

logger = logging.getLogger(__name__)


class LLMClient(abc.ABC):
    """Abstract base class for LLM API clients."""
    
    @abc.abstractmethod
    def query(self, prompt: str, **kwargs) -> str:
        """Send a query to the LLM and return the response."""
        pass
    
    @abc.abstractmethod
    def check_connection(self) -> bool:
        """Check if the connection to the LLM is working."""
        pass


class BaseClient(LLMClient):
    """Base class for LLM API clients with common functionality."""
    
    def __init__(self, args: Optional[argparse.Namespace] = None, 
                 config: Optional[Dict[str, Any]] = None):
        self.session = requests.Session()
        self.timeout = 60
        self.rate_limit_enabled = False
        self.max_requests_per_minute = 60
        self.cooldown_period = 5
        self.request_timestamps = []
        
        self._load_config(args, config)
        
    def _load_config(self, args: Optional[argparse.Namespace] = None, 
                    config: Optional[Dict[str, Any]] = None):
        """Load configuration from args and config."""
        if args and getattr(args, 'timeout', None) is not None:
            self.timeout = args.timeout
        elif config and config.get('timeout') is not None:
            self.timeout = config['timeout']

        if args and getattr(args, 'cooldown_period', None) is not None:
            self.cooldown_period = args.cooldown_period
        elif config and 'security_testing' in config and 'rate_limit_protection' in config['security_testing']:
            self.cooldown_period = config['security_testing']['rate_limit_protection'].get('cooldown_period', 5)        

        if config and 'rate_limiting' in config:
            self.rate_limit_enabled = config['rate_limiting'].get('enabled', False)
            self.max_requests_per_minute = config['rate_limiting'].get('max_requests_per_minute', 60)
        
        if config and 'security_testing' in config and 'rate_limit_protection' in config['security_testing']:
            rate_limit_protection = config['security_testing']['rate_limit_protection']
            self.max_requests_per_minute = rate_limit_protection.get('max_requests_per_minute', self.max_requests_per_minute)
            self.cooldown_period = rate_limit_protection.get('cooldown_period', 5)
            self.rate_limit_enabled = True
            
        self.args = args
        self.config = config  
        
    def _sanitize_response(self, response: str) -> str:
        """Sanitize the response according to config settings."""
        sanitization_config = config_manager.main_config.get('security_testing', {}).get('response_sanitization', {})
        max_length = sanitization_config.get('max_response_length', 5000)
        truncate = sanitization_config.get('truncate_if_exceed', True)

        if truncate and len(response) > max_length:
            return response[:max_length]
        return response
    
    def _check_rate_limit(self):
        """Check if rate limit is exceeded and wait if necessary."""
        if not self.rate_limit_enabled:
            return
        
        now = time.time()
        self.request_timestamps = [t for t in self.request_timestamps if now - t < 60]
    
        if len(self.request_timestamps) >= self.max_requests_per_minute:
            wait_time = self.cooldown_period
            logger.warning(f"Rate limit exceeded. Waiting {wait_time} seconds.")
            time.sleep(wait_time)
            self.request_timestamps = self.request_timestamps[-(self.max_requests_per_minute // 2):]
    
        self.request_timestamps.append(now)


class OpenAIClient(BaseClient):
    """Client for the OpenAI API."""
    
    def __init__(self, api_key: str, 
                 args: Optional[argparse.Namespace] = None, 
                 config: Optional[Dict[str, Any]] = None) -> None:
        super().__init__(args, config)
        
        self.api_key = api_key
        
        openai_config = (config or {}).get('api_providers', {}).get('openai', {})
        self.base_url = openai_config.get('base_url', "https://api.openai.com/v1/chat/completions")
        
        if args and getattr(args, 'model', None):
            self.model = args.model
        else:
            self.model = (config or {}).get('default_model', 'gpt-4')
            
        allowed_models = openai_config.get('models', [])
        if allowed_models and self.model not in allowed_models:
            logger.warning(f"Model '{self.model}' not in allowed models list: {allowed_models}")
            if allowed_models:
                logger.warning(f"Falling back to first allowed model: {allowed_models[0]}")
                self.model = allowed_models[0]

        self.default_params: Dict[str, Any] = openai_config.get('default_params', {
            'temperature': 0.7,
            'max_tokens': 1000,
            'top_p': 1.0,
            'frequency_penalty': 0.0,
            'presence_penalty': 0.0
        })

        self.session.headers.update({
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        })

        logger.info(f"OpenAIClient: configured with model={self.model}, timeout={self.timeout}")
    
    def query(self, prompt: str, **kwargs) -> str:
        """Send a query to the OpenAI API and return the response."""
        self._check_rate_limit()
        
        logger.debug(f"Query to OpenAI API:")
        logger.debug(f"- Model: {self.model}")
        logger.debug(f"- Prompt length: {len(prompt)} characters")
        
        safe_params = {k: v for k, v in kwargs.items() if 'key' not in k.lower() and 'token' not in k.lower()}
        logger.debug(f"- Parameters: {safe_params}")

        query_params = {**self.default_params, **kwargs}
        payload = {
            "model": self.model,
            "messages": [{"role": "user", "content": prompt}],
            **query_params
        }

        try:
            logger.debug(f"Sending request to OpenAI API at {self.base_url}")
            start_time = time.time()
            response = self.session.post(self.base_url, json=payload, timeout=self.timeout)
            elapsed_time = time.time() - start_time
            logger.debug(f"Response received in {elapsed_time:.2f} seconds")
        
            response.raise_for_status()

            response_json = response.json()
            if 'choices' not in response_json or not response_json['choices']:
                logger.error(f"Unexpected response format: {response_json}")
                return "Error: Unexpected response format from OpenAI API"
                
            full_response = response_json["choices"][0]["message"]["content"]
            logger.debug(f"Response length: {len(full_response)} characters")
        
            sanitized_response = self._sanitize_response(full_response)
            if len(sanitized_response) != len(full_response):
                logger.info(f"Response sanitized: original {len(full_response)} chars → sanitized {len(sanitized_response)} chars")
        
            return sanitized_response
        except requests.exceptions.Timeout:
            logger.error(f"OpenAI API request timed out after {self.timeout} seconds")
            return f"Error: Request timed out after {self.timeout} seconds"
            
        except requests.exceptions.HTTPError as e:
            logger.error(f"HTTP error from OpenAI API: {e}")
            error_message = self._extract_openai_error(e)
            return f"Error: {error_message}"
            
        except requests.exceptions.RequestException as e:
            logger.error(f"Error querying OpenAI API: {e}")
            return f"Error: {str(e)}"
            
        except json.JSONDecodeError:
            logger.error("Failed to parse JSON response from OpenAI API")
            return "Error: Invalid JSON response from API"
            
        except Exception as e:
            logger.error(f"Unexpected error in OpenAI API query: {e}")
            import traceback
            logger.error(traceback.format_exc())
            return f"Error: {str(e)}"
    
    def _extract_openai_error(self, e: requests.exceptions.HTTPError) -> str:
        """Extract detailed error information from OpenAI API errors."""
        error_message = "Unknown error"
        
        if hasattr(e, 'response') and e.response:
            try:
                error_json = e.response.json()
                if 'error' in error_json:
                    error_message = error_json['error'].get('message', str(e))
                    error_type = error_json['error'].get('type', 'unknown')
                    logger.error(f"OpenAI API error type: {error_type}, message: {error_message}")
            except:
                error_message = f"HTTP error {e.response.status_code}"
        
        return error_message

    def check_connection(self) -> bool:
        """Check if the connection to the OpenAI API is working."""
        logger.info(f"Testing connection to OpenAI API with model {self.model}...")
        try:
            response = self.query("Hello, this is a connection test.", max_tokens=10)
            if not response.startswith("Error:"):
                logger.info("Connection to OpenAI API successful!")
                return True
            else:
                logger.error(f"Connection test failed with response: {response}")
                return False
        except Exception as e:
            logger.error(f"Connection check failed: {e}")
            return False


class CustomAPIClient(BaseClient):
    """Client for custom LLM API endpoints."""
    
    def __init__(self, api_url: str, 
                 api_key: Optional[str] = None, 
                 args: Optional[argparse.Namespace] = None, 
                 config: Optional[Dict[str, Any]] = None) -> None:
        super().__init__(args, config)
        
        self.api_url = api_url
        self.api_key = api_key
        
        custom_api_config = (config or {}).get('api_providers', {}).get('custom_api', {})
        
        self.headers = custom_api_config.get('headers', {})
        self.session.headers.update(self.headers)
        
        if self.api_key:
            self.session.headers.update({"Authorization": f"Bearer {self.api_key}"})
        
        self.session.headers.update({"Content-Type": "application/json"})
        
        self.default_params: Dict[str, Any] = custom_api_config.get('default_params', {
            'temperature': 0.5,
            'max_tokens': 1200
        })

        logger.info(f"CustomAPIClient: final timeout = {self.timeout}")
        logger.info(f"CustomAPIClient: rate limiting enabled = {self.rate_limit_enabled}")
        if self.rate_limit_enabled:
            logger.info(f"CustomAPIClient: max requests per minute = {self.max_requests_per_minute}")
            logger.info(f"CustomAPIClient: cooldown period = {self.cooldown_period}")

    def query(self, prompt: str, **kwargs) -> str:
        """Send a query to the custom LLM API and return the response."""
        self._check_rate_limit()

        query_params = {**self.default_params, **kwargs}

        logger.debug(f"Query to Custom API:")
        logger.debug(f"- Endpoint: {self.api_url}")
        logger.debug(f"- Prompt length: {len(prompt)} characters")

        safe_params = {k: v for k, v in query_params.items() if 'key' not in k.lower() and 'token' not in k.lower()}
        logger.debug(f"- Parameters: {safe_params}")

        payload = self._prepare_payload(prompt, query_params)

        try:
            start_time = time.time()
            logger.debug(f"Sending request to {self.api_url}")
    
            response = self.session.post(self.api_url, json=payload, timeout=self.timeout)
    
            elapsed_time = time.time() - start_time
            logger.debug(f"Response received in {elapsed_time:.2f} seconds")
    
            response.raise_for_status()

            response_data = response.json()
            logger.debug(f"Response format: {list(response_data.keys())}")

            content = self._extract_response_content(response_data)
    
            if content:
                sanitized_response = self._sanitize_response(content)
                if len(sanitized_response) != len(content):
                    logger.info(f"Response sanitized: original {len(content)} chars → sanitized {len(sanitized_response)} chars")
        
                return sanitized_response
            else:
                logger.warning("Empty content received from API")
                return "Error: Empty response from API"
        
        except requests.exceptions.Timeout:
            logger.error(f"Request timed out after {self.timeout} seconds")
            return f"HTTP_ERROR_TIMEOUT: Request timed out after {self.timeout} seconds"
    
        except requests.exceptions.HTTPError as e:
            logger.error(f"HTTP error from custom API: {e}")
            status_code = e.response.status_code if hasattr(e, 'response') else 'unknown'
            error_text = e.response.text if hasattr(e, 'response') else str(e)
            return f"HTTP_ERROR_{status_code}: {error_text}"
    
        except requests.exceptions.RequestException as e:
            logger.error(f"Error querying custom API: {e}")
            error_detail = self._extract_request_error_details(e)
            if error_detail:
                logger.error(f"API error response: {error_detail}")
    
            return f"HTTP_ERROR_REQUEST: {str(e)}"
    
        except json.JSONDecodeError:
            logger.error("Failed to parse JSON response from API")
            return "HTTP_ERROR_PARSE: Invalid JSON response from API"
    
        except Exception as e:
            logger.error(f"Unexpected error in API query: {e}")
            import traceback
            logger.error(traceback.format_exc())
            return f"HTTP_ERROR_UNKNOWN: {str(e)}"
    
    def _prepare_payload(self, prompt: str, query_params: Dict[str, Any]) -> Dict[str, Any]:
        """Prepare the request payload based on endpoint type."""
        if self.api_url.endswith('/chat/completions'):
            payload = {
                "messages": [{"role": "user", "content": prompt}],
            }
            logger.debug("- Using OpenAI-compatible chat completions format")
            for key, value in query_params.items():
                if key != "messages":
                    payload[key] = value
        else:
            payload = {
                "prompt": prompt,
            }
            logger.debug("- Using legacy completion format")
            for key, value in query_params.items():
                if key != "prompt":
                    payload[key] = value
        
        return payload
    
    def _extract_response_content(self, response_data: Dict[str, Any]) -> Optional[str]:
        """Extract content from different response formats."""
        content = None
        
        if 'choices' in response_data and response_data['choices']:
            logger.debug("Processing response in OpenAI format")
            
            if isinstance(response_data['choices'][0], dict):
                if 'message' in response_data['choices'][0]:
                    content = response_data['choices'][0].get('message', {}).get('content', '')
                    logger.debug("Detected chat completion format")
                else:
                    content = response_data['choices'][0].get('text', '')
                    logger.debug("Detected regular completion format")
            else:
                content = str(response_data['choices'][0])
                logger.debug("Detected choices as non-dict format")
            
        elif 'text' in response_data:
            content = response_data['text']
            logger.debug("Detected simple text format")
        
        elif 'response' in response_data:
            content = response_data['response']
            logger.debug("Detected custom response format")
            
        else:
            content = str(response_data)
            logger.warning("Unknown response format. Using fallback format handling")
        
        return content
    
    def _extract_request_error_details(self, e: requests.exceptions.RequestException) -> str:
        """Extract error details from request exceptions."""
        if hasattr(e, 'response') and e.response:
            try:
                return e.response.json() if e.response.content else "No error details"
            except:
                return f"API error status: {e.response.status_code}"
        return "Unknown error"

    def check_connection(self) -> bool:
        """Check if the connection to the custom API is working."""
        test_prompts = [
            "Hello, this is a connection test.",
            "Verify system connection",
            "Test API functionality"
        ]
        
        logger.info("Testing connection to custom API endpoint...")
        
        for prompt in test_prompts:
            try:
                logger.info(f"Sending test prompt: '{prompt}'")
                response = self.query(prompt, max_tokens=10)
                if response and not response.startswith("Error:"):
                    logger.info("Connection test successful!")
                    return True
                else:
                    logger.warning(f"Test prompt failed with response: {response}")
            except Exception as e:
                logger.warning(f"Connection test failed for prompt '{prompt}': {e}")
        
        logger.error("All connection tests failed. Please check your API endpoint and credentials.")
        return False


def create_llm_client(args: argparse.Namespace, config: Dict[str, Any]) -> LLMClient:
    """Create an LLM client based on command-line arguments and configuration."""
    api_key = args.api_key or config.get("api_key") or os.environ.get("LLM_API_KEY")
    target_url = args.target_url or config.get("target_url")
    
    if not (api_key or target_url):
        logger.error("You must provide either an API key or a target URL")
        raise ValueError("You must provide either an API key or a target URL")
        
    if not args.model:
        args.model = config.get('default_model', 'gpt-4')
        
    config['default_model'] = args.model
    
    if target_url:
        logger.info(f"Using custom API endpoint: {target_url}")
        return CustomAPIClient(target_url, api_key, args, config)
    else:
        logger.info(f"Using OpenAI API with model: {args.model}")
        return OpenAIClient(api_key, args, config)