isa-model 0.0.1__py3-none-any.whl

isa_model/inference/services/audio/openai_tts_service.py

@@ -0,0 +1,71 @@
+from typing import Dict, Any
+import tempfile
+import os
+from openai import AsyncOpenAI
+from tenacity import retry, stop_after_attempt, wait_exponential
+from isa_model.inference.services.base_service import BaseService
+from isa_model.inference.providers.base_provider import BaseProvider
+import logging
+
+logger = logging.getLogger(__name__)
+
+class YYDSAudioService(BaseService):
+    """Audio model service wrapper for YYDS"""
+    
+    def __init__(self, provider: 'BaseProvider', model_name: str):
+        super().__init__(provider, model_name)
+        # 初始化 AsyncOpenAI 客户端
+        self._client = AsyncOpenAI(
+            api_key=self.config.get('api_key'),
+            base_url=self.config.get('base_url')
+        )
+        self.language = self.config.get('language', None)
+    
+    @property
+    def client(self) -> AsyncOpenAI:
+        """获取底层的 OpenAI 客户端"""
+        return self._client
+    
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=4, max=10),
+        reraise=True
+    )
+    async def transcribe(self, audio_data: bytes) -> Dict[str, Any]:
+        """转写音频数据
+        
+        Args:
+            audio_data: 音频二进制数据
+            
+        Returns:
+            Dict[str, Any]: 包含转写文本的字典
+        """
+        try:
+            # 创建临时文件存储音频数据
+            with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_file:
+                temp_file.write(audio_data)
+                temp_file.flush()
+                
+                # 以二进制模式打开文件用于 API 请求
+                with open(temp_file.name, 'rb') as audio_file:
+                    # 只在有效的 ISO-639-1 语言代码时包含 language 参数
+                    params = {
+                        'model': self.model_name,
+                        'file': audio_file,
+                    }
+                    if self.language and isinstance(self.language, str):
+                        params['language'] = self.language
+                        
+                    response = await self._client.audio.transcriptions.create(**params)
+                    
+                # 清理临时文件
+                os.unlink(temp_file.name)
+                
+                # 返回包含转写文本的字典
+                return {
+                    "text": response.text
+                }
+                
+        except Exception as e:
+            logger.error(f"Error in audio transcription: {e}")
+            raise

isa_model/inference/services/base_service.py

@@ -0,0 +1,106 @@
+from abc import ABC, abstractmethod
+from typing import Dict, Any, List, Union, AsyncGenerator, TypeVar, Optional
+from isa_model.inference.providers.base_provider import BaseProvider
+
+T = TypeVar('T')  # Generic type for responses
+
+class BaseService(ABC):
+    """Base class for all AI services"""
+    
+    def __init__(self, provider: 'BaseProvider', model_name: str):
+        self.provider = provider
+        self.model_name = model_name
+        self.config = provider.get_config()
+        
+    def __await__(self):
+        """Make the service awaitable"""
+        yield
+        return self
+
+class BaseLLMService(BaseService):
+    """Base class for LLM services"""
+    
+    @abstractmethod
+    async def ainvoke(self, prompt: Union[str, List[Dict[str, str]], Any]) -> T:
+        """Universal invocation method"""
+        pass
+    
+    @abstractmethod
+    async def achat(self, messages: List[Dict[str, str]]) -> T:
+        """Chat completion method"""
+        pass
+    
+    @abstractmethod
+    async def acompletion(self, prompt: str) -> T:
+        """Text completion method"""
+        pass
+    
+    @abstractmethod
+    async def agenerate(self, messages: List[Dict[str, str]], n: int = 1) -> List[T]:
+        """Generate multiple completions"""
+        pass
+    
+    @abstractmethod
+    async def astream_chat(self, messages: List[Dict[str, str]]) -> AsyncGenerator[str, None]:
+        """Stream chat responses"""
+        pass
+    
+    @abstractmethod
+    def get_token_usage(self) -> Any:
+        """Get total token usage statistics"""
+        pass
+    
+    @abstractmethod
+    def get_last_token_usage(self) -> Dict[str, int]:
+        """Get token usage from last request"""
+        pass
+
+class BaseEmbeddingService(BaseService):
+    """Base class for embedding services"""
+    
+    @abstractmethod
+    async def create_text_embedding(self, text: str) -> List[float]:
+        """Create embedding for single text"""
+        pass
+    
+    @abstractmethod
+    async def create_text_embeddings(self, texts: List[str]) -> List[List[float]]:
+        """Create embeddings for multiple texts"""
+        pass
+    
+    @abstractmethod
+    async def create_chunks(self, text: str, metadata: Optional[Dict] = None) -> List[Dict]:
+        """Create text chunks with embeddings"""
+        pass
+    
+    @abstractmethod
+    async def compute_similarity(self, embedding1: List[float], embedding2: List[float]) -> float:
+        """Compute similarity between two embeddings"""
+        pass
+    
+    @abstractmethod
+    async def close(self):
+        """Cleanup resources"""
+        pass
+
+class BaseRerankService(BaseService):
+    """Base class for reranking services"""
+    
+    @abstractmethod
+    async def rerank(
+        self,
+        query: str,
+        documents: List[Dict],
+        top_k: int = 5
+    ) -> List[Dict]:
+        """Rerank documents based on query relevance"""
+        pass
+    
+    @abstractmethod
+    async def rerank_texts(
+        self,
+        query: str,
+        texts: List[str]
+    ) -> List[Dict]:
+        """Rerank raw texts based on query relevance"""
+        pass

isa_model/inference/services/embedding/ollama_embed_service.py

@@ -0,0 +1,97 @@
+import logging
+import httpx
+import asyncio
+from typing import List, Dict, Any, Optional
+
+# 保留您指定的导入和框架结构
+from isa_model.inference.services.base_service import BaseEmbeddingService
+from isa_model.inference.providers.base_provider import BaseProvider
+
+logger = logging.getLogger(__name__)
+
+class OllamaEmbedService(BaseEmbeddingService):
+    """
+    Ollama embedding service.
+    此类遵循基础服务架构，但使用其自己的 HTTP 客户端与 Ollama API 通信，
+    而不依赖于注入的 backend 对象。
+    """
+    
+    def __init__(self, provider: 'BaseProvider', model_name: str = "bge-m3"):
+        # 保持对基类和 provider 的兼容
+        super().__init__(provider, model_name)
+        
+        # 从基类继承的 self.config 中获取配置
+        host = self.config.get("host", "localhost")
+        port = self.config.get("port", 11434)
+        
+        # 创建并持有自己的 httpx 客户端实例
+        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}")
+    
+    async def create_text_embedding(self, text: str) -> List[float]:
+        """为单个文本创建 embedding"""
+        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"]
+            
+        except httpx.RequestError as e:
+            logger.error(f"An error occurred while requesting {e.request.url!r}: {e}")
+            raise
+        except Exception as e:
+            logger.error(f"Error creating text embedding: {e}")
+            raise
+    
+    async def create_text_embeddings(self, texts: List[str]) -> List[List[float]]:
+        """为多个文本并发地创建 embeddings"""
+        if not texts:
+            return []
+        
+        tasks = [self.create_text_embedding(text) for text in texts]
+        embeddings = await asyncio.gather(*tasks)
+        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)]
+        
+        if not chunk_texts:
+            return []
+
+        embeddings = await self.create_text_embeddings(chunk_texts)
+        
+        chunks = [
+            {
+                "text": chunk_text,
+                "embedding": emb,
+                "metadata": metadata or {}
+            }
+            for chunk_text, emb in zip(chunk_texts, embeddings)
+        ]
+            
+        return chunks
+    
+    async def compute_similarity(self, embedding1: List[float], embedding2: List[float]) -> float:
+        """计算两个嵌入向量之间的余弦相似度"""
+        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
+        
+        if norm1 * norm2 == 0:
+            return 0.0
+            
+        return dot_product / (norm1 * norm2)
+    
+    async def close(self):
+        """关闭内置的 HTTP 客户端"""
+        await self.client.aclose()
+        logger.info("OllamaEmbedService's internal client has been closed.")

isa_model/inference/services/llm/__init__.py

@@ -0,0 +1,12 @@
+"""
+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
+
+__all__ = [
+    "OllamaLLMService",
+    "OpenAILLMService",
+] 

isa_model/inference/services/llm/base_llm_service.py

@@ -0,0 +1,134 @@
+from abc import ABC, abstractmethod
+from typing import Dict, Any, List, Union, Optional, AsyncGenerator, TypeVar
+from isa_model.inference.services.base_service import BaseService
+
+T = TypeVar('T')  # Generic type for responses
+
+class BaseLLMService(BaseService):
+    """Base class for Large Language Model services"""
+    
+    @abstractmethod
+    async def ainvoke(self, prompt: Union[str, List[Dict[str, str]], Any]) -> T:
+        """
+        Universal invocation method that handles different input types
+        
+        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
+        
+        Args:
+            messages: List of message dictionaries with 'role' and 'content' keys
+                     Example: [{"role": "user", "content": "Hello"}]
+            
+        Returns:
+            Chat completion response
+        """
+        pass
+    
+    @abstractmethod
+    async def acompletion(self, prompt: str) -> T:
+        """
+        Text completion method for simple prompt completion
+        
+        Args:
+            prompt: Input text prompt
+            
+        Returns:
+            Text completion response
+        """
+        pass
+    
+    @abstractmethod
+    async def agenerate(self, messages: List[Dict[str, str]], n: int = 1) -> List[T]:
+        """
+        Generate multiple completions for the same input
+        
+        Args:
+            messages: List of message dictionaries
+            n: Number of completions to generate
+            
+        Returns:
+            List of completion responses
+        """
+        pass
+    
+    @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
+    
+    @abstractmethod
+    async def astream_completion(self, prompt: str) -> AsyncGenerator[str, None]:
+        """
+        Stream completion responses token by token
+        
+        Args:
+            prompt: Input text prompt
+            
+        Yields:
+            Individual tokens or chunks of the response
+        """
+        pass
+    
+    @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
+        """
+        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
+        """
+        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
+        """
+        pass
+    
+    @abstractmethod
+    async def close(self):
+        """Cleanup resources and close connections"""
+        pass

isa_model/inference/services/llm/ollama_llm_service.py

@@ -0,0 +1,99 @@
+import logging
+from typing import Dict, Any, List, Union, AsyncGenerator, Optional
+from isa_model.inference.services.base_service import BaseLLMService
+from isa_model.inference.providers.base_provider import BaseProvider
+
+logger = logging.getLogger(__name__)
+
+class OllamaLLMService(BaseLLMService):
+    """Ollama LLM service using backend client"""
+    
+    def __init__(self, provider: 'BaseProvider', model_name: str = "llama3.1"):
+        super().__init__(provider, model_name)
+            
+        self.last_token_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
+        logger.info(f"Initialized OllamaLLMService with model {model_name}")
+    
+    async def ainvoke(self, prompt: Union[str, List[Dict[str, str]], Any]):
+        """Universal invocation method"""
+        if isinstance(prompt, str):
+            return await self.acompletion(prompt)
+        elif isinstance(prompt, list):
+            return await self.achat(prompt)
+        else:
+            raise ValueError("Prompt must be string or list of messages")
+    
+    async def achat(self, messages: List[Dict[str, str]]):
+        """Chat completion method"""
+        try:
+            payload = {
+                "model": self.model_name,
+                "messages": messages,
+                "stream": False
+            }
+            response = await self.backend.post("/api/chat", payload)
+            
+            # Update token usage if available
+            if "eval_count" in response:
+                self.last_token_usage = {
+                    "prompt_tokens": response.get("prompt_eval_count", 0),
+                    "completion_tokens": response.get("eval_count", 0),
+                    "total_tokens": response.get("prompt_eval_count", 0) + response.get("eval_count", 0)
+                }
+            
+            return response["message"]["content"]
+            
+        except Exception as e:
+            logger.error(f"Error in chat completion: {e}")
+            raise
+    
+    async def acompletion(self, prompt: str):
+        """Text completion method"""
+        try:
+            payload = {
+                "model": self.model_name,
+                "prompt": prompt,
+                "stream": False
+            }
+            response = await self.backend.post("/api/generate", payload)
+            
+            # Update token usage if available
+            if "eval_count" in response:
+                self.last_token_usage = {
+                    "prompt_tokens": response.get("prompt_eval_count", 0),
+                    "completion_tokens": response.get("eval_count", 0),
+                    "total_tokens": response.get("prompt_eval_count", 0) + response.get("eval_count", 0)
+                }
+            
+            return response["response"]
+            
+        except Exception as e:
+            logger.error(f"Error in text completion: {e}")
+            raise
+    
+    async def agenerate(self, messages: List[Dict[str, str]], n: int = 1) -> List[str]:
+        """Generate multiple completions"""
+        results = []
+        for _ in range(n):
+            result = await self.achat(messages)
+            results.append(result)
+        return results
+    
+    async def astream_chat(self, messages: List[Dict[str, str]]) -> AsyncGenerator[str, None]:
+        """Stream chat responses"""
+        # Note: This would require modifying the backend to support streaming
+        # For now, return the full response
+        response = await self.achat(messages)
+        yield response
+    
+    def get_token_usage(self):
+        """Get total token usage statistics"""
+        return self.last_token_usage
+    
+    def get_last_token_usage(self) -> Dict[str, int]:
+        """Get token usage from last request"""
+        return self.last_token_usage
+        
+    async def close(self):
+        """Close the backend client"""
+        await self.backend.close() 

isa_model/inference/services/llm/openai_llm_service.py

@@ -0,0 +1,138 @@
+import logging
+import os
+from typing import Dict, Any, List, Union, AsyncGenerator, Optional
+
+# 使用官方 OpenAI 库和 dotenv
+from openai import AsyncOpenAI
+from dotenv import load_dotenv
+
+from isa_model.inference.services.base_service import BaseLLMService
+from isa_model.inference.providers.base_provider import BaseProvider
+
+# 加载 .env.local 文件中的环境变量
+load_dotenv(dotenv_path='.env.local')
+
+logger = logging.getLogger(__name__)
+
+class OpenAILLMService(BaseLLMService):
+    """OpenAI LLM service implementation"""
+    
+    def __init__(self, provider: 'BaseProvider', model_name: str = "gpt-3.5-turbo"):
+        super().__init__(provider, model_name)
+        
+        # 从provider配置初始化 AsyncOpenAI 客户端
+        try:
+            api_key = provider.config.get("api_key") or os.getenv("OPENAI_API_KEY")
+            base_url = provider.config.get("api_base") or os.getenv("OPENAI_API_BASE")
+            
+            self.client = AsyncOpenAI(
+                api_key=api_key,
+                base_url=base_url
+            )
+        except TypeError as e:
+            logger.error("初始化 OpenAI 客户端失败。请检查您的 .env.local 文件中是否正确设置了 OPENAI_API_KEY。")
+            raise ValueError("OPENAI_API_KEY 未设置。") from e
+            
+        self.last_token_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
+        logger.info(f"Initialized OpenAILLMService with model {self.model_name} and endpoint {self.client.base_url}")
+    
+    async def ainvoke(self, prompt: Union[str, List[Dict[str, str]], Any]) -> str:
+        """Universal invocation method"""
+        if isinstance(prompt, str):
+            return await self.acompletion(prompt)
+        elif isinstance(prompt, list):
+            return await self.achat(prompt)
+        else:
+            raise ValueError("Prompt must be a string or a list of messages")
+    
+    async def achat(self, messages: List[Dict[str, str]]) -> str:
+        """Chat completion method"""
+        try:
+            temperature = self.config.get("temperature", 0.7)
+            max_tokens = self.config.get("max_tokens", 1024)
+            
+            response = await self.client.chat.completions.create(
+                model=self.model_name,
+                messages=messages,
+                temperature=temperature,
+                max_tokens=max_tokens
+            )
+            
+            if response.usage:
+                self.last_token_usage = {
+                    "prompt_tokens": response.usage.prompt_tokens,
+                    "completion_tokens": response.usage.completion_tokens,
+                    "total_tokens": response.usage.total_tokens
+                }
+            
+            return response.choices[0].message.content or ""
+            
+        except Exception as e:
+            logger.error(f"Error in chat completion: {e}")
+            raise
+    
+    async def acompletion(self, prompt: str) -> str:
+        """Text completion method (using chat API)"""
+        messages = [{"role": "user", "content": prompt}]
+        return await self.achat(messages)
+    
+    async def agenerate(self, messages: List[Dict[str, str]], n: int = 1) -> List[str]:
+        """Generate multiple completions"""
+        try:
+            temperature = self.config.get("temperature", 0.7)
+            max_tokens = self.config.get("max_tokens", 1024)
+            
+            response = await self.client.chat.completions.create(
+                model=self.model_name,
+                messages=messages,
+                temperature=temperature,
+                max_tokens=max_tokens,
+                n=n
+            )
+            
+            if response.usage:
+                self.last_token_usage = {
+                    "prompt_tokens": response.usage.prompt_tokens,
+                    "completion_tokens": response.usage.completion_tokens,
+                    "total_tokens": response.usage.total_tokens
+                }
+            
+            return [choice.message.content or "" for choice in response.choices]
+        except Exception as e:
+            logger.error(f"Error in generate: {e}")
+            raise
+    
+    async def astream_chat(self, messages: List[Dict[str, str]]) -> AsyncGenerator[str, None]:
+        """Stream chat responses"""
+        try:
+            temperature = self.config.get("temperature", 0.7)
+            max_tokens = self.config.get("max_tokens", 1024)
+            
+            stream = await self.client.chat.completions.create(
+                model=self.model_name,
+                messages=messages,
+                temperature=temperature,
+                max_tokens=max_tokens,
+                stream=True
+            )
+            
+            async for chunk in stream:
+                content = chunk.choices[0].delta.content
+                if content:
+                    yield content
+                
+        except Exception as e:
+            logger.error(f"Error in stream chat: {e}")
+            raise
+    
+    def get_token_usage(self) -> Dict[str, int]:
+        """Get total token usage statistics"""
+        return self.last_token_usage
+    
+    def get_last_token_usage(self) -> Dict[str, int]:
+        """Get token usage from last request"""
+        return self.last_token_usage
+        
+    async def close(self):
+        """Close the backend client"""
+        await self.client.aclose()