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

isa_model/inference/services/audio/replicate_tts_service.py

@@ -4,42 +4,45 @@ import replicate
 from tenacity import retry, stop_after_attempt, wait_exponential
 
 from isa_model.inference.services.audio.base_tts_service import BaseTTSService
-from isa_model.inference.providers.base_provider import BaseProvider
-from isa_model.inference.billing_tracker import ServiceType
 
 logger = logging.getLogger(__name__)
 
 class ReplicateTTSService(BaseTTSService):
     """
-    Replicate Text-to-Speech service using Kokoro model.
+    Replicate Text-to-Speech service using Kokoro model with unified architecture.
     High-quality voice synthesis with multiple voice options.
     """
     
-    def __init__(self, provider: 'BaseProvider', model_name: str = "jaaari/kokoro-82m:f559560eb822dc509045f3921a1921234918b91739db4bf3daab2169b71c7a13"):
-        super().__init__(provider, model_name)
+    def __init__(self, provider_name: str, model_name: str = "jaaari/kokoro-82m:f559560eb822dc509045f3921a1921234918b91739db4bf3daab2169b71c7a13", **kwargs):
+        super().__init__(provider_name, model_name, **kwargs)
         
-        # Get full configuration from provider (including sensitive data)
-        provider_config = provider.get_full_config()
+        # Get configuration from centralized config manager
+        provider_config = self.get_provider_config()
         
         # Set up Replicate API token from provider configuration
-        self.api_token = provider_config.get('api_token') or provider_config.get('replicate_api_token')
-        if not self.api_token:
-            raise ValueError("Replicate API token not found in provider configuration")
-        
-        # Set environment variable for replicate library
-        import os
-        os.environ['REPLICATE_API_TOKEN'] = self.api_token
-        
-        # Available voices for Kokoro model
-        self.available_voices = [
-            "af_bella", "af_nicole", "af_sarah", "af_sky", "am_adam", "am_michael"
-        ]
-        
-        # Default settings
-        self.default_voice = "af_nicole"
-        self.default_speed = 1.0
-        
-        logger.info(f"Initialized ReplicateTTSService with model '{self.model_name}'")
+        try:
+            self.api_token = provider_config.get('api_key') or provider_config.get('replicate_api_token')
+            if not self.api_token:
+                raise ValueError("Replicate API token not found in provider configuration")
+            
+            # Set environment variable for replicate library
+            import os
+            os.environ['REPLICATE_API_TOKEN'] = self.api_token
+            
+            # Available voices for Kokoro model
+            self.available_voices = [
+                "af_bella", "af_nicole", "af_sarah", "af_sky", "am_adam", "am_michael"
+            ]
+            
+            # Default settings
+            self.default_voice = "af_nicole"
+            self.default_speed = 1.0
+            
+            logger.info(f"Initialized ReplicateTTSService with model '{self.model_name}'")
+            
+        except Exception as e:
+            logger.error(f"Failed to initialize Replicate client: {e}")
+            raise ValueError(f"Failed to initialize Replicate client: {e}") from e
     
     @retry(
         stop=stop_after_attempt(3),
@@ -51,8 +54,8 @@ class ReplicateTTSService(BaseTTSService):
         text: str, 
         voice: Optional[str] = None,
         speed: float = 1.0,
-        pitch: Optional[float] = None,
-        volume: Optional[float] = None
+        pitch: float = 1.0,
+        format: str = "wav"
     ) -> Dict[str, Any]:
         """Synthesize speech from text using Kokoro model"""
         try:
@@ -99,8 +102,8 @@ class ReplicateTTSService(BaseTTSService):
             estimated_duration_seconds = (words / 150.0) * 60.0 / speed
             
             # Track usage for billing
-            self._track_usage(
-                service_type=ServiceType.AUDIO_TTS,
+            await self._track_usage(
+                service_type="audio_tts",
                 operation="synthesize_speech",
                 input_tokens=0,
                 output_tokens=0,
@@ -115,15 +118,24 @@ class ReplicateTTSService(BaseTTSService):
                 }
             )
             
+            # Download audio data for return format consistency
+            import aiohttp
+            async with aiohttp.ClientSession() as session:
+                async with session.get(audio_url) as response:
+                    response.raise_for_status()
+                    audio_data = await response.read()
+            
             result = {
-                "audio_url": audio_url,
-                "text": text,
-                "voice": selected_voice,
-                "speed": speed,
-                "duration_seconds": estimated_duration_seconds,
+                "audio_data": audio_data,
+                "format": "wav",  # Kokoro typically outputs WAV
+                "duration": estimated_duration_seconds,
+                "sample_rate": 22050,
+                "audio_url": audio_url,  # Keep URL for reference
                 "metadata": {
                     "model": self.model_name,
                     "provider": "replicate",
+                    "voice": selected_voice,
+                    "speed": speed,
                     "voice_options": self.available_voices
                 }
             }
@@ -137,36 +149,29 @@ class ReplicateTTSService(BaseTTSService):
     
     async def synthesize_speech_to_file(
         self, 
-        text: str, 
+        text: str,
         output_path: str,
         voice: Optional[str] = None,
         speed: float = 1.0,
-        pitch: Optional[float] = None,
-        volume: Optional[float] = None
+        pitch: float = 1.0,
+        format: str = "wav"
     ) -> Dict[str, Any]:
         """Synthesize speech and save to file"""
         try:
-            # Get audio URL
-            result = await self.synthesize_speech(text, voice, speed, pitch, volume)
-            audio_url = result["audio_url"]
+            # Get synthesis result
+            result = await self.synthesize_speech(text, voice, speed, pitch, format)
+            audio_data = result["audio_data"]
             
-            # Download and save audio
-            import aiohttp
-            import aiofiles
+            # Save audio data to file
+            with open(output_path, 'wb') as f:
+                f.write(audio_data)
             
-            async with aiohttp.ClientSession() as session:
-                async with session.get(audio_url) as response:
-                    response.raise_for_status()
-                    audio_data = await response.read()
-                    
-                    async with aiofiles.open(output_path, 'wb') as f:
-                        await f.write(audio_data)
-            
-            result["output_path"] = output_path
-            result["file_size"] = len(audio_data)
-            
-            logger.info(f"Audio saved to: {output_path}")
-            return result
+            return {
+                "file_path": output_path,
+                "duration": result["duration"],
+                "sample_rate": result["sample_rate"],
+                "file_size": len(audio_data)
+            }
             
         except Exception as e:
             logger.error(f"Error saving audio to file: {e}")

isa_model/inference/services/base_service.py

@@ -1,19 +1,50 @@
 from abc import ABC, abstractmethod
 from typing import Dict, Any, List, Union, AsyncGenerator, TypeVar, Optional
-from isa_model.inference.providers.base_provider import BaseProvider
-from isa_model.inference.billing_tracker import track_usage, ServiceType, Provider
+from ...core.models.model_manager import ModelManager
+from ...core.config.config_manager import ConfigManager
+from ...core.types import Provider, ServiceType
 
 T = TypeVar('T')  # Generic type for responses
 
 class BaseService(ABC):
-    """Base class for all AI services"""
+    """Base class for all AI services - now uses centralized managers"""
     
-    def __init__(self, provider: 'BaseProvider', model_name: str):
-        self.provider = provider
+    def __init__(self, 
+                 provider_name: str, 
+                 model_name: str, 
+                 model_manager: Optional[ModelManager] = None,
+                 config_manager: Optional[ConfigManager] = None):
+        self.provider_name = provider_name
         self.model_name = model_name
-        self.config = provider.get_full_config()
+        self.model_manager = model_manager or ModelManager()
+        self.config_manager = config_manager or ConfigManager()
         
-    def _track_usage(
+        # Validate provider is configured
+        if not self.config_manager.is_provider_enabled(provider_name):
+            raise ValueError(f"Provider {provider_name} is not configured or enabled")
+    
+    def get_api_key(self) -> str:
+        """Get API key for the provider"""
+        api_key = self.config_manager.get_provider_api_key(self.provider_name)
+        if not api_key:
+            raise ValueError(f"No API key configured for provider {self.provider_name}")
+        return api_key
+    
+    def get_provider_config(self) -> Dict[str, Any]:
+        """Get provider configuration"""
+        config = self.config_manager.get_provider_config(self.provider_name)
+        if not config:
+            return {}
+        
+        return {
+            "api_key": config.api_key,
+            "api_base_url": config.api_base_url,
+            "organization": config.organization,
+            "rate_limit_rpm": config.rate_limit_rpm,
+            "rate_limit_tpm": config.rate_limit_tpm,
+        }
+        
+    async def _track_usage(
         self,
         service_type: Union[str, ServiceType],
         operation: str,
@@ -23,23 +54,30 @@ class BaseService(ABC):
         output_units: Optional[float] = None,
         metadata: Optional[Dict[str, Any]] = None
     ):
-        """Track usage for billing purposes"""
+        """Track usage for billing purposes using centralized billing tracker"""
         try:
-            # Determine provider name - try multiple attributes
-            provider_name = getattr(self.provider, 'name', None) or \
-                          getattr(self.provider, 'provider_name', None) or \
-                          getattr(self.provider, '__class__', type(None)).__name__.lower().replace('provider', '') or \
-                          'unknown'
+            # Calculate cost using centralized pricing
+            cost_usd = None
+            if input_tokens is not None and output_tokens is not None:
+                cost_usd = self.model_manager.calculate_cost(
+                    provider=self.provider_name,
+                    model_name=self.model_name,
+                    input_tokens=input_tokens,
+                    output_tokens=output_tokens
+                )
             
-            track_usage(
-                provider=provider_name,
-                service_type=service_type,
-                model_name=self.model_name,
+            # Track usage through model manager
+            self.model_manager.billing_tracker.track_model_usage(
+                model_id=self.model_name,
+                operation_type="inference",
+                provider=self.provider_name,
+                service_type=service_type if isinstance(service_type, str) else service_type.value,
                 operation=operation,
                 input_tokens=input_tokens,
                 output_tokens=output_tokens,
                 input_units=input_units,
                 output_units=output_units,
+                cost_usd=cost_usd,
                 metadata=metadata
             )
         except Exception as e:

isa_model/inference/services/embedding/base_embed_service.py

@@ -3,7 +3,71 @@ from typing import Dict, Any, List, Union, Optional
 from isa_model.inference.services.base_service import BaseService
 
 class BaseEmbedService(BaseService):
-    """Base class for embedding services"""
+    """Base class for embedding services with unified task dispatch"""
+    
+    async def invoke(
+        self, 
+        input_data: Union[str, List[str]],
+        task: Optional[str] = None,
+        **kwargs
+    ) -> Union[List[float], List[List[float]], List[Dict[str, Any]], Dict[str, Any]]:
+        """
+        统一的任务分发方法 - Base类提供通用实现
+        
+        Args:
+            input_data: 输入数据，可以是:
+                - str: 单个文本
+                - List[str]: 多个文本（批量处理）
+            task: 任务类型，支持多种embedding任务
+            **kwargs: 任务特定的附加参数
+            
+        Returns:
+            Various types depending on task
+        """
+        task = task or "embed"
+        
+        # ==================== 嵌入生成类任务 ====================
+        if task == "embed":
+            if isinstance(input_data, list):
+                return await self.create_text_embeddings(input_data)
+            else:
+                return await self.create_text_embedding(input_data)
+        elif task == "embed_batch":
+            if not isinstance(input_data, list):
+                input_data = [input_data]
+            return await self.create_text_embeddings(input_data)
+        elif task == "chunk_and_embed":
+            if isinstance(input_data, list):
+                raise ValueError("chunk_and_embed task requires single text input")
+            return await self.create_chunks(input_data, kwargs.get("metadata"))
+        elif task == "similarity":
+            embedding1 = kwargs.get("embedding1")
+            embedding2 = kwargs.get("embedding2")
+            if not embedding1 or not embedding2:
+                raise ValueError("similarity task requires embedding1 and embedding2 parameters")
+            similarity = await self.compute_similarity(embedding1, embedding2)
+            return {"similarity": similarity}
+        elif task == "find_similar":
+            query_embedding = kwargs.get("query_embedding")
+            candidate_embeddings = kwargs.get("candidate_embeddings")
+            if not query_embedding or not candidate_embeddings:
+                raise ValueError("find_similar task requires query_embedding and candidate_embeddings parameters")
+            return await self.find_similar_texts(
+                query_embedding, 
+                candidate_embeddings, 
+                kwargs.get("top_k", 5)
+            )
+        else:
+            raise NotImplementedError(f"{self.__class__.__name__} does not support task: {task}")
+    
+    def get_supported_tasks(self) -> List[str]:
+        """
+        获取支持的任务列表
+        
+        Returns:
+            List of supported task names
+        """
+        return ["embed", "embed_batch", "chunk_and_embed", "similarity", "find_similar"]
     
     @abstractmethod
     async def create_text_embedding(self, text: str) -> List[float]:

isa_model/inference/services/embedding/ollama_embed_service.py

@@ -3,44 +3,65 @@ import httpx
 import asyncio
 from typing import List, Dict, Any, Optional
 
-# 保留您指定的导入和框架结构
 from isa_model.inference.services.embedding.base_embed_service import BaseEmbedService
-from isa_model.inference.providers.base_provider import BaseProvider
 
 logger = logging.getLogger(__name__)
 
 class OllamaEmbedService(BaseEmbedService):
     """
-    Ollama embedding service.
-    此类遵循基础服务架构，但使用其自己的 HTTP 客户端与 Ollama API 通信，
-    而不依赖于注入的 backend 对象。
+    Ollama embedding service with unified architecture.
+    Uses direct HTTP client communication with Ollama API.
     """
     
-    def __init__(self, provider: 'BaseProvider', model_name: str = "bge-m3"):
-        # 保持对基类和 provider 的兼容
-        super().__init__(provider, model_name)
+    def __init__(self, provider_name: str, model_name: str = "bge-m3", **kwargs):
+        super().__init__(provider_name, model_name, **kwargs)
         
-        # 从基类继承的 self.config 中获取配置
-        host = self.config.get("host", "localhost")
-        port = self.config.get("port", 11434)
+        # Get configuration from centralized config manager
+        provider_config = self.get_provider_config()
         
-        # 创建并持有自己的 httpx 客户端实例
-        base_url = f"http://{host}:{port}"
-        self.client = httpx.AsyncClient(base_url=base_url, timeout=30.0)
+        # Initialize HTTP client with provider configuration
+        try:
+            host = provider_config.get("host", "localhost")
+            port = provider_config.get("port", 11434)
+            base_url = f"http://{host}:{port}"
+            
+            self.client = httpx.AsyncClient(base_url=base_url, timeout=30.0)
             
-        logger.info(f"Initialized OllamaEmbedService with model '{self.model_name}' at {base_url}")
+            logger.info(f"Initialized OllamaEmbedService with model '{self.model_name}' at {base_url}")
+            
+        except Exception as e:
+            logger.error(f"Failed to initialize Ollama client: {e}")
+            raise ValueError(f"Failed to initialize Ollama client: {e}") from e
     
     async def create_text_embedding(self, text: str) -> List[float]:
-        """为单个文本创建 embedding"""
+        """Create embedding for single text"""
         try:
             payload = {
                 "model": self.model_name,
                 "prompt": text
             }
-            # 使用自己的 client 实例，而不是 self.backend
+            
             response = await self.client.post("/api/embeddings", json=payload)
-            response.raise_for_status() # 检查请求是否成功
-            return response.json()["embedding"]
+            response.raise_for_status()
+            
+            result = response.json()
+            embedding = result["embedding"]
+            
+            # Track usage for billing (estimate token usage for Ollama)
+            estimated_tokens = len(text.split()) * 1.3  # Rough estimation
+            await self._track_usage(
+                service_type="embedding",
+                operation="create_text_embedding",
+                input_tokens=int(estimated_tokens),
+                output_tokens=0,
+                metadata={
+                    "model": self.model_name,
+                    "text_length": len(text),
+                    "estimated_tokens": int(estimated_tokens)
+                }
+            )
+            
+            return embedding
             
         except httpx.RequestError as e:
             logger.error(f"An error occurred while requesting {e.request.url!r}: {e}")
@@ -50,41 +71,70 @@ class OllamaEmbedService(BaseEmbedService):
             raise
     
     async def create_text_embeddings(self, texts: List[str]) -> List[List[float]]:
-        """为多个文本并发地创建 embeddings"""
+        """Create embeddings for multiple texts concurrently"""
         if not texts:
             return []
         
         tasks = [self.create_text_embedding(text) for text in texts]
         embeddings = await asyncio.gather(*tasks)
+        
+        # Track batch usage for billing
+        total_estimated_tokens = sum(len(text.split()) * 1.3 for text in texts)
+        await self._track_usage(
+            service_type="embedding",
+            operation="create_text_embeddings",
+            input_tokens=int(total_estimated_tokens),
+            output_tokens=0,
+            metadata={
+                "model": self.model_name,
+                "batch_size": len(texts),
+                "total_text_length": sum(len(t) for t in texts),
+                "estimated_tokens": int(total_estimated_tokens)
+            }
+        )
+        
         return embeddings
     
     async def create_chunks(self, text: str, metadata: Optional[Dict] = None) -> List[Dict]:
-        """将文本分块并为每个块创建 embedding"""
-        chunk_size = 200  # 单词数量
-        words = text.split()
-        chunk_texts = [" ".join(words[i:i+chunk_size]) for i in range(0, len(words), chunk_size)]
+        """Create text chunks with embeddings"""
+        chunk_size = 200  # words
+        overlap = 50     # word overlap between chunks
         
-        if not chunk_texts:
+        words = text.split()
+        if not words:
             return []
-
-        embeddings = await self.create_text_embeddings(chunk_texts)
         
-        chunks = [
-            {
+        chunks = []
+        chunk_texts = []
+        
+        for i in range(0, len(words), chunk_size - overlap):
+            chunk_words = words[i:i + chunk_size]
+            chunk_text = " ".join(chunk_words)
+            chunk_texts.append(chunk_text)
+            
+            chunks.append({
                 "text": chunk_text,
-                "embedding": emb,
+                "start_index": i,
+                "end_index": min(i + chunk_size, len(words)),
                 "metadata": metadata or {}
-            }
-            for chunk_text, emb in zip(chunk_texts, embeddings)
-        ]
-            
+            })
+        
+        # Get embeddings for all chunks
+        embeddings = await self.create_text_embeddings(chunk_texts)
+        
+        # Add embeddings to chunks
+        for chunk, embedding in zip(chunks, embeddings):
+            chunk["embedding"] = embedding
+        
         return chunks
     
     async def compute_similarity(self, embedding1: List[float], embedding2: List[float]) -> float:
-        """计算两个嵌入向量之间的余弦相似度"""
+        """Compute cosine similarity between two embeddings"""
+        import math
+        
         dot_product = sum(a * b for a, b in zip(embedding1, embedding2))
-        norm1 = sum(a * a for a in embedding1) ** 0.5
-        norm2 = sum(b * b for b in embedding2) ** 0.5
+        norm1 = math.sqrt(sum(a * a for a in embedding1))
+        norm2 = math.sqrt(sum(b * b for b in embedding2))
         
         if norm1 * norm2 == 0:
             return 0.0
@@ -99,9 +149,13 @@ class OllamaEmbedService(BaseEmbedService):
     ) -> List[Dict[str, Any]]:
         """Find most similar texts based on embeddings"""
         similarities = []
+        
         for i, candidate in enumerate(candidate_embeddings):
             similarity = await self.compute_similarity(query_embedding, candidate)
-            similarities.append({"index": i, "similarity": similarity})
+            similarities.append({
+                "index": i, 
+                "similarity": similarity
+            })
         
         # Sort by similarity in descending order and return top_k
         similarities.sort(key=lambda x: x["similarity"], reverse=True)
@@ -109,15 +163,21 @@ class OllamaEmbedService(BaseEmbedService):
     
     def get_embedding_dimension(self) -> int:
         """Get the dimension of embeddings produced by this service"""
-        # BGE-M3 produces 1024-dimensional embeddings
-        return 1024
+        # Model-specific dimensions
+        model_dimensions = {
+            "bge-m3": 1024,
+            "bge-large": 1024,
+            "all-minilm": 384,
+            "nomic-embed-text": 768
+        }
+        return model_dimensions.get(self.model_name, 1024)
     
     def get_max_input_length(self) -> int:
         """Get maximum input text length supported"""
-        # BGE-M3 supports up to 8192 tokens
+        # Most Ollama embedding models support up to 8192 tokens
         return 8192
 
     async def close(self):
-        """关闭内置的 HTTP 客户端"""
+        """Cleanup resources"""
         await self.client.aclose()
-        logger.info("OllamaEmbedService's internal client has been closed.")
+        logger.info("OllamaEmbedService client has been closed.")

isa_model/inference/services/embedding/openai_embed_service.py

@@ -5,8 +5,6 @@ from openai import AsyncOpenAI
 from tenacity import retry, stop_after_attempt, wait_exponential
 
 from isa_model.inference.services.embedding.base_embed_service import BaseEmbedService
-from isa_model.inference.providers.base_provider import BaseProvider
-from isa_model.inference.billing_tracker import ServiceType
 
 logger = logging.getLogger(__name__)
 
@@ -16,11 +14,11 @@ class OpenAIEmbedService(BaseEmbedService):
     Provides high-quality embeddings for production use.
     """
     
-    def __init__(self, provider: 'BaseProvider', model_name: str = "text-embedding-3-small"):
-        super().__init__(provider, model_name)
+    def __init__(self, provider_name: str, model_name: str = "text-embedding-3-small", **kwargs):
+        super().__init__(provider_name, model_name, **kwargs)
         
-        # Get full configuration from provider (including sensitive data)
-        provider_config = provider.get_full_config()
+        # Get configuration from centralized config manager
+        provider_config = self.get_provider_config()
         
         # Initialize AsyncOpenAI client with provider configuration
         try:
@@ -67,8 +65,8 @@ class OpenAIEmbedService(BaseEmbedService):
             usage = getattr(response, 'usage', None)
             if usage:
                 total_tokens = getattr(usage, 'total_tokens', 0)
-                self._track_usage(
-                    service_type=ServiceType.EMBEDDING,
+                await self._track_usage(
+                    service_type="embedding",
                     operation="create_text_embedding",
                     input_tokens=total_tokens,
                     output_tokens=0,
@@ -112,8 +110,8 @@ class OpenAIEmbedService(BaseEmbedService):
             usage = getattr(response, 'usage', None)
             if usage:
                 total_tokens = getattr(usage, 'total_tokens', 0)
-                self._track_usage(
-                    service_type=ServiceType.EMBEDDING,
+                await self._track_usage(
+                    service_type="embedding",
                     operation="create_text_embeddings",
                     input_tokens=total_tokens,
                     output_tokens=0,