isa-model 0.0.2__py3-none-any.whl → 0.3.1__py3-none-any.whl

isa_model/__init__.py

@@ -2,4 +2,4 @@
 isA_Model - A simple interface for AI model integration
 """
 
-__version__ = "0.1.0" 
+__version__ = "0.0.1"

isa_model/core/model_manager.py

@@ -2,7 +2,7 @@ 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.utils import HfHubHTTPError
+from huggingface_hub.errors import HfHubHTTPError
 from .model_storage import ModelStorage, LocalModelStorage
 from .model_registry import ModelRegistry, ModelType, ModelCapability
 
@@ -11,19 +11,81 @@ 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},
+        }
+    }
+    
     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) -> Path:
+                       force_download: bool = False) -> Optional[Path]:
         """
         Get model files, downloading if necessary
         
@@ -36,7 +98,7 @@ class ModelManager:
             force_download: Force re-download even if cached
             
         Returns:
-            Path to the model files
+            Path to the model files or None if failed
         """
         # Check if model is already downloaded
         if not force_download:
@@ -80,7 +142,10 @@ class ModelManager:
             
         except HfHubHTTPError as e:
             logger.error(f"Failed to download model {model_id}: {e}")
-            raise
+            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"""

isa_model/core/model_registry.py

@@ -3,6 +3,9 @@ from enum import Enum
 import logging
 from pathlib import Path
 import json
+import sqlite3
+from datetime import datetime
+import threading
 
 logger = logging.getLogger(__name__)
 
@@ -29,27 +32,45 @@ class ModelType(str, Enum):
     VISION = "vision"
 
 class ModelRegistry:
-    """Registry for model metadata and capabilities"""
+    """SQLite-based registry for model metadata and capabilities"""
     
-    def __init__(self, registry_file: str = "./models/model_registry.json"):
-        self.registry_file = Path(registry_file)
-        self.registry: Dict[str, Dict[str, Any]] = {}
-        self._load_registry()
+    def __init__(self, db_path: str = "./models/model_registry.db"):
+        self.db_path = Path(db_path)
+        self.db_path.parent.mkdir(parents=True, exist_ok=True)
+        self._lock = threading.Lock()
+        self._initialize_database()
     
-    def _load_registry(self):
-        """Load model registry from file"""
-        if self.registry_file.exists():
-            with open(self.registry_file, 'r') as f:
-                self.registry = json.load(f)
-        else:
-            self.registry = {}
-            self._save_registry()
-    
-    def _save_registry(self):
-        """Save model registry to file"""
-        self.registry_file.parent.mkdir(parents=True, exist_ok=True)
-        with open(self.registry_file, 'w') as f:
-            json.dump(self.registry, f, indent=2)
+    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,
@@ -58,14 +79,30 @@ class ModelRegistry:
                       metadata: Dict[str, Any]) -> bool:
         """Register a model with its capabilities and metadata"""
         try:
-            self.registry[model_id] = {
-                "type": model_type,
-                "capabilities": [cap.value for cap in capabilities],
-                "metadata": metadata
-            }
-            self._save_registry()
+            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
@@ -73,43 +110,233 @@ class ModelRegistry:
     def unregister_model(self, model_id: str) -> bool:
         """Unregister a model"""
         try:
-            if model_id in self.registry:
-                del self.registry[model_id]
-                self._save_registry()
-                logger.info(f"Unregistered model {model_id}")
-                return True
-            return False
+            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"""
-        return self.registry.get(model_id)
+        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"""
-        return {
-            model_id: info
-            for model_id, info in self.registry.items()
-            if info["type"] == model_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"""
-        return {
-            model_id: info
-            for model_id, info in self.registry.items()
-            if capability.value in info["capabilities"]
-        }
+        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"""
-        model_info = self.get_model_info(model_id)
-        if not model_info:
+        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
-        return capability.value in model_info["capabilities"]
     
     def list_models(self) -> Dict[str, Dict[str, Any]]:
         """List all registered models"""
-        return self.registry 
+        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 {}