isa-model 0.3.0__py3-none-any.whl → 0.3.2__py3-none-any.whl

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/base_llm_service.py

@@ -1,244 +1,140 @@
 from abc import ABC, abstractmethod
-from typing import Dict, Any, List, Union, Optional, AsyncGenerator, TypeVar, Callable
+from typing import Dict, Any, List, Union, Optional, AsyncGenerator, Callable
 from isa_model.inference.services.base_service import BaseService
-
-T = TypeVar('T')  # Generic type for responses
+from isa_model.inference.services.llm.llm_adapter import AdapterManager
 
 class BaseLLMService(BaseService):
-    """Base class for Large Language Model services"""
+    """Base class for Large Language Model services with unified invoke interface"""
     
-    @abstractmethod
-    async def ainvoke(self, prompt: Union[str, List[Dict[str, str]], Any]) -> T:
-        """
-        Universal invocation method that handles different input types
+    def __init__(self, provider, model_name: str):
+        super().__init__(provider, model_name)
+        self._bound_tools: List[Any] = []  # 改为存储原始工具对象
+        self._tool_mappings: Dict[str, tuple] = {}  # 工具名到(工具, 适配器)的映射
         
-        Args:
-            prompt: Can be a string, list of messages, or other format
-            
-        Returns:
-            Model response in the appropriate format
-        """
-        pass
-    
-    @abstractmethod
-    async def achat(self, messages: List[Dict[str, str]]) -> T:
-        """
-        Chat completion method using message format
+        # 初始化适配器管理器
+        self.adapter_manager = AdapterManager()
         
-        Args:
-            messages: List of message dictionaries with 'role' and 'content' keys
-                     Example: [{"role": "user", "content": "Hello"}]
-            
-        Returns:
-            Chat completion response
-        """
-        pass
+        # Get streaming config from provider config
+        self.streaming = self.config.get("streaming", False)
     
-    @abstractmethod
-    async def acompletion(self, prompt: str) -> T:
+    def bind_tools(self, tools: List[Any], **kwargs) -> 'BaseLLMService':
         """
-        Text completion method for simple prompt completion
+        Bind tools to this LLM service for function calling
         
         Args:
-            prompt: Input text prompt
+            tools: List of tools to bind (functions, LangChain tools, etc.)
+            **kwargs: Additional tool binding parameters
             
         Returns:
-            Text completion response
+            Self for method chaining
         """
-        pass
+        self._bound_tools = tools
+        return self
     
-    @abstractmethod
-    async def agenerate(self, messages: List[Dict[str, str]], n: int = 1) -> List[T]:
-        """
-        Generate multiple completions for the same input
+    async def _prepare_tools_for_request(self) -> List[Dict[str, Any]]:
+        """准备工具用于请求"""
+        if not self._bound_tools:
+            return []
         
-        Args:
-            messages: List of message dictionaries
-            n: Number of completions to generate
-            
-        Returns:
-            List of completion responses
-        """
-        pass
+        schemas, self._tool_mappings = await self.adapter_manager.convert_tools_to_schemas(self._bound_tools)
+        return schemas
     
-    @abstractmethod
-    async def astream_chat(self, messages: List[Dict[str, str]]) -> AsyncGenerator[str, None]:
-        """
-        Stream chat responses token by token
-        
-        Args:
-            messages: List of message dictionaries
-            
-        Yields:
-            Individual tokens or chunks of the response
-        """
-        pass
+    def _prepare_messages(self, input_data: Union[str, List[Dict[str, str]], Any]) -> List[Dict[str, str]]:
+        """使用适配器管理器转换消息格式"""
+        return self.adapter_manager.convert_messages(input_data)
+    
+    def _format_response(self, response: str, original_input: Any) -> Union[str, Any]:
+        """使用适配器管理器格式化响应"""
+        return self.adapter_manager.format_response(response, original_input)
+    
+    async def _execute_tool_call(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
+        """使用适配器管理器执行工具调用"""
+        return await self.adapter_manager.execute_tool(tool_name, arguments, self._tool_mappings)
     
     @abstractmethod
-    async def astream_completion(self, prompt: str) -> AsyncGenerator[str, None]:
+    async def ainvoke(self, input_data: Union[str, List[Dict[str, str]], Any]) -> Union[str, Any]:
         """
-        Stream completion responses token by token
+        Universal async invocation method that handles different input types
         
         Args:
-            prompt: Input text prompt
+            input_data: Can be:
+                - str: Simple text prompt
+                - list: Message history like [{"role": "user", "content": "hello"}]
+                - Any: LangChain message objects or other formats
             
-        Yields:
-            Individual tokens or chunks of the response
+        Returns:
+            Model response (string for simple cases, object for complex cases)
         """
         pass
     
-    def bind_tools(self, tools: List[Union[Dict[str, Any], Callable]], **kwargs) -> 'BaseLLMService':
+    def invoke(self, input_data: Union[str, List[Dict[str, str]], Any]) -> Union[str, Any]:
         """
-        Bind tools to this LLM service for function calling (LangChain interface)
+        Synchronous wrapper for ainvoke
         
         Args:
-            tools: List of tools to bind. Can be:
-                  - Dictionary with tool schema
-                  - Callable functions (will be converted to schema)
-            **kwargs: Additional tool binding parameters
+            input_data: Same as ainvoke
             
         Returns:
-            A new instance of the service with tools bound
-            
-        Example:
-            def get_weather(location: str) -> str:
-                '''Get weather for a location'''
-                return f"Weather in {location}: Sunny, 25°C"
-            
-            llm_with_tools = llm.bind_tools([get_weather])
-            response = await llm_with_tools.ainvoke("What's the weather in Paris?")
-        """
-        # Create a copy of the current service
-        bound_service = self._create_bound_copy()
-        bound_service._bound_tools = self._convert_tools_to_schema(tools)
-        bound_service._tool_binding_kwargs = kwargs
-        return bound_service
-    
-    def _create_bound_copy(self) -> 'BaseLLMService':
-        """Create a copy of this service for tool binding"""
-        # Default implementation - subclasses should override if needed
-        bound_service = self.__class__(self.provider, self.model_name)
-        bound_service.config = self.config.copy()
-        return bound_service
-    
-    def _convert_tools_to_schema(self, tools: List[Union[Dict[str, Any], Callable]]) -> List[Dict[str, Any]]:
-        """Convert tools to OpenAI function calling schema"""
-        schemas = []
-        for tool in tools:
-            if callable(tool):
-                schema = self._function_to_schema(tool)
-            elif isinstance(tool, dict):
-                schema = tool
-            else:
-                raise ValueError(f"Tool must be callable or dict, got {type(tool)}")
-            schemas.append(schema)
-        return schemas
-    
-    def _function_to_schema(self, func: Callable) -> Dict[str, Any]:
-        """Convert a Python function to OpenAI function schema"""
-        import inspect
-        import json
-        from typing import get_type_hints
-        
-        sig = inspect.signature(func)
-        type_hints = get_type_hints(func)
-        
-        properties = {}
-        required = []
-        
-        for param_name, param in sig.parameters.items():
-            param_type = type_hints.get(param_name, str)
-            
-            # Convert Python types to JSON schema types
-            if param_type == str:
-                prop_type = "string"
-            elif param_type == int:
-                prop_type = "integer"
-            elif param_type == float:
-                prop_type = "number"
-            elif param_type == bool:
-                prop_type = "boolean"
-            elif param_type == list:
-                prop_type = "array"
-            elif param_type == dict:
-                prop_type = "object"
-            else:
-                prop_type = "string"  # Default fallback
-            
-            properties[param_name] = {"type": prop_type}
-            
-            # Add parameter to required if it has no default value
-            if param.default == inspect.Parameter.empty:
-                required.append(param_name)
-        
-        return {
-            "type": "function",
-            "function": {
-                "name": func.__name__,
-                "description": func.__doc__ or f"Function {func.__name__}",
-                "parameters": {
-                    "type": "object",
-                    "properties": properties,
-                    "required": required
-                }
-            }
-        }
+            Model response
+        """
+        import asyncio
+        try:
+            # Try to get current event loop
+            loop = asyncio.get_running_loop()
+            # If we're in an event loop, create a new thread
+            import concurrent.futures
+            with concurrent.futures.ThreadPoolExecutor() as executor:
+                future = executor.submit(asyncio.run, self.ainvoke(input_data))
+                return future.result()
+        except RuntimeError:
+            # No event loop running, create a new one
+            return asyncio.run(self.ainvoke(input_data))
     
     def _has_bound_tools(self) -> bool:
         """Check if this service has bound tools"""
-        return hasattr(self, '_bound_tools') and self._bound_tools
+        return bool(self._bound_tools)
     
-    def _get_bound_tools(self) -> List[Dict[str, Any]]:
-        """Get the bound tools schema"""
-        return getattr(self, '_bound_tools', [])
-    
-    def _execute_tool_call(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
-        """Execute a tool call by name with arguments"""
-        # This is a placeholder - subclasses should implement actual tool execution
-        raise NotImplementedError("Tool execution not implemented in base class")
+    def _get_bound_tools(self) -> List[Any]:
+        """Get the bound tools"""
+        return self._bound_tools
     
     @abstractmethod
     def get_token_usage(self) -> Dict[str, Any]:
-        """
-        Get cumulative token usage statistics for this service instance
-        
-        Returns:
-            Dict containing token usage information:
-            - total_tokens: Total tokens used
-            - prompt_tokens: Tokens used for prompts
-            - completion_tokens: Tokens used for completions
-            - requests_count: Number of requests made
-        """
+        """Get cumulative token usage statistics"""
         pass
     
     @abstractmethod
     def get_last_token_usage(self) -> Dict[str, int]:
-        """
-        Get token usage from the last request
-        
-        Returns:
-            Dict containing last request token usage:
-            - prompt_tokens: Tokens in last prompt
-            - completion_tokens: Tokens in last completion
-            - total_tokens: Total tokens in last request
-        """
+        """Get token usage from the last request"""
         pass
     
     @abstractmethod
     def get_model_info(self) -> Dict[str, Any]:
-        """
-        Get information about the current model
-        
-        Returns:
-            Dict containing model information:
-            - name: Model name
-            - max_tokens: Maximum context length
-            - supports_streaming: Whether streaming is supported
-            - supports_functions: Whether function calling is supported
-        """
+        """Get information about the current model"""
         pass
     
     @abstractmethod
     async def close(self):
         """Cleanup resources and close connections"""
         pass
+    
+    def get_last_usage_with_cost(self) -> Dict[str, Any]:
+        """Get last request usage with cost information"""
+        usage = self.get_last_token_usage()
+        
+        # Calculate cost using provider
+        if hasattr(self.provider, 'calculate_cost'):
+            cost = getattr(self.provider, 'calculate_cost')(
+                self.model_name,
+                usage["prompt_tokens"],
+                usage["completion_tokens"]
+            )
+        else:
+            cost = 0.0
+        
+        return {
+            **usage,
+            "cost_usd": cost,
+            "model": self.model_name,
+            "provider": getattr(self.provider, 'name', 'unknown')
+        }