agent-api-server 2.1.7__py3-none-any.whl

agent_api_server/shared/get_model_info.py

@@ -0,0 +1,491 @@
+import logging
+import json
+from typing import Dict, Any, Optional, Union
+from dataclasses import dataclass
+from llm_sdk.model_providers.base import ConfigType
+from langchain_core.runnables.config import RunnableConfig
+from langchain_deepseek.chat_models import DEFAULT_API_BASE
+from agent_api_server.configs import global_config
+from agent_api_server.dynamic_llm.dynamic_llm import DynamicLLM
+from agent_api_server.shared.ase import AESCipher
+
+logger = logging.getLogger(__name__)
+
+# Mapping of providers to their credential keys
+PROVIDER_CREDENTIAL_MAP = {
+    'azure_openai': 'openai_api_key',
+    'xinference': 'api_key',
+    'tongyi': 'dashscope_api_key',
+    'openai_api_compatible': 'api_key',
+    'deepseek': 'api_key',
+    # ollama does not require API key
+}
+
+# Default values for special providers
+SPECIAL_PROVIDER_DEFAULTS = {
+    'xinference': 'sk-xxxxxx',
+    'ollama': '',
+}
+
+# Provider-specific base URLs
+PROVIDER_BASE_URLS = {
+    'deepseek': DEFAULT_API_BASE,
+    'tongyi': 'https://dashscope.aliyuncs.com/compatible-mode/v1',
+}
+
+
+@dataclass
+class LLMConfig:
+    """Configuration for LLM provider settings."""
+    provider: Optional[str] = None
+    model: Optional[str] = None
+    credentials: Optional[Union[str, Dict[str, Any]]] = None
+    agent_id: Optional[str] = None
+    ts_tenant: Optional[str] = None
+
+
+@dataclass
+class CredentialsInfo:
+    """Structured credential information."""
+    api_key: str = ''
+    base_url: str = ''
+
+
+class ConfigKeyBuilder:
+    """Builds configuration keys based on tool and tenant."""
+
+    @staticmethod
+    def get_key(tool_name: str, field: str, ts_tenant: Optional[str] = None) -> str:
+        """
+        Generate configuration key in format: TOOL_CHAT_FIELD[_TENANT]
+
+        Args:
+            tool_name: Name of the tool
+            field: Configuration field (e.g., PROVIDER, MODEL)
+            ts_tenant: Optional tenant identifier
+
+        Returns:
+            Formatted configuration key
+        """
+        parts = []
+
+        if tool_name.lower() != "default":
+            parts.append(tool_name.upper())
+
+        parts.extend([ConfigType.CHAT.value, field])
+
+        if ts_tenant:
+            parts.append(ts_tenant)
+
+        return "_".join(parts)
+
+
+class CredentialsParser:
+    """Parses credential data from various formats."""
+
+    @staticmethod
+    def parse_credentials(credentials: Union[str, Dict[str, Any]],
+                          tool_name: str) -> Optional[Dict[str, Any]]:
+        """
+        Parse credentials from string or dict format.
+
+        Args:
+            credentials: Credentials data (string JSON or dict)
+            tool_name: Name of the tool for logging
+
+        Returns:
+            Parsed credentials as dict, or None on error
+        """
+        if isinstance(credentials, dict):
+            return credentials
+
+        if not isinstance(credentials, str):
+            logger.error(f"Unexpected credentials type for tool {tool_name}: {type(credentials)}")
+            return None
+
+        credentials_str = credentials.strip()
+        if not credentials_str:
+            return {}
+
+        if credentials_str.startswith(('{', '[')):
+            try:
+                return json.loads(credentials_str)
+            except json.JSONDecodeError as e:
+                logger.error(f"Failed to parse JSON credentials for tool {tool_name}: {e}")
+                logger.debug(f"Raw credentials (first 100 chars): {credentials_str[:100]}")
+                return None
+
+        return {"raw_key": credentials_str}
+
+
+class SecurityManager:
+    """Manages cryptographic operations for API keys."""
+
+    def __init__(self):
+        self._aes_cache: Dict[str, AESCipher] = {}
+
+    def _get_aes_client(self, app_id: str) -> AESCipher:
+        """Get or create AES cipher client for the given app_id."""
+        if app_id not in self._aes_cache:
+            self._aes_cache[app_id] = AESCipher(app_id)
+        return self._aes_cache[app_id]
+
+    def decrypt_api_key(self, encrypted_key: str, app_id: str) -> Optional[str]:
+        """
+        Decrypt API key using AES cipher.
+
+        Args:
+            encrypted_key: Encrypted API key
+            app_id: Application ID for decryption
+
+        Returns:
+            Decrypted key, or None if decryption fails
+        """
+        if not encrypted_key or not app_id:
+            return None
+
+        try:
+            aes_client = self._get_aes_client(app_id)
+            decrypted_key = aes_client.decrypt(encrypted_key)
+            logger.debug(f"Successfully decrypted API key (length: {len(decrypted_key)})")
+            return decrypted_key
+        except Exception as e:
+            logger.error(f"Failed to decrypt API key: {str(e)}")
+            return None
+
+
+class LLMConfigManager:
+    """Main manager for LLM configuration and credential retrieval."""
+
+    def __init__(self):
+        self.key_builder = ConfigKeyBuilder()
+        self.credentials_parser = CredentialsParser()
+        self.security_manager = SecurityManager()
+
+    def get_llm_config(self, tool_name: str, context: Dict[str, Any]) -> LLMConfig:
+        """
+        Extract LLM configuration from context.
+
+        Args:
+            tool_name: Name of the tool
+            context: Configuration context dictionary
+
+        Returns:
+            LLMConfig object with extracted values
+        """
+        configurable = context.get("configurable", {})
+        ts_tenant = configurable.get("ts_tenant") or configurable.get("TSTenant")
+        graph_name = configurable.get("graph_name")
+
+        return LLMConfig(
+            provider=configurable.get(self.key_builder.get_key(tool_name, "PROVIDER", ts_tenant)),
+            model=configurable.get(self.key_builder.get_key(tool_name, "MODEL", ts_tenant)),
+            credentials=configurable.get(self.key_builder.get_key(tool_name, "CREDENTIALS", ts_tenant)),
+            agent_id=graph_name,
+            ts_tenant=ts_tenant
+        )
+
+    def get_api_version(self, tool_name: str, config: Optional[RunnableConfig]) -> str:
+        if not config:
+            logger.debug(f"No config provided for tool: {tool_name}")
+            return ''
+
+        llm_config = self.get_llm_config(tool_name, context=config or {})
+        credentials_dict = self.credentials_parser.parse_credentials(
+            llm_config.credentials, tool_name
+        )
+        if not credentials_dict:
+            return ''
+
+        return credentials_dict.get('openai_api_version')
+
+    def get_base_model_name(self, tool_name: str, config: Optional[RunnableConfig]) -> str:
+        if not config:
+            logger.debug(f"No config provided for tool: {tool_name}")
+            return ''
+
+        llm_config = self.get_llm_config(tool_name, context=config or {})
+        credentials_dict = self.credentials_parser.parse_credentials(
+            llm_config.credentials, tool_name
+        )
+        if not credentials_dict:
+            return ''
+
+        return credentials_dict.get('base_model_name')
+
+    def get_provider(self, tool_name: str, config: Optional[RunnableConfig]) -> str:
+        """Get provider name from configuration."""
+        llm_config = self.get_llm_config(tool_name, context=config or {})
+        return llm_config.provider or ''
+
+    def get_model_name(self, tool_name: str, config: Optional[RunnableConfig]) -> str:
+        """Get model name from configuration."""
+        llm_config = self.get_llm_config(tool_name, context=config or {})
+        return llm_config.model or ''
+
+    def _process_xinference_base_url(self, credentials_dict: Dict[str, Any], app_id: str) -> str:
+        """
+        Process Xinference base URL, handling decryption and formatting.
+
+        Args:
+            credentials_dict: Parsed credentials dictionary
+            app_id: Application ID for decryption
+
+        Returns:
+            Formatted base URL, or empty string on error
+        """
+        server_url = credentials_dict.get("server_url")
+        if not server_url:
+            logger.error("server_url is required for xinference provider")
+            return ''
+
+        decrypted_url = self.security_manager.decrypt_api_key(server_url, app_id)
+        if not decrypted_url:
+            logger.warning("Failed to decrypt xinference server_url")
+            return ''
+
+        base_url = decrypted_url.rstrip("/")
+        if not base_url.endswith("/v1"):
+            base_url = base_url + "/v1"
+
+        return base_url
+
+    @staticmethod
+    def _get_base_url_from_credentials(credentials_dict: Dict[str, Any], provider: str) -> str:
+        """
+        Extract base URL from credentials for specific providers.
+
+        Args:
+            credentials_dict: Parsed credentials dictionary
+            provider: Provider name
+
+        Returns:
+            Base URL if found, empty string otherwise
+        """
+        if provider == 'azure_openai':
+            return credentials_dict.get("openai_api_base", "")
+        elif provider == 'openai_api_compatible':
+            return credentials_dict.get("endpoint_url", "")
+        elif provider == 'ollama':
+            return credentials_dict.get("base_url", "")
+        return ''
+
+    def get_base_url(self, tool_name: str, config: Optional[RunnableConfig]) -> str:
+        """
+        Get base URL for the LLM provider.
+
+        Args:
+            tool_name: Name of the tool
+            config: Runnable configuration
+
+        Returns:
+            Base URL string, or empty string if not found
+        """
+        if not config:
+            logger.debug(f"No config provided for tool: {tool_name}")
+            return ''
+
+        llm_config = self.get_llm_config(tool_name, context=config)
+        if not llm_config.provider:
+            logger.debug(f"No provider found for tool: {tool_name}")
+            return ''
+
+        # Check provider-specific defaults
+        if llm_config.provider in PROVIDER_BASE_URLS:
+            return PROVIDER_BASE_URLS[llm_config.provider]
+
+        if not llm_config.credentials:
+            logger.debug(f"No credentials found for tool: {tool_name}")
+            return ''
+
+        credentials_dict = self.credentials_parser.parse_credentials(
+            llm_config.credentials, tool_name
+        )
+        if not credentials_dict:
+            return ''
+
+        # Special handling for xinference
+        if llm_config.provider == 'xinference':
+            if not llm_config.agent_id:
+                logger.warning("No agent_id found in config for xinference")
+                return ''
+            return self._process_xinference_base_url(credentials_dict, llm_config.agent_id)
+
+        return self._get_base_url_from_credentials(credentials_dict, llm_config.provider)
+
+    @staticmethod
+    def _get_raw_api_key_from_credentials(credentials: Union[str, Dict[str, Any]],
+                                          credential_key: str) -> Optional[str]:
+        """
+        Extract raw API key from credentials.
+
+        Args:
+            credentials: Credentials data
+            credential_key: Key to look for in credentials
+
+        Returns:
+            Raw API key, or None if not found
+        """
+        if isinstance(credentials, dict):
+            return credentials.get(credential_key)
+
+        if not isinstance(credentials, str):
+            return None
+
+        credentials_str = credentials.strip()
+        if not credentials_str:
+            return None
+
+        if credentials_str.startswith(('{', '[')):
+            try:
+                cred_dict = json.loads(credentials_str)
+                if isinstance(cred_dict, dict):
+                    return cred_dict.get(credential_key)
+                logger.warning(f"Parsed credentials is not a dict: {type(cred_dict)}")
+                return None
+            except json.JSONDecodeError:
+                return None
+        else:
+            return credentials_str
+
+    def get_api_key(self, tool_name: str, config: Optional[RunnableConfig]) -> str:
+        """
+        Get API key for the LLM provider, handling decryption.
+
+        Args:
+            tool_name: Name of the tool
+            config: Runnable configuration
+
+        Returns:
+            API key string (decrypted if encrypted, raw otherwise)
+        """
+        if not config:
+            logger.debug(f"No config provided for tool: {tool_name}")
+            return ''
+
+        llm_config = self.get_llm_config(tool_name, context=config)
+
+        # Check for required fields
+        if not llm_config.provider:
+            logger.debug(f"No provider found for tool: {tool_name}")
+            return ''
+
+        provider = llm_config.provider
+
+        # Return defaults for special providers if no credentials
+        if provider in SPECIAL_PROVIDER_DEFAULTS and not llm_config.credentials:
+            return SPECIAL_PROVIDER_DEFAULTS[provider]
+
+        # Get credential key for the provider
+        credential_key = PROVIDER_CREDENTIAL_MAP.get(provider)
+        if not credential_key:
+            logger.debug(f"No credential mapping for provider: {provider}")
+            return ''
+
+        # Extract raw key from credentials
+        raw_key = self._get_raw_api_key_from_credentials(
+            llm_config.credentials, credential_key
+        )
+
+        # Handle missing keys
+        if not raw_key:
+            if provider in SPECIAL_PROVIDER_DEFAULTS:
+                return SPECIAL_PROVIDER_DEFAULTS.get(provider, '')
+            return ''
+
+        # Return default keys directly
+        if (provider in SPECIAL_PROVIDER_DEFAULTS and
+                raw_key == SPECIAL_PROVIDER_DEFAULTS[provider]):
+            return raw_key
+
+        # Attempt decryption for encrypted keys
+        if llm_config.agent_id:
+            decrypted_key = self.security_manager.decrypt_api_key(raw_key, llm_config.agent_id)
+            if decrypted_key is not None:
+                return decrypted_key
+            else:
+                # Return raw key if decryption fails
+                logger.warning(f"Decryption failed for {provider}, using raw key")
+                return raw_key or ''
+        else:
+            # No agent_id for decryption, return raw key
+            logger.debug(f"No agent_id for decryption, using raw key for {provider}")
+            return raw_key
+
+
+# Global instance
+config_manager = LLMConfigManager()
+
+def get_chat_model_information(tool_name: str, config: Optional[RunnableConfig]) -> Dict:
+    configurable = config.get("configurable")
+    graph_name = configurable.get('graph_name')
+    use_sys_llm = configurable.get('use_sys_llm')
+
+    if use_sys_llm == 'true':
+        from llm_sdk.llm import ModelInstanceFactory
+        _model = ModelInstanceFactory.get_model_instance(
+            model_managment_url=global_config.MODEL_MANAGER_SERVICE_URL,
+            config_type=ConfigType.CHAT,
+            token=global_config.CLIENT_TOKEN,
+            app_id=graph_name,
+            app_name=graph_name,
+        )
+
+        return _model.unified_credential()
+    else:
+        chat_llm = DynamicLLM(tool_name=tool_name, config_type=ConfigType.CHAT)
+        llm_config = chat_llm._get_llm_config_from_context(config or {})
+
+        if chat_llm._is_credentials_empty(config=llm_config):
+            _ = chat_llm._get_model_instance(llm_config, config)
+
+        model_provider = get_chat_llm_model_provider(tool_name=tool_name, config=config)
+        api_key = get_chat_llm_model_api_key(tool_name=tool_name, config=config)
+        base_url = get_chat_llm_model_base_url(tool_name=tool_name, config=config)
+        api_version = get_chat_llm_model_api_version(tool_name=tool_name, config=config)
+
+        if model_provider == 'azure_openai':
+            model_name = get_chat_llm_model_base_model_name(tool_name=tool_name, config=config)
+        else:
+            model_name = get_chat_llm_model_name(tool_name=tool_name, config=config)
+
+        return {
+            "openai_api_key": api_key,
+            "openai_api_type": model_provider,
+            "openai_api_deployment": model_name,
+            "openai_api_version": api_version,
+            "openai_api_base": base_url
+        }
+
+def get_chat_llm_model_provider(tool_name: str, config: Optional[RunnableConfig]) -> str:
+    """Public interface to get LLM provider."""
+    return config_manager.get_provider(tool_name, config)
+
+
+def get_chat_llm_model_name(tool_name: str, config: Optional[RunnableConfig]) -> str:
+    """Public interface to get LLM model name."""
+    return config_manager.get_model_name(tool_name, config)
+
+
+def get_chat_llm_model_base_url(tool_name: str, config: Optional[RunnableConfig]) -> str:
+    """Public interface to get LLM base URL."""
+    return config_manager.get_base_url(tool_name, config)
+
+
+def get_chat_llm_model_api_key(tool_name: str, config: Optional[RunnableConfig]) -> str:
+    """Public interface to get LLM API key."""
+    return config_manager.get_api_key(tool_name, config)
+
+def get_chat_llm_model_api_version(tool_name: str, config: Optional[RunnableConfig]) -> str:
+    provider = config_manager.get_provider(tool_name, config)
+    if provider == 'azure_openai':
+        return config_manager.get_api_version(tool_name, config)
+    else:
+        return ''
+
+def get_chat_llm_model_base_model_name(tool_name: str, config: Optional[RunnableConfig]) -> str:
+    provider = config_manager.get_provider(tool_name, config)
+    if provider == 'azure_openai':
+        return config_manager.get_base_model_name(tool_name, config)
+    else:
+        return ''