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

isa_model/inference/services/base_service.py

@@ -1,6 +1,7 @@
 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
 
 T = TypeVar('T')  # Generic type for responses
 
@@ -10,7 +11,41 @@ class BaseService(ABC):
     def __init__(self, provider: 'BaseProvider', model_name: str):
         self.provider = provider
         self.model_name = model_name
-        self.config = provider.get_config()
+        self.config = provider.get_full_config()
+        
+    def _track_usage(
+        self,
+        service_type: Union[str, ServiceType],
+        operation: str,
+        input_tokens: Optional[int] = None,
+        output_tokens: Optional[int] = None,
+        input_units: Optional[float] = None,
+        output_units: Optional[float] = None,
+        metadata: Optional[Dict[str, Any]] = None
+    ):
+        """Track usage for billing purposes"""
+        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'
+            
+            track_usage(
+                provider=provider_name,
+                service_type=service_type,
+                model_name=self.model_name,
+                operation=operation,
+                input_tokens=input_tokens,
+                output_tokens=output_tokens,
+                input_units=input_units,
+                output_units=output_units,
+                metadata=metadata
+            )
+        except Exception as e:
+            # Don't let billing tracking break the service
+            import logging
+            logging.getLogger(__name__).warning(f"Failed to track usage: {e}")
         
     def __await__(self):
         """Make the service awaitable"""

isa_model/inference/services/embedding/base_embed_service.py

@@ -0,0 +1,112 @@
+from abc import ABC, abstractmethod
+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"""
+    
+    @abstractmethod
+    async def create_text_embedding(self, text: str) -> List[float]:
+        """
+        Create embedding for single text
+        
+        Args:
+            text: Input text to embed
+            
+        Returns:
+            List of float values representing the embedding vector
+        """
+        pass
+    
+    @abstractmethod
+    async def create_text_embeddings(self, texts: List[str]) -> List[List[float]]:
+        """
+        Create embeddings for multiple texts
+        
+        Args:
+            texts: List of input texts to embed
+            
+        Returns:
+            List of embedding vectors, one for each input text
+        """
+        pass
+    
+    @abstractmethod
+    async def create_chunks(self, text: str, metadata: Optional[Dict] = None) -> List[Dict]:
+        """
+        Create text chunks with embeddings
+        
+        Args:
+            text: Input text to chunk and embed
+            metadata: Optional metadata to attach to chunks
+            
+        Returns:
+            List of dictionaries containing:
+            - text: The chunk text
+            - embedding: The embedding vector
+            - metadata: Associated metadata
+            - start_index: Start position in original text
+            - end_index: End position in original text
+        """
+        pass
+    
+    @abstractmethod
+    async def compute_similarity(self, embedding1: List[float], embedding2: List[float]) -> float:
+        """
+        Compute similarity between two embeddings
+        
+        Args:
+            embedding1: First embedding vector
+            embedding2: Second embedding vector
+            
+        Returns:
+            Similarity score (typically cosine similarity, range -1 to 1)
+        """
+        pass
+    
+    @abstractmethod
+    async def find_similar_texts(
+        self, 
+        query_embedding: List[float], 
+        candidate_embeddings: List[List[float]], 
+        top_k: int = 5
+    ) -> List[Dict[str, Any]]:
+        """
+        Find most similar texts based on embeddings
+        
+        Args:
+            query_embedding: Query embedding vector
+            candidate_embeddings: List of candidate embedding vectors
+            top_k: Number of top similar results to return
+            
+        Returns:
+            List of dictionaries containing:
+            - index: Index in candidate_embeddings
+            - similarity: Similarity score
+        """
+        pass
+    
+    @abstractmethod
+    def get_embedding_dimension(self) -> int:
+        """
+        Get the dimension of embeddings produced by this service
+        
+        Returns:
+            Integer dimension of embedding vectors
+        """
+        pass
+    
+    @abstractmethod
+    def get_max_input_length(self) -> int:
+        """
+        Get maximum input text length supported
+        
+        Returns:
+            Maximum number of characters/tokens supported
+        """
+        pass
+    
+    @abstractmethod
+    async def close(self):
+        """Cleanup resources"""
+        pass

isa_model/inference/services/embedding/ollama_embed_service.py

@@ -4,12 +4,12 @@ import asyncio
 from typing import List, Dict, Any, Optional
 
 # 保留您指定的导入和框架结构
-from isa_model.inference.services.base_service import BaseEmbeddingService
+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(BaseEmbeddingService):
+class OllamaEmbedService(BaseEmbedService):
     """
     Ollama embedding service.
     此类遵循基础服务架构，但使用其自己的 HTTP 客户端与 Ollama API 通信，
@@ -91,6 +91,32 @@ class OllamaEmbedService(BaseEmbeddingService):
             
         return dot_product / (norm1 * norm2)
     
+    async def find_similar_texts(
+        self, 
+        query_embedding: List[float], 
+        candidate_embeddings: List[List[float]], 
+        top_k: int = 5
+    ) -> 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})
+        
+        # Sort by similarity in descending order and return top_k
+        similarities.sort(key=lambda x: x["similarity"], reverse=True)
+        return similarities[:top_k]
+    
+    def get_embedding_dimension(self) -> int:
+        """Get the dimension of embeddings produced by this service"""
+        # BGE-M3 produces 1024-dimensional embeddings
+        return 1024
+    
+    def get_max_input_length(self) -> int:
+        """Get maximum input text length supported"""
+        # BGE-M3 supports up to 8192 tokens
+        return 8192
+
     async def close(self):
         """关闭内置的 HTTP 客户端"""
         await self.client.aclose()

isa_model/inference/services/embedding/openai_embed_service.py

@@ -0,0 +1,223 @@
+import logging
+import asyncio
+from typing import List, Dict, Any, Optional
+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__)
+
+class OpenAIEmbedService(BaseEmbedService):
+    """
+    OpenAI embedding service using text-embedding-3-small as default.
+    Provides high-quality embeddings for production use.
+    """
+    
+    def __init__(self, provider: 'BaseProvider', model_name: str = "text-embedding-3-small"):
+        super().__init__(provider, model_name)
+        
+        # Get full configuration from provider (including sensitive data)
+        provider_config = provider.get_full_config()
+        
+        # Initialize AsyncOpenAI client with provider configuration
+        try:
+            if not provider_config.get("api_key"):
+                raise ValueError("OpenAI API key not found in provider configuration")
+            
+            self.client = AsyncOpenAI(
+                api_key=provider_config["api_key"],
+                base_url=provider_config.get("base_url", "https://api.openai.com/v1"),
+                organization=provider_config.get("organization")
+            )
+            
+            logger.info(f"Initialized OpenAIEmbedService with model '{self.model_name}'")
+            
+        except Exception as e:
+            logger.error(f"Failed to initialize OpenAI client: {e}")
+            raise ValueError(f"Failed to initialize OpenAI client. Check your API key configuration: {e}") from e
+        
+        # Model-specific configurations
+        self.dimensions = provider_config.get('dimensions', None)  # Optional dimension reduction
+        self.encoding_format = provider_config.get('encoding_format', 'float')
+    
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=4, max=10),
+        reraise=True
+    )
+    async def create_text_embedding(self, text: str) -> List[float]:
+        """Create embedding for single text"""
+        try:
+            kwargs = {
+                "model": self.model_name,
+                "input": text,
+                "encoding_format": self.encoding_format
+            }
+            
+            # Add dimensions parameter if specified (for text-embedding-3-small/large)
+            if self.dimensions and "text-embedding-3" in self.model_name:
+                kwargs["dimensions"] = self.dimensions
+            
+            response = await self.client.embeddings.create(**kwargs)
+            
+            # Track usage for billing
+            usage = getattr(response, 'usage', None)
+            if usage:
+                total_tokens = getattr(usage, 'total_tokens', 0)
+                self._track_usage(
+                    service_type=ServiceType.EMBEDDING,
+                    operation="create_text_embedding",
+                    input_tokens=total_tokens,
+                    output_tokens=0,
+                    metadata={
+                        "model": self.model_name,
+                        "dimensions": self.dimensions,
+                        "text_length": len(text)
+                    }
+                )
+            
+            return response.data[0].embedding
+            
+        except Exception as e:
+            logger.error(f"Error creating text embedding: {e}")
+            raise
+    
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=4, max=10),
+        reraise=True
+    )
+    async def create_text_embeddings(self, texts: List[str]) -> List[List[float]]:
+        """Create embeddings for multiple texts"""
+        if not texts:
+            return []
+        
+        try:
+            kwargs = {
+                "model": self.model_name,
+                "input": texts,
+                "encoding_format": self.encoding_format
+            }
+            
+            # Add dimensions parameter if specified
+            if self.dimensions and "text-embedding-3" in self.model_name:
+                kwargs["dimensions"] = self.dimensions
+            
+            response = await self.client.embeddings.create(**kwargs)
+            
+            # Track usage for billing
+            usage = getattr(response, 'usage', None)
+            if usage:
+                total_tokens = getattr(usage, 'total_tokens', 0)
+                self._track_usage(
+                    service_type=ServiceType.EMBEDDING,
+                    operation="create_text_embeddings",
+                    input_tokens=total_tokens,
+                    output_tokens=0,
+                    metadata={
+                        "model": self.model_name,
+                        "dimensions": self.dimensions,
+                        "batch_size": len(texts),
+                        "total_text_length": sum(len(t) for t in texts)
+                    }
+                )
+            
+            return [data.embedding for data in response.data]
+            
+        except Exception as e:
+            logger.error(f"Error creating text embeddings: {e}")
+            raise
+    
+    async def create_chunks(self, text: str, metadata: Optional[Dict] = None) -> List[Dict]:
+        """Create text chunks with embeddings"""
+        # Chunk size optimized for OpenAI models (roughly 512 tokens)
+        chunk_size = 400  # words
+        overlap = 50      # word overlap between chunks
+        
+        words = text.split()
+        if not words:
+            return []
+        
+        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,
+                "start_index": i,
+                "end_index": min(i + chunk_size, len(words)),
+                "metadata": metadata or {}
+            })
+        
+        # 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 = 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
+            
+        return dot_product / (norm1 * norm2)
+    
+    async def find_similar_texts(
+        self, 
+        query_embedding: List[float], 
+        candidate_embeddings: List[List[float]], 
+        top_k: int = 5
+    ) -> 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
+            })
+        
+        # Sort by similarity in descending order and return top_k
+        similarities.sort(key=lambda x: x["similarity"], reverse=True)
+        return similarities[:top_k]
+    
+    def get_embedding_dimension(self) -> int:
+        """Get the dimension of embeddings produced by this service"""
+        if self.dimensions:
+            return self.dimensions
+        
+        # Default dimensions for OpenAI models
+        model_dimensions = {
+            "text-embedding-3-small": 1536,
+            "text-embedding-3-large": 3072,
+            "text-embedding-ada-002": 1536
+        }
+        
+        return model_dimensions.get(self.model_name, 1536)
+    
+    def get_max_input_length(self) -> int:
+        """Get maximum input text length supported"""
+        # OpenAI embedding models support up to 8192 tokens
+        return 8192
+    
+    async def close(self):
+        """Cleanup resources"""
+        await self.client.close()
+        logger.info("OpenAIEmbedService client has been closed.")

isa_model/inference/services/llm/__init__.py

@@ -5,8 +5,10 @@ LLM Services - Business logic services for Language Models
 # Import LLM services here when created
 from .ollama_llm_service import OllamaLLMService
 from .openai_llm_service import OpenAILLMService
+from .triton_llm_service import TritonLLMService
 
 __all__ = [
     "OllamaLLMService",
     "OpenAILLMService",
+    "TritonLLMService"
 ]