isa-model 0.0.1__py3-none-any.whl

isa_model/inference/providers/base_provider.py

@@ -0,0 +1,30 @@
+from abc import ABC, abstractmethod
+from typing import Dict, List, Any, Optional
+
+from isa_model.inference.base import ModelType, Capability
+
+class BaseProvider(ABC):
+    """Base class for all AI providers"""
+    
+    def __init__(self, config: Optional[Dict[str, Any]] = None):
+        self.config = config or {}
+    
+    @abstractmethod
+    def get_capabilities(self) -> Dict[ModelType, List[Capability]]:
+        """Get provider capabilities by model type"""
+        pass
+    
+    @abstractmethod
+    def get_models(self, model_type: ModelType) -> List[str]:
+        """Get available models for given type"""
+        pass
+    
+    @abstractmethod
+    def get_config(self) -> Dict[str, Any]:
+        """Get provider configuration"""
+        return self.config
+    
+    @abstractmethod
+    def is_reasoning_model(self, model_name: str) -> bool:
+        """Check if the model is optimized for reasoning tasks"""
+        pass

isa_model/inference/providers/model_cache_manager.py

@@ -0,0 +1,341 @@
+from typing import Dict, List, Any, Optional
+import aiohttp
+import logging
+import asyncio
+from collections import OrderedDict
+import os
+import json
+import hashlib
+from pathlib import Path
+from isa_model.inference.base import ModelType
+
+logger = logging.getLogger(__name__)
+
+class ModelCacheManager:
+    """管理Triton服务器模型的加载/卸载，支持轮询模式"""
+    
+    def __init__(self, cache_size: int = 5, model_repository: str = "/models"):
+        """
+        初始化模型缓存管理器
+        
+        Args:
+            cache_size: 最大缓存模型数量
+            model_repository: 模型库路径
+        """
+        self.cache_size = cache_size
+        self.model_repository = model_repository
+        
+        # LRU缓存使用OrderedDict
+        self.model_cache = OrderedDict()
+        
+        # 服务器配置
+        self.server_config = {
+            "polling_enabled": True,  # 默认启用轮询模式（适合多模型场景）
+            "triton_url": "localhost:8000",
+            "openai_api_url": "localhost:9000"
+        }
+        
+        # 模型类型映射
+        self.model_type_map = {
+            ModelType.LLM: "llm",
+            ModelType.EMBEDDING: "embedding",
+            ModelType.VISION: "vision",
+            ModelType.RERANK: "rerank"
+        }
+        
+        logger.info(f"初始化ModelCacheManager，缓存大小: {cache_size}，模型库: {model_repository}")
+    
+    async def detect_server_mode(self):
+        """检测Triton服务器是否运行在轮询模式"""
+        try:
+            # 尝试加载任意模型以检测模式
+            models = await self._get_repository_models()
+            if not models:
+                logger.warning("无法获取模型列表，无法检测服务器模式")
+                return
+                
+            test_model = models[0]
+            url = f"http://{self.server_config['triton_url']}/v2/repository/models/{test_model}/load"
+            
+            async with aiohttp.ClientSession() as session:
+                async with session.post(url) as response:
+                    response_text = await response.text()
+                    
+                    if response.status == 503 and "polling is enabled" in response_text:
+                        self.server_config["polling_enabled"] = True
+                        logger.info("检测到Triton服务器运行在轮询模式（多模型模式）")
+                    elif response.status == 200:
+                        self.server_config["polling_enabled"] = False
+                        logger.info("检测到Triton服务器运行在手动加载模式（单模型模式）")
+                    else:
+                        logger.warning(f"无法确定服务器模式，状态码: {response.status}")
+        except Exception as e:
+            logger.error(f"检测服务器模式时出错: {e}")
+    
+    async def load_model(self, model_name: str, model_type: ModelType) -> bool:
+        """
+        加载模型到Triton服务器
+        
+        Args:
+            model_name: 模型名称
+            model_type: 模型类型
+            
+        Returns:
+            bool: 成功返回True，失败返回False
+        """
+        # 如果是第一次加载，检测服务器模式
+        if not hasattr(self, '_mode_detected'):
+            await self.detect_server_mode()
+            self._mode_detected = True
+        
+        if model_name in self.model_cache:
+            # 模型已加载，移到LRU缓存末尾
+            self.model_cache.move_to_end(model_name)
+            logger.info(f"模型 {model_name} 已在缓存中，移至末尾")
+            return True
+        
+        try:
+            # 检查模型是否已加载到服务器
+            is_loaded = await self._check_model_loaded(model_name)
+            if is_loaded:
+                logger.info(f"模型 {model_name} 已在服务器中加载")
+                self.model_cache[model_name] = {
+                    "type": model_type,
+                    "load_time": asyncio.get_event_loop().time()
+                }
+                return True
+            
+            # 如果在轮询模式下，我们不能手动加载模型
+            if self.server_config["polling_enabled"]:
+                # 检查模型是否存在
+                exists = await self._check_model_exists(model_name)
+                if exists:
+                    logger.warning(f"服务器在轮询模式下，无法手动加载模型 {model_name}，但模型存在")
+                    # 我们假设模型将通过轮询加载
+                    return True
+                else:
+                    logger.error(f"模型 {model_name} 不存在于服务器存储库中")
+                    return False
+            else:
+                # 在非轮询模式下，可以手动加载
+                # 如果缓存已满，卸载最少使用的模型
+                if len(self.model_cache) >= self.cache_size:
+                    lru_model, _ = self.model_cache.popitem(last=False)
+                    await self._unload_from_triton(lru_model)
+                    logger.info(f"从缓存中卸载LRU模型 {lru_model}")
+                
+                # 加载新模型
+                success = await self._load_to_triton(model_name)
+                if success:
+                    self.model_cache[model_name] = {
+                        "type": model_type,
+                        "load_time": asyncio.get_event_loop().time()
+                    }
+                    logger.info(f"成功加载模型 {model_name}")
+                    return True
+                else:
+                    logger.error(f"加载模型 {model_name} 失败")
+                    return False
+                
+        except Exception as e:
+            logger.error(f"加载模型 {model_name} 时出错: {e}")
+            return False
+    
+    async def unload_model(self, model_name: str) -> bool:
+        """卸载模型"""
+        # 如果在轮询模式下，我们不能手动卸载模型
+        if self.server_config["polling_enabled"]:
+            logger.warning(f"服务器在轮询模式下，无法手动卸载模型 {model_name}")
+            return True
+        
+        if model_name not in self.model_cache:
+            logger.warning(f"模型 {model_name} 未在缓存中，无需卸载")
+            return True
+        
+        try:
+            # 卸载模型
+            success = await self._unload_from_triton(model_name)
+            if success:
+                # 从缓存中移除
+                self.model_cache.pop(model_name, None)
+                logger.info(f"成功卸载模型 {model_name}")
+                return True
+            else:
+                logger.error(f"卸载模型 {model_name} 失败")
+                return False
+                
+        except Exception as e:
+            logger.error(f"卸载模型 {model_name} 时出错: {e}")
+            return False
+    
+    async def _load_to_triton(self, model_name: str) -> bool:
+        """向Triton服务器发送加载模型请求"""
+        try:
+            logger.info(f"尝试加载模型 {model_name} 到Triton服务器")
+            
+            url = f"http://{self.server_config['triton_url']}/v2/repository/models/{model_name}/load"
+            
+            async with aiohttp.ClientSession() as session:
+                async with session.post(url) as response:
+                    response_text = await response.text()
+                    
+                    if response.status == 200:
+                        logger.info(f"成功加载模型 {model_name}")
+                        return True
+                    elif response.status == 400:
+                        # 模型可能已加载
+                        logger.info(f"模型 {model_name} 可能已加载: {response_text}")
+                        return True
+                    elif response.status == 503 and "polling is enabled" in response_text:
+                        # 检测到轮询模式
+                        self.server_config["polling_enabled"] = True
+                        logger.warning(f"服务器在轮询模式下，无法手动加载模型: {response_text}")
+                        # 检查模型是否存在
+                        return await self._check_model_exists(model_name)
+                    else:
+                        logger.error(f"加载模型 {model_name} 失败: Status {response.status}, Response: {response_text}")
+                        return False
+                    
+        except Exception as e:
+            logger.error(f"向Triton API发送加载模型 {model_name} 请求时出错: {e}")
+            return False
+    
+    async def _check_model_loaded(self, model_name: str) -> bool:
+        """检查模型是否已加载"""
+        try:
+            url = f"http://{self.server_config['triton_url']}/v2/models/{model_name}/ready"
+            
+            async with aiohttp.ClientSession() as session:
+                async with session.get(url) as response:
+                    if response.status == 200:
+                        logger.info(f"模型 {model_name} 已加载")
+                        return True
+                    else:
+                        logger.info(f"模型 {model_name} 未加载，状态码: {response.status}")
+                        return False
+        except Exception as e:
+            logger.error(f"检查模型 {model_name} 是否加载时出错: {e}")
+            return False
+    
+    async def _check_model_exists(self, model_name: str) -> bool:
+        """检查模型是否存在于存储库中"""
+        try:
+            url = f"http://{self.server_config['triton_url']}/v2/repository/index"
+            
+            async with aiohttp.ClientSession() as session:
+                async with session.post(url) as response:
+                    if response.status == 200:
+                        models = await response.json()
+                        model_names = [model["name"] for model in models]
+                        exists = model_name in model_names
+                        logger.info(f"模型 {model_name} {'存在' if exists else '不存在'}于存储库中")
+                        logger.info(f"可用模型: {model_names}")
+                        return exists
+                    else:
+                        logger.error(f"检查模型存在性失败: {response.status}")
+                        return False
+        except Exception as e:
+            logger.error(f"检查模型存在性时出错: {e}")
+            return False
+    
+    async def _unload_from_triton(self, model_name: str) -> bool:
+        """从Triton服务器卸载模型"""
+        try:
+            url = f"http://{self.server_config['triton_url']}/v2/repository/models/{model_name}/unload"
+            
+            async with aiohttp.ClientSession() as session:
+                async with session.post(url) as response:
+                    response_text = await response.text()
+                    
+                    if response.status == 200:
+                        logger.info(f"成功卸载模型 {model_name}")
+                        return True
+                    elif response.status == 503 and "polling is enabled" in response_text:
+                        # 检测到轮询模式
+                        self.server_config["polling_enabled"] = True
+                        logger.warning(f"服务器在轮询模式下，无法手动卸载模型: {response_text}")
+                        return True
+                    else:
+                        logger.error(f"卸载模型 {model_name} 失败: Status {response.status}, Response: {response_text}")
+                        return False
+        except Exception as e:
+            logger.error(f"向Triton API发送卸载模型 {model_name} 请求时出错: {e}")
+            return False
+    
+    def list_available_models(self, model_type: ModelType = None) -> List[str]:
+        """
+        列出可用模型
+        
+        Args:
+            model_type: 按模型类型筛选
+            
+        Returns:
+            模型名称列表
+        """
+        try:
+            # 获取模型列表
+            models = asyncio.run(self._get_repository_models())
+            
+            if not models:
+                logger.warning("在存储库中未找到模型或无法连接到服务器")
+                return []
+            
+            # 如果未指定模型类型，返回所有模型
+            if model_type is None:
+                return models
+            
+            # 基于命名约定的简单过滤器
+            if model_type == ModelType.LLM:
+                # 返回包含关键字的模型
+                llm_keywords = ["llama", "mistral", "gemma", "qwen", "phi", "gpt", "falcon"]
+                return [m for m in models if any(kw in m.lower() for kw in llm_keywords)]
+            elif model_type == ModelType.EMBEDDING:
+                embed_keywords = ["embed", "bge", "e5", "text-embedding"]
+                return [m for m in models if any(kw in m.lower() for kw in embed_keywords)]
+            elif model_type == ModelType.VISION:
+                vision_keywords = ["clip", "vision", "multimodal", "image"]
+                return [m for m in models if any(kw in m.lower() for kw in vision_keywords)]
+            elif model_type == ModelType.RERANK:
+                rerank_keywords = ["rerank", "cross-encoder"]
+                return [m for m in models if any(kw in m.lower() for kw in rerank_keywords)]
+            else:
+                return []
+                
+        except Exception as e:
+            logger.error(f"列出模型时出错: {e}")
+            return []
+    
+    async def _get_repository_models(self) -> List[str]:
+        """从Triton服务器获取模型列表"""
+        try:
+            url = f"http://{self.server_config['triton_url']}/v2/repository/index"
+            
+            async with aiohttp.ClientSession() as session:
+                async with session.post(url) as response:
+                    if response.status == 200:
+                        models = await response.json()
+                        return [model["name"] for model in models]
+                    else:
+                        logger.error(f"获取模型失败: Status {response.status}")
+                        return []
+        except Exception as e:
+            logger.error(f"获取存储库模型时出错: {e}")
+            return []
+    
+    async def get_openai_models(self) -> List[Dict[str, Any]]:
+        """获取OpenAI兼容API中的可用模型"""
+        try:
+            url = f"http://{self.server_config['openai_api_url']}/v1/models"
+            
+            async with aiohttp.ClientSession() as session:
+                async with session.get(url) as response:
+                    if response.status == 200:
+                        result = await response.json()
+                        logger.info(f"从OpenAI API获取到 {len(result.get('data', []))} 个模型")
+                        return result.get("data", [])
+                    else:
+                        logger.error(f"获取OpenAI模型失败: Status {response.status}")
+                        return []
+        except Exception as e:
+            logger.error(f"获取OpenAI模型时出错: {e}")
+            return [] 

isa_model/inference/providers/ollama_provider.py

@@ -0,0 +1,73 @@
+from isa_model.inference.providers.base_provider import BaseProvider
+from isa_model.inference.base import ModelType, Capability
+from typing import Dict, List, Any
+import logging
+
+logger = logging.getLogger(__name__)
+
+class OllamaProvider(BaseProvider):
+    """Provider for Ollama API"""
+    
+    def __init__(self, config=None):
+        """
+        Initialize the Ollama Provider
+        
+        Args:
+            config (dict, optional): Configuration for the provider
+                - base_url: Base URL for Ollama API (default: http://localhost:11434)
+                - timeout: Timeout for API calls in seconds
+        """
+        default_config = {
+            "base_url": "http://localhost:11434",
+            "timeout": 60,
+            "stream": True,
+            "temperature": 0.7,
+            "top_p": 0.9,
+            "max_tokens": 2048,
+            "keep_alive": "5m"
+        }
+        
+        # Merge default config with provided config
+        merged_config = {**default_config, **(config or {})}
+        
+        super().__init__(config=merged_config)
+        self.name = "ollama"
+        
+        logger.info(f"Initialized OllamaProvider with URL: {self.config['base_url']}")
+    
+    def get_capabilities(self) -> Dict[ModelType, List[Capability]]:
+        """Get provider capabilities by model type"""
+        return {
+            ModelType.LLM: [
+                Capability.CHAT, 
+                Capability.COMPLETION
+            ],
+            ModelType.EMBEDDING: [
+                Capability.EMBEDDING
+            ],
+            ModelType.VISION: [
+                Capability.IMAGE_UNDERSTANDING
+            ]
+        }
+    
+    def get_models(self, model_type: ModelType) -> List[str]:
+        """Get available models for given type"""
+        # Placeholder: In real implementation, this would query Ollama API
+        if model_type == ModelType.LLM:
+            return ["llama3", "mistral", "phi3", "llama3.1", "codellama", "gemma"]
+        elif model_type == ModelType.EMBEDDING:
+            return ["bge-m3", "nomic-embed-text"]
+        elif model_type == ModelType.VISION:
+            return ["llava", "bakllava", "llama3-vision"]
+        else:
+            return []
+    
+    def get_config(self) -> Dict[str, Any]:
+        """Get provider configuration"""
+        return self.config
+    
+    def is_reasoning_model(self, model_name: str) -> bool:
+        """Check if the model is optimized for reasoning tasks"""
+        # Default implementation: consider larger models as reasoning-capable
+        reasoning_models = ["llama3.1", "llama3", "claude3", "gpt4", "mixtral", "gemma", "palm2"]
+        return any(rm in model_name.lower() for rm in reasoning_models)

isa_model/inference/providers/openai_provider.py

@@ -0,0 +1,101 @@
+from isa_model.inference.providers.base_provider import BaseProvider
+from isa_model.inference.base import ModelType, Capability
+from typing import Dict, List, Any
+import logging
+import os
+
+logger = logging.getLogger(__name__)
+
+class OpenAIProvider(BaseProvider):
+    """Provider for OpenAI API"""
+    
+    def __init__(self, config=None):
+        """
+        Initialize the OpenAI Provider
+        
+        Args:
+            config (dict, optional): Configuration for the provider
+                - api_key: OpenAI API key (can be passed here or via environment variable)
+                - api_base: Base URL for OpenAI API (default: https://api.openai.com/v1)
+                - timeout: Timeout for API calls in seconds
+        """
+        default_config = {
+            "api_key": "",  # Will be set from config or environment
+            "api_base": os.environ.get("OPENAI_API_BASE", "https://api.openai.com/v1"),
+            "timeout": 60,
+            "stream": True,
+            "temperature": 0.7,
+            "top_p": 0.9,
+            "max_tokens": 1024
+        }
+        
+        # Merge default config with provided config
+        merged_config = {**default_config, **(config or {})}
+        
+        # Set API key from config first, then fallback to environment variable
+        if not merged_config["api_key"]:
+            merged_config["api_key"] = os.environ.get("OPENAI_API_KEY", "")
+        
+        super().__init__(config=merged_config)
+        self.name = "openai"
+        
+        logger.info(f"Initialized OpenAIProvider with URL: {self.config['api_base']}")
+        
+        # Only warn if no API key is provided at all
+        if not self.config["api_key"]:
+            logger.info("OpenAI API key not provided. You can set it via OPENAI_API_KEY environment variable or pass it in the config when creating services.")
+    
+    def set_api_key(self, api_key: str):
+        """
+        Set the API key after initialization
+        
+        Args:
+            api_key: OpenAI API key
+        """
+        self.config["api_key"] = api_key
+        logger.info("OpenAI API key updated")
+    
+    def get_capabilities(self) -> Dict[ModelType, List[Capability]]:
+        """Get provider capabilities by model type"""
+        return {
+            ModelType.LLM: [
+                Capability.CHAT, 
+                Capability.COMPLETION
+            ],
+            ModelType.EMBEDDING: [
+                Capability.EMBEDDING
+            ],
+            ModelType.VISION: [
+                Capability.IMAGE_GENERATION,
+                Capability.MULTIMODAL_UNDERSTANDING
+            ],
+            ModelType.AUDIO: [
+                Capability.SPEECH_TO_TEXT
+            ]
+        }
+    
+    def get_models(self, model_type: ModelType) -> List[str]:
+        """Get available models for given type"""
+        if model_type == ModelType.LLM:
+            return ["gpt-4o", "gpt-4o-mini", "gpt-4-turbo", "gpt-4", "gpt-3.5-turbo"]
+        elif model_type == ModelType.EMBEDDING:
+            return ["text-embedding-3-large", "text-embedding-3-small", "text-embedding-ada-002"]
+        elif model_type == ModelType.VISION:
+            return ["gpt-4o", "gpt-4-vision-preview"]
+        elif model_type == ModelType.AUDIO:
+            return ["whisper-1"]
+        else:
+            return []
+    
+    def get_config(self) -> Dict[str, Any]:
+        """Get provider configuration"""
+        # Return a copy without sensitive information
+        config_copy = self.config.copy()
+        if "api_key" in config_copy:
+            config_copy["api_key"] = "***" if config_copy["api_key"] else ""
+        return config_copy
+    
+    def is_reasoning_model(self, model_name: str) -> bool:
+        """Check if the model is optimized for reasoning tasks"""
+        reasoning_models = ["gpt-4", "gpt-4o", "gpt-4-turbo"]
+        return any(rm in model_name.lower() for rm in reasoning_models) 

isa_model/inference/providers/replicate_provider.py

@@ -0,0 +1,107 @@
+from isa_model.inference.providers.base_provider import BaseProvider
+from isa_model.inference.base import ModelType, Capability
+from typing import Dict, List, Any
+import logging
+import os
+
+logger = logging.getLogger(__name__)
+
+class ReplicateProvider(BaseProvider):
+    """Provider for Replicate API"""
+    
+    def __init__(self, config=None):
+        """
+        Initialize the Replicate Provider
+        
+        Args:
+            config (dict, optional): Configuration for the provider
+                - api_token: Replicate API token (can be passed here or via environment variable)
+                - timeout: Timeout for API calls in seconds
+        """
+        default_config = {
+            "api_token": "",  # Will be set from config or environment
+            "timeout": 60,
+            "stream": True,
+            "max_tokens": 1024
+        }
+        
+        # Merge default config with provided config
+        merged_config = {**default_config, **(config or {})}
+        
+        # Set API token from config first, then fallback to environment variable
+        if not merged_config["api_token"]:
+            merged_config["api_token"] = os.environ.get("REPLICATE_API_TOKEN", "")
+        
+        super().__init__(config=merged_config)
+        self.name = "replicate"
+        
+        logger.info(f"Initialized ReplicateProvider")
+        
+        # Only warn if no API token is provided at all
+        if not self.config["api_token"]:
+            logger.info("Replicate API token not provided. You can set it via REPLICATE_API_TOKEN environment variable or pass it in the config when creating services.")
+    
+    def set_api_token(self, api_token: str):
+        """
+        Set the API token after initialization
+        
+        Args:
+            api_token: Replicate API token
+        """
+        self.config["api_token"] = api_token
+        logger.info("Replicate API token updated")
+    
+    def get_capabilities(self) -> Dict[ModelType, List[Capability]]:
+        """Get provider capabilities by model type"""
+        return {
+            ModelType.LLM: [
+                Capability.CHAT, 
+                Capability.COMPLETION
+            ],
+            ModelType.VISION: [
+                Capability.IMAGE_GENERATION,
+                Capability.MULTIMODAL_UNDERSTANDING
+            ],
+            ModelType.AUDIO: [
+                Capability.SPEECH_TO_TEXT,
+                Capability.TEXT_TO_SPEECH
+            ]
+        }
+    
+    def get_models(self, model_type: ModelType) -> List[str]:
+        """Get available models for given type"""
+        if model_type == ModelType.LLM:
+            return [
+                "meta/llama-3-70b-instruct",
+                "meta/llama-3-8b-instruct",
+                "anthropic/claude-3-opus-20240229",
+                "anthropic/claude-3-sonnet-20240229"
+            ]
+        elif model_type == ModelType.VISION:
+            return [
+                "stability-ai/sdxl",
+                "stability-ai/stable-diffusion-3-medium",
+                "meta/llama-3-70b-vision",
+                "anthropic/claude-3-opus-20240229",
+                "anthropic/claude-3-sonnet-20240229"
+            ]
+        elif model_type == ModelType.AUDIO:
+            return [
+                "openai/whisper",
+                "suno-ai/bark"
+            ]
+        else:
+            return []
+    
+    def get_config(self) -> Dict[str, Any]:
+        """Get provider configuration"""
+        # Return a copy without sensitive information
+        config_copy = self.config.copy()
+        if "api_token" in config_copy:
+            config_copy["api_token"] = "***" if config_copy["api_token"] else ""
+        return config_copy
+    
+    def is_reasoning_model(self, model_name: str) -> bool:
+        """Check if the model is optimized for reasoning tasks"""
+        reasoning_models = ["llama-3-70b", "claude-3-opus", "claude-3-sonnet"]
+        return any(rm in model_name.lower() for rm in reasoning_models)