isa-model 0.3.5__py3-none-any.whl → 0.3.6__py3-none-any.whl

isa_model/core/model_manager.py

@@ -1,213 +0,0 @@
-from typing import Dict, Optional, List, Any
-import logging
-from pathlib import Path
-from huggingface_hub import hf_hub_download, snapshot_download
-from huggingface_hub.errors import HfHubHTTPError
-from .model_storage import ModelStorage, LocalModelStorage
-from .model_registry import ModelRegistry, ModelType, ModelCapability
-
-logger = logging.getLogger(__name__)
-
-class ModelManager:
-    """Model management service for handling model downloads, versions, and caching"""
-    
-    # 统一的模型计费信息 (per 1M tokens)
-    MODEL_PRICING = {
-        # OpenAI Models
-        "openai": {
-            "gpt-4o-mini": {"input": 0.15, "output": 0.6},
-            "gpt-4.1-mini": {"input": 0.4, "output": 1.6},
-            "gpt-4.1-nano": {"input": 0.1, "output": 0.4},
-            "gpt-4o": {"input": 5.0, "output": 15.0},
-            "gpt-4-turbo": {"input": 10.0, "output": 30.0},
-            "gpt-4": {"input": 30.0, "output": 60.0},
-            "gpt-3.5-turbo": {"input": 0.5, "output": 1.5},
-            "text-embedding-3-small": {"input": 0.02, "output": 0.0},
-            "text-embedding-3-large": {"input": 0.13, "output": 0.0},
-            "whisper-1": {"input": 6.0, "output": 0.0},
-            "tts-1": {"input": 15.0, "output": 0.0},
-            "tts-1-hd": {"input": 30.0, "output": 0.0},
-        },
-        # Ollama Models (免费本地模型)
-        "ollama": {
-            "llama3.2:3b-instruct-fp16": {"input": 0.0, "output": 0.0},
-            "llama3.2-vision:latest": {"input": 0.0, "output": 0.0},
-            "bge-m3": {"input": 0.0, "output": 0.0},
-        },
-        # Replicate Models
-        "replicate": {
-            "black-forest-labs/flux-schnell": {"input": 3.0, "output": 0.0},  # $3 per 1000 images
-            "black-forest-labs/flux-kontext-pro": {"input": 40.0, "output": 0.0},  # $0.04 per image = $40 per 1000 images
-            "meta/meta-llama-3-8b-instruct": {"input": 0.05, "output": 0.25},
-            "kokoro-82m": {"input": 0.0, "output": 0.4},  # ~$0.0004 per second
-            "jaaari/kokoro-82m:f559560eb822dc509045f3921a1921234918b91739db4bf3daab2169b71c7a13": {"input": 0.0, "output": 0.4},
-        },
-        # YYDS Models
-        "yyds": {
-            "claude-sonnet-4-20250514": {"input": 4.5, "output": 22.5},  # $0.0045/1K = $4.5/1M, $0.0225/1K = $22.5/1M
-            "claude-3-5-sonnet-20241022": {"input": 3.0, "output": 15.0},  # $0.003/1K = $3.0/1M, $0.015/1K = $15.0/1M
-        }
-    }
-    
-    def __init__(self, 
-                 storage: Optional[ModelStorage] = None,
-                 registry: Optional[ModelRegistry] = None):
-        self.storage = storage or LocalModelStorage()
-        self.registry = registry or ModelRegistry()
-    
-    def get_model_pricing(self, provider: str, model_name: str) -> Dict[str, float]:
-        """获取模型定价信息"""
-        return self.MODEL_PRICING.get(provider, {}).get(model_name, {"input": 0.0, "output": 0.0})
-    
-    def calculate_cost(self, provider: str, model_name: str, input_tokens: int, output_tokens: int) -> float:
-        """计算请求成本"""
-        pricing = self.get_model_pricing(provider, model_name)
-        input_cost = (input_tokens / 1_000_000) * pricing["input"]
-        output_cost = (output_tokens / 1_000_000) * pricing["output"]
-        return input_cost + output_cost
-    
-    def get_cheapest_model(self, provider: str, model_type: str = "llm") -> Optional[str]:
-        """获取最便宜的模型"""
-        provider_models = self.MODEL_PRICING.get(provider, {})
-        if not provider_models:
-            return None
-        
-        # 计算每个模型的平均成本 (假设输入输出比例 1:1)
-        cheapest_model = None
-        lowest_cost = float('inf')
-        
-        for model_name, pricing in provider_models.items():
-            avg_cost = (pricing["input"] + pricing["output"]) / 2
-            if avg_cost < lowest_cost:
-                lowest_cost = avg_cost
-                cheapest_model = model_name
-        
-        return cheapest_model
-
-    async def get_model(self, 
-                       model_id: str, 
-                       repo_id: str,
-                       model_type: ModelType,
-                       capabilities: List[ModelCapability],
-                       revision: Optional[str] = None,
-                       force_download: bool = False) -> Optional[Path]:
-        """
-        Get model files, downloading if necessary
-        
-        Args:
-            model_id: Unique identifier for the model
-            repo_id: Hugging Face repository ID
-            model_type: Type of model (LLM, embedding, etc.)
-            capabilities: List of model capabilities
-            revision: Specific model version/tag
-            force_download: Force re-download even if cached
-            
-        Returns:
-            Path to the model files or None if failed
-        """
-        # Check if model is already downloaded
-        if not force_download:
-            model_path = await self.storage.load_model(model_id)
-            if model_path:
-                logger.info(f"Using cached model {model_id}")
-                return model_path
-        
-        try:
-            # Download model files
-            logger.info(f"Downloading model {model_id} from {repo_id}")
-            model_dir = Path(f"./models/temp/{model_id}")
-            model_dir.mkdir(parents=True, exist_ok=True)
-            
-            snapshot_download(
-                repo_id=repo_id,
-                revision=revision,
-                local_dir=model_dir,
-                local_dir_use_symlinks=False
-            )
-            
-            # Save model and metadata
-            metadata = {
-                "repo_id": repo_id,
-                "revision": revision,
-                "downloaded_at": str(Path(model_dir).stat().st_mtime)
-            }
-            
-            # Register model
-            self.registry.register_model(
-                model_id=model_id,
-                model_type=model_type,
-                capabilities=capabilities,
-                metadata=metadata
-            )
-            
-            # Save model files
-            await self.storage.save_model(model_id, str(model_dir), metadata)
-            
-            return await self.storage.load_model(model_id)
-            
-        except HfHubHTTPError as e:
-            logger.error(f"Failed to download model {model_id}: {e}")
-            return None
-        except Exception as e:
-            logger.error(f"Unexpected error downloading model {model_id}: {e}")
-            return None
-    
-    async def list_models(self) -> List[Dict[str, Any]]:
-        """List all downloaded models with their metadata"""
-        models = await self.storage.list_models()
-        return [
-            {
-                "model_id": model_id,
-                **metadata,
-                **(self.registry.get_model_info(model_id) or {})
-            }
-            for model_id, metadata in models.items()
-        ]
-    
-    async def remove_model(self, model_id: str) -> bool:
-        """Remove a model and its metadata"""
-        try:
-            # Remove from storage
-            storage_success = await self.storage.delete_model(model_id)
-            
-            # Unregister from registry
-            registry_success = self.registry.unregister_model(model_id)
-            
-            return storage_success and registry_success
-            
-        except Exception as e:
-            logger.error(f"Failed to remove model {model_id}: {e}")
-            return False
-    
-    async def get_model_info(self, model_id: str) -> Optional[Dict[str, Any]]:
-        """Get information about a specific model"""
-        storage_info = await self.storage.get_metadata(model_id)
-        registry_info = self.registry.get_model_info(model_id)
-        
-        if not storage_info and not registry_info:
-            return None
-            
-        return {
-            **(storage_info or {}),
-            **(registry_info or {})
-        }
-    
-    async def update_model(self, 
-                          model_id: str, 
-                          repo_id: str,
-                          model_type: ModelType,
-                          capabilities: List[ModelCapability],
-                          revision: Optional[str] = None) -> bool:
-        """Update a model to a new version"""
-        try:
-            return bool(await self.get_model(
-                model_id=model_id,
-                repo_id=repo_id,
-                model_type=model_type,
-                capabilities=capabilities,
-                revision=revision,
-                force_download=True
-            ))
-        except Exception as e:
-            logger.error(f"Failed to update model {model_id}: {e}")
-            return False 

isa_model/core/model_registry.py

@@ -1,375 +0,0 @@
-from typing import Dict, List, Optional, Any
-from enum import Enum
-import logging
-from pathlib import Path
-import json
-import sqlite3
-from datetime import datetime
-import threading
-
-logger = logging.getLogger(__name__)
-
-class ModelCapability(str, Enum):
-    """Model capabilities"""
-    TEXT_GENERATION = "text_generation"
-    CHAT = "chat"
-    EMBEDDING = "embedding"
-    RERANKING = "reranking"
-    REASONING = "reasoning"
-    IMAGE_GENERATION = "image_generation"
-    IMAGE_ANALYSIS = "image_analysis"
-    AUDIO_TRANSCRIPTION = "audio_transcription"
-    IMAGE_UNDERSTANDING = "image_understanding"
-    UI_DETECTION = "ui_detection"
-    OCR = "ocr"
-    TABLE_DETECTION = "table_detection"
-    TABLE_STRUCTURE_RECOGNITION = "table_structure_recognition"
-
-class ModelType(str, Enum):
-    """Model types"""
-    LLM = "llm"
-    EMBEDDING = "embedding"
-    RERANK = "rerank"
-    IMAGE = "image"
-    AUDIO = "audio"
-    VIDEO = "video"
-    VISION = "vision"
-
-class ModelRegistry:
-    """Model registry with SQLite or Supabase backend"""
-    
-    def __init__(self, db_path: str = "./models/model_registry.db", use_supabase: bool = None):
-        # Auto-detect Supabase if environment variables are set
-        if use_supabase is None:
-            import os
-            use_supabase = bool(os.getenv("SUPABASE_URL") and os.getenv("SUPABASE_ANON_KEY"))
-        
-        self.use_supabase = use_supabase
-        
-        if self.use_supabase:
-            try:
-                from .storage.supabase_storage import SupabaseModelRegistry
-                self.backend = SupabaseModelRegistry()
-                logger.info("Using Supabase backend for model registry")
-            except ImportError as e:
-                logger.warning(f"Supabase not available, falling back to SQLite: {e}")
-                self.use_supabase = False
-        
-        if not self.use_supabase:
-            # Use SQLite backend
-            self.db_path = Path(db_path)
-            self.db_path.parent.mkdir(parents=True, exist_ok=True)
-            self._lock = threading.Lock()
-            self._initialize_database()
-            self.backend = None
-            logger.info("Using SQLite backend for model registry")
-    
-    def _initialize_database(self):
-        """Initialize SQLite database with required tables"""
-        with sqlite3.connect(self.db_path) as conn:
-            conn.execute("""
-                CREATE TABLE IF NOT EXISTS models (
-                    model_id TEXT PRIMARY KEY,
-                    model_type TEXT NOT NULL,
-                    metadata TEXT,
-                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
-                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
-                )
-            """)
-            
-            conn.execute("""
-                CREATE TABLE IF NOT EXISTS model_capabilities (
-                    model_id TEXT,
-                    capability TEXT,
-                    PRIMARY KEY (model_id, capability),
-                    FOREIGN KEY (model_id) REFERENCES models(model_id) ON DELETE CASCADE
-                )
-            """)
-            
-            conn.execute("""
-                CREATE INDEX IF NOT EXISTS idx_model_type ON models(model_type)
-            """)
-            
-            conn.execute("""
-                CREATE INDEX IF NOT EXISTS idx_capability ON model_capabilities(capability)
-            """)
-            
-            conn.commit()
-    
-    def register_model(self, 
-                      model_id: str,
-                      model_type: ModelType,
-                      capabilities: List[ModelCapability],
-                      metadata: Dict[str, Any]) -> bool:
-        """Register a model with its capabilities and metadata"""
-        if self.use_supabase:
-            return self.backend.register_model(
-                model_id=model_id,
-                model_type=model_type.value,
-                capabilities=[cap.value for cap in capabilities],
-                metadata=metadata
-            )
-        
-        # SQLite implementation
-        try:
-            with self._lock:
-                with sqlite3.connect(self.db_path) as conn:
-                    # Insert or update model
-                    conn.execute("""
-                        INSERT OR REPLACE INTO models 
-                        (model_id, model_type, metadata, updated_at)
-                        VALUES (?, ?, ?, CURRENT_TIMESTAMP)
-                    """, (model_id, model_type.value, json.dumps(metadata)))
-                    
-                    # Clear existing capabilities
-                    conn.execute("DELETE FROM model_capabilities WHERE model_id = ?", (model_id,))
-                    
-                    # Insert new capabilities
-                    for capability in capabilities:
-                        conn.execute("""
-                            INSERT INTO model_capabilities (model_id, capability)
-                            VALUES (?, ?)
-                        """, (model_id, capability.value))
-                    
-                    conn.commit()
-                    
-            logger.info(f"Registered model {model_id}")
-            return True
-            
-        except Exception as e:
-            logger.error(f"Failed to register model {model_id}: {e}")
-            return False
-    
-    def unregister_model(self, model_id: str) -> bool:
-        """Unregister a model"""
-        try:
-            with self._lock:
-                with sqlite3.connect(self.db_path) as conn:
-                    cursor = conn.execute("DELETE FROM models WHERE model_id = ?", (model_id,))
-                    conn.commit()
-                    
-                    if cursor.rowcount > 0:
-                        logger.info(f"Unregistered model {model_id}")
-                        return True
-                    return False
-                    
-        except Exception as e:
-            logger.error(f"Failed to unregister model {model_id}: {e}")
-            return False
-    
-    def get_model_info(self, model_id: str) -> Optional[Dict[str, Any]]:
-        """Get model information"""
-        try:
-            with sqlite3.connect(self.db_path) as conn:
-                conn.row_factory = sqlite3.Row
-                
-                # Get model info
-                model_row = conn.execute("""
-                    SELECT model_id, model_type, metadata, created_at, updated_at
-                    FROM models WHERE model_id = ?
-                """, (model_id,)).fetchone()
-                
-                if not model_row:
-                    return None
-                
-                # Get capabilities
-                capabilities = conn.execute("""
-                    SELECT capability FROM model_capabilities WHERE model_id = ?
-                """, (model_id,)).fetchall()
-                
-                model_info = {
-                    "model_id": model_row["model_id"],
-                    "type": model_row["model_type"],
-                    "capabilities": [cap["capability"] for cap in capabilities],
-                    "metadata": json.loads(model_row["metadata"]) if model_row["metadata"] else {},
-                    "created_at": model_row["created_at"],
-                    "updated_at": model_row["updated_at"]
-                }
-                
-                return model_info
-                
-        except Exception as e:
-            logger.error(f"Failed to get model info for {model_id}: {e}")
-            return None
-    
-    def get_models_by_type(self, model_type: ModelType) -> Dict[str, Dict[str, Any]]:
-        """Get all models of a specific type"""
-        try:
-            with sqlite3.connect(self.db_path) as conn:
-                conn.row_factory = sqlite3.Row
-                
-                models = conn.execute("""
-                    SELECT model_id, model_type, metadata, created_at, updated_at
-                    FROM models WHERE model_type = ?
-                """, (model_type.value,)).fetchall()
-                
-                result = {}
-                for model in models:
-                    model_id = model["model_id"]
-                    
-                    # Get capabilities for this model
-                    capabilities = conn.execute("""
-                        SELECT capability FROM model_capabilities WHERE model_id = ?
-                    """, (model_id,)).fetchall()
-                    
-                    result[model_id] = {
-                        "type": model["model_type"],
-                        "capabilities": [cap["capability"] for cap in capabilities],
-                        "metadata": json.loads(model["metadata"]) if model["metadata"] else {},
-                        "created_at": model["created_at"],
-                        "updated_at": model["updated_at"]
-                    }
-                
-                return result
-                
-        except Exception as e:
-            logger.error(f"Failed to get models by type {model_type}: {e}")
-            return {}
-    
-    def get_models_by_capability(self, capability: ModelCapability) -> Dict[str, Dict[str, Any]]:
-        """Get all models with a specific capability"""
-        try:
-            with sqlite3.connect(self.db_path) as conn:
-                conn.row_factory = sqlite3.Row
-                
-                models = conn.execute("""
-                    SELECT DISTINCT m.model_id, m.model_type, m.metadata, m.created_at, m.updated_at
-                    FROM models m
-                    JOIN model_capabilities mc ON m.model_id = mc.model_id
-                    WHERE mc.capability = ?
-                """, (capability.value,)).fetchall()
-                
-                result = {}
-                for model in models:
-                    model_id = model["model_id"]
-                    
-                    # Get all capabilities for this model
-                    capabilities = conn.execute("""
-                        SELECT capability FROM model_capabilities WHERE model_id = ?
-                    """, (model_id,)).fetchall()
-                    
-                    result[model_id] = {
-                        "type": model["model_type"],
-                        "capabilities": [cap["capability"] for cap in capabilities],
-                        "metadata": json.loads(model["metadata"]) if model["metadata"] else {},
-                        "created_at": model["created_at"],
-                        "updated_at": model["updated_at"]
-                    }
-                
-                return result
-                
-        except Exception as e:
-            logger.error(f"Failed to get models by capability {capability}: {e}")
-            return {}
-    
-    def has_capability(self, model_id: str, capability: ModelCapability) -> bool:
-        """Check if a model has a specific capability"""
-        try:
-            with sqlite3.connect(self.db_path) as conn:
-                result = conn.execute("""
-                    SELECT 1 FROM model_capabilities 
-                    WHERE model_id = ? AND capability = ?
-                """, (model_id, capability.value)).fetchone()
-                
-                return result is not None
-                
-        except Exception as e:
-            logger.error(f"Failed to check capability for {model_id}: {e}")
-            return False
-    
-    def list_models(self) -> Dict[str, Dict[str, Any]]:
-        """List all registered models"""
-        try:
-            with sqlite3.connect(self.db_path) as conn:
-                conn.row_factory = sqlite3.Row
-                
-                models = conn.execute("""
-                    SELECT model_id, model_type, metadata, created_at, updated_at
-                    FROM models ORDER BY created_at DESC
-                """).fetchall()
-                
-                result = {}
-                for model in models:
-                    model_id = model["model_id"]
-                    
-                    # Get capabilities for this model
-                    capabilities = conn.execute("""
-                        SELECT capability FROM model_capabilities WHERE model_id = ?
-                    """, (model_id,)).fetchall()
-                    
-                    result[model_id] = {
-                        "type": model["model_type"],
-                        "capabilities": [cap["capability"] for cap in capabilities],
-                        "metadata": json.loads(model["metadata"]) if model["metadata"] else {},
-                        "created_at": model["created_at"],
-                        "updated_at": model["updated_at"]
-                    }
-                
-                return result
-                
-        except Exception as e:
-            logger.error(f"Failed to list models: {e}")
-            return {}
-    
-    def get_stats(self) -> Dict[str, Any]:
-        """Get registry statistics"""
-        try:
-            with sqlite3.connect(self.db_path) as conn:
-                # Count total models
-                total_models = conn.execute("SELECT COUNT(*) FROM models").fetchone()[0]
-                
-                # Count by type
-                type_counts = dict(conn.execute("""
-                    SELECT model_type, COUNT(*) FROM models GROUP BY model_type
-                """).fetchall())
-                
-                # Count by capability
-                capability_counts = dict(conn.execute("""
-                    SELECT capability, COUNT(*) FROM model_capabilities GROUP BY capability
-                """).fetchall())
-                
-                return {
-                    "total_models": total_models,
-                    "models_by_type": type_counts,
-                    "models_by_capability": capability_counts
-                }
-                
-        except Exception as e:
-            logger.error(f"Failed to get stats: {e}")
-            return {}
-    
-    def search_models(self, query: str) -> Dict[str, Dict[str, Any]]:
-        """Search models by name or metadata"""
-        try:
-            with sqlite3.connect(self.db_path) as conn:
-                conn.row_factory = sqlite3.Row
-                
-                models = conn.execute("""
-                    SELECT model_id, model_type, metadata, created_at, updated_at
-                    FROM models 
-                    WHERE model_id LIKE ? OR metadata LIKE ?
-                    ORDER BY created_at DESC
-                """, (f"%{query}%", f"%{query}%")).fetchall()
-                
-                result = {}
-                for model in models:
-                    model_id = model["model_id"]
-                    
-                    # Get capabilities for this model
-                    capabilities = conn.execute("""
-                        SELECT capability FROM model_capabilities WHERE model_id = ?
-                    """, (model_id,)).fetchall()
-                    
-                    result[model_id] = {
-                        "type": model["model_type"],
-                        "capabilities": [cap["capability"] for cap in capabilities],
-                        "metadata": json.loads(model["metadata"]) if model["metadata"] else {},
-                        "created_at": model["created_at"],
-                        "updated_at": model["updated_at"]
-                    }
-                
-                return result
-                
-        except Exception as e:
-            logger.error(f"Failed to search models with query '{query}': {e}")
-            return {}