isa-model 0.1.0__py3-none-any.whl

isa_model/inference/services/llm/openai_llm_service.py

@@ -0,0 +1,129 @@
+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
+from isa_model.inference.backends.openai_client import OpenAIBackendClient
+
+logger = logging.getLogger(__name__)
+
+class OpenAILLMService(BaseLLMService):
+    """OpenAI LLM service implementation"""
+    
+    def __init__(self, provider: 'BaseProvider', model_name: str = "gpt-3.5-turbo", backend: Optional[OpenAIBackendClient] = None):
+        super().__init__(provider, model_name)
+        
+        # Use provided backend or create new one
+        if backend:
+            self.backend = backend
+        else:
+            api_key = self.config.get("api_key", "")
+            api_base = self.config.get("api_base", "https://api.openai.com/v1")
+            self.backend = OpenAIBackendClient(api_key, api_base)
+            
+        self.last_token_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
+        logger.info(f"Initialized OpenAILLMService 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:
+            temperature = self.config.get("temperature", 0.7)
+            max_tokens = self.config.get("max_tokens", 1024)
+            
+            payload = {
+                "model": self.model_name,
+                "messages": messages,
+                "temperature": temperature,
+                "max_tokens": max_tokens
+            }
+            response = await self.backend.post("/chat/completions", payload)
+            
+            # Update token usage
+            self.last_token_usage = response.get("usage", {
+                "prompt_tokens": 0,
+                "completion_tokens": 0,
+                "total_tokens": 0
+            })
+            
+            return response["choices"][0]["message"]["content"]
+            
+        except Exception as e:
+            logger.error(f"Error in chat completion: {e}")
+            raise
+    
+    async def acompletion(self, prompt: str):
+        """Text completion method (using chat API since completions is deprecated)"""
+        try:
+            messages = [{"role": "user", "content": prompt}]
+            return await self.achat(messages)
+        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"""
+        try:
+            temperature = self.config.get("temperature", 0.7)
+            max_tokens = self.config.get("max_tokens", 1024)
+            
+            payload = {
+                "model": self.model_name,
+                "messages": messages,
+                "temperature": temperature,
+                "max_tokens": max_tokens,
+                "n": n
+            }
+            response = await self.backend.post("/chat/completions", payload)
+            
+            # Update token usage
+            self.last_token_usage = response.get("usage", {
+                "prompt_tokens": 0,
+                "completion_tokens": 0,
+                "total_tokens": 0
+            })
+            
+            return [choice["message"]["content"] 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)
+            
+            payload = {
+                "model": self.model_name,
+                "messages": messages,
+                "temperature": temperature,
+                "max_tokens": max_tokens,
+                "stream": True
+            }
+            
+            async for chunk in self.backend.stream_chat(payload):
+                yield chunk
+                
+        except Exception as e:
+            logger.error(f"Error in stream chat: {e}")
+            raise
+    
+    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/replicate_llm_service.py

@@ -0,0 +1,179 @@
+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
+from isa_model.inference.backends.replicate_client import ReplicateBackendClient
+
+logger = logging.getLogger(__name__)
+
+class ReplicateLLMService(BaseLLMService):
+    """Replicate LLM service implementation"""
+    
+    def __init__(self, provider: 'BaseProvider', model_name: str = "meta/llama-3-8b-instruct", backend: Optional[ReplicateBackendClient] = None):
+        super().__init__(provider, model_name)
+        
+        # Use provided backend or create new one
+        if backend:
+            self.backend = backend
+        else:
+            api_token = self.config.get("api_token", "")
+            self.backend = ReplicateBackendClient(api_token)
+            
+        # Parse model name for Replicate format (owner/model)
+        if "/" not in model_name:
+            logger.warning(f"Model name {model_name} is not in Replicate format (owner/model). Using as-is.")
+        
+        # Store version separately if provided
+        self.model_version = None
+        if ":" in model_name:
+            self.model_name, self.model_version = model_name.split(":", 1)
+            
+        self.last_token_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
+        logger.info(f"Initialized ReplicateLLMService 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:
+            temperature = self.config.get("temperature", 0.7)
+            max_tokens = self.config.get("max_tokens", 1024)
+            
+            # Convert to Replicate format
+            prompt = self._convert_messages_to_prompt(messages)
+            
+            # Prepare input data
+            input_data = {
+                "prompt": prompt,
+                "temperature": temperature,
+                "max_new_tokens": max_tokens,
+                "system_prompt": self._extract_system_prompt(messages)
+            }
+            
+            # Call Replicate API
+            prediction = await self.backend.create_prediction(
+                self.model_name, 
+                self.model_version, 
+                input_data
+            )
+            
+            # Get output - could be a list of strings or a single string
+            output = prediction.get("output", "")
+            if isinstance(output, list):
+                output = "".join(output)
+            
+            # Approximate token usage - Replicate doesn't provide token counts
+            approx_prompt_tokens = len(prompt) // 4  # Very rough approximation
+            approx_completion_tokens = len(output) // 4
+            self.last_token_usage = {
+                "prompt_tokens": approx_prompt_tokens,
+                "completion_tokens": approx_completion_tokens,
+                "total_tokens": approx_prompt_tokens + approx_completion_tokens
+            }
+            
+            return output
+            
+        except Exception as e:
+            logger.error(f"Error in chat completion: {e}")
+            raise
+    
+    def _convert_messages_to_prompt(self, messages: List[Dict[str, str]]) -> str:
+        """Convert chat messages to Replicate prompt format"""
+        prompt = ""
+        
+        for msg in messages:
+            role = msg.get("role", "user")
+            content = msg.get("content", "")
+            
+            # Skip system prompts - handled separately
+            if role == "system":
+                continue
+                
+            if role == "user":
+                prompt += f"Human: {content}\n\n"
+            elif role == "assistant":
+                prompt += f"Assistant: {content}\n\n"
+            else:
+                # Default to user for unknown roles
+                prompt += f"Human: {content}\n\n"
+        
+        # Add final assistant prefix for the model to continue
+        prompt += "Assistant: "
+        
+        return prompt
+    
+    def _extract_system_prompt(self, messages: List[Dict[str, str]]) -> str:
+        """Extract system prompt from messages"""
+        for msg in messages:
+            if msg.get("role") == "system":
+                return msg.get("content", "")
+        return ""
+    
+    async def acompletion(self, prompt: str):
+        """Text completion method"""
+        try:
+            # For simple completion, use chat format with a single user message
+            messages = [{"role": "user", "content": prompt}]
+            return await self.achat(messages)
+        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"""
+        # Replicate doesn't support multiple outputs in one call,
+        # so we make multiple calls
+        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"""
+        try:
+            temperature = self.config.get("temperature", 0.7)
+            max_tokens = self.config.get("max_tokens", 1024)
+            
+            # Convert to Replicate format
+            prompt = self._convert_messages_to_prompt(messages)
+            
+            # Prepare input data
+            input_data = {
+                "prompt": prompt,
+                "temperature": temperature,
+                "max_new_tokens": max_tokens,
+                "system_prompt": self._extract_system_prompt(messages),
+                "stream": True
+            }
+            
+            # Call Replicate API with streaming
+            async for chunk in self.backend.stream_prediction(
+                self.model_name, 
+                self.model_version, 
+                input_data
+            ):
+                yield chunk
+                
+        except Exception as e:
+            logger.error(f"Error in stream chat: {e}")
+            raise
+    
+    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/triton_llm_service.py

@@ -0,0 +1,230 @@
+import json
+import logging
+import asyncio
+from typing import Dict, List, Any, AsyncGenerator, Optional, Union
+
+import numpy as np
+
+from isa_model.inference.services.base_service import BaseLLMService
+from isa_model.inference.providers.triton_provider import TritonProvider
+
+logger = logging.getLogger(__name__)
+
+
+class TritonLLMService(BaseLLMService):
+    """
+    LLM service that uses Triton Inference Server to run inference.
+    """
+    
+    def __init__(self, provider: TritonProvider, model_name: str):
+        """
+        Initialize the Triton LLM service.
+        
+        Args:
+            provider: The Triton provider
+            model_name: Name of the model in Triton (e.g., "Llama3-8B")
+        """
+        super().__init__(provider, model_name)
+        self.client = None
+        self.tokenizer = None
+        self.token_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
+        self.last_token_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
+        
+    async def _initialize_client(self):
+        """Initialize the Triton client"""
+        if self.client is None:
+            self.client = self.provider.create_client()
+            
+            # Check if model is ready
+            if not self.provider.is_model_ready(self.model_name):
+                logger.error(f"Model {self.model_name} is not ready on Triton server")
+                raise RuntimeError(f"Model {self.model_name} is not ready on Triton server")
+                
+            logger.info(f"Initialized Triton client for model: {self.model_name}")
+    
+    async def ainvoke(self, prompt: Union[str, List[Dict[str, str]], Any]) -> str:
+        """
+        Universal invocation method.
+        
+        Args:
+            prompt: Text prompt or chat messages
+        
+        Returns:
+            Generated text
+        """
+        if isinstance(prompt, str):
+            return await self.acompletion(prompt)
+        elif isinstance(prompt, list) and all(isinstance(m, dict) for m in prompt):
+            return await self.achat(prompt)
+        else:
+            raise ValueError("Prompt must be either a string or a list of message dictionaries")
+    
+    async def achat(self, messages: List[Dict[str, str]]) -> str:
+        """
+        Chat completion method.
+        
+        Args:
+            messages: List of message dictionaries with 'role' and 'content'
+            
+        Returns:
+            Generated chat response
+        """
+        # Format chat messages into a single prompt
+        formatted_prompt = self._format_chat_messages(messages)
+        return await self.acompletion(formatted_prompt)
+    
+    async def acompletion(self, prompt: str) -> str:
+        """
+        Text completion method.
+        
+        Args:
+            prompt: Text prompt
+            
+        Returns:
+            Generated text completion
+        """
+        await self._initialize_client()
+        
+        try:
+            import tritonclient.http as httpclient
+            
+            # Create input tensors
+            input_text = np.array([prompt], dtype=np.object_)
+            inputs = [httpclient.InferInput("TEXT", input_text.shape, "BYTES")]
+            inputs[0].set_data_from_numpy(input_text)
+            
+            # Default parameters
+            generation_params = {
+                "max_new_tokens": self.config.get("max_new_tokens", 512),
+                "temperature": self.config.get("temperature", 0.7),
+                "top_p": self.config.get("top_p", 0.9),
+                "do_sample": self.config.get("temperature", 0.7) > 0
+            }
+            
+            # Add parameters as input tensor
+            param_json = json.dumps(generation_params)
+            param_data = np.array([param_json], dtype=np.object_)
+            param_input = httpclient.InferInput("PARAMETERS", param_data.shape, "BYTES")
+            param_input.set_data_from_numpy(param_data)
+            inputs.append(param_input)
+            
+            # Create output tensor
+            outputs = [httpclient.InferRequestedOutput("TEXT")]
+            
+            # Send the request
+            response = await asyncio.to_thread(
+                self.client.infer,
+                self.model_name,
+                inputs,
+                outputs=outputs
+            )
+            
+            # Process the response
+            output = response.as_numpy("TEXT")
+            response_text = output[0].decode('utf-8')
+            
+            # Update token usage (estimated since we don't have actual token counts)
+            prompt_tokens = len(prompt) // 4  # Rough estimate
+            completion_tokens = len(response_text) // 4  # Rough estimate
+            total_tokens = prompt_tokens + completion_tokens
+            
+            self.last_token_usage = {
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": completion_tokens,
+                "total_tokens": total_tokens
+            }
+            
+            # Update total token usage
+            self.token_usage["prompt_tokens"] += prompt_tokens
+            self.token_usage["completion_tokens"] += completion_tokens
+            self.token_usage["total_tokens"] += total_tokens
+            
+            return response_text
+            
+        except Exception as e:
+            logger.error(f"Error during Triton inference: {str(e)}")
+            raise
+    
+    async def agenerate(self, messages: List[Dict[str, str]], n: int = 1) -> List[str]:
+        """
+        Generate multiple completions.
+        
+        Args:
+            messages: List of message dictionaries
+            n: Number of completions to generate
+            
+        Returns:
+            List of generated 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.
+        
+        Args:
+            messages: List of message dictionaries
+            
+        Yields:
+            Generated text chunks
+        """
+        # For Triton, we don't have true streaming, so we generate the full response
+        # and then simulate streaming
+        full_response = await self.achat(messages)
+        
+        # Simulate streaming by yielding words
+        words = full_response.split()
+        for i in range(len(words)):
+            chunk = ' '.join(words[:i+1])
+            yield chunk
+            await asyncio.sleep(0.05)  # Small delay to simulate streaming
+    
+    def get_token_usage(self) -> Dict[str, int]:
+        """
+        Get total token usage statistics.
+        
+        Returns:
+            Dictionary with token usage statistics
+        """
+        return self.token_usage
+    
+    def get_last_token_usage(self) -> Dict[str, int]:
+        """
+        Get token usage from last request.
+        
+        Returns:
+            Dictionary with token usage statistics from last request
+        """
+        return self.last_token_usage
+    
+    def _format_chat_messages(self, messages: List[Dict[str, str]]) -> str:
+        """
+        Format chat messages into a single prompt for models that don't support chat natively.
+        
+        Args:
+            messages: List of message dictionaries
+            
+        Returns:
+            Formatted prompt
+        """
+        formatted_prompt = ""
+        
+        for message in messages:
+            role = message.get("role", "user").lower()
+            content = message.get("content", "")
+            
+            if role == "system":
+                formatted_prompt += f"System: {content}\n\n"
+            elif role == "user":
+                formatted_prompt += f"User: {content}\n\n"
+            elif role == "assistant":
+                formatted_prompt += f"Assistant: {content}\n\n"
+            else:
+                formatted_prompt += f"{role.capitalize()}: {content}\n\n"
+        
+        formatted_prompt += "Assistant: "
+        return formatted_prompt 

isa_model/inference/services/others/table_transformer_service.py

@@ -0,0 +1,61 @@
+from typing import Dict, Any, List, Union, Optional
+from ...base_service import BaseService
+from ...base_provider import BaseProvider
+from transformers import TableTransformerForObjectDetection, DetrImageProcessor
+import torch
+from PIL import Image
+import numpy as np
+
+class TableTransformerService(BaseService):
+    """Table detection service using Microsoft's Table Transformer"""
+    
+    def __init__(self, provider: 'BaseProvider', model_name: str = "microsoft/table-transformer-detection"):
+        super().__init__(provider, model_name)
+        self.processor = DetrImageProcessor.from_pretrained(model_name)
+        self.model = TableTransformerForObjectDetection.from_pretrained(model_name)
+        if torch.cuda.is_available():
+            self.model = self.model.cuda()
+    
+    async def detect_tables(self, image_path: str) -> Dict[str, Any]:
+        """Detect tables in image"""
+        try:
+            # Load and process image
+            image = Image.open(image_path)
+            inputs = self.processor(images=image, return_tensors="pt")
+            
+            if torch.cuda.is_available():
+                inputs = {k: v.cuda() for k, v in inputs.items()}
+            
+            # Run inference
+            outputs = self.model(**inputs)
+            
+            # Convert outputs to image size
+            target_sizes = torch.tensor([image.size[::-1]])
+            results = self.processor.post_process_object_detection(
+                outputs, target_sizes=target_sizes, threshold=0.7
+            )[0]
+            
+            tables = []
+            for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+                if label == 1:  # Table class
+                    tables.append({
+                        "confidence": score.item(),
+                        "bbox": box.tolist(),
+                        "type": "table"
+                    })
+            
+            return {
+                "tables": tables,
+                "image_size": image.size
+            }
+            
+        except Exception as e:
+            raise RuntimeError(f"Table detection failed: {e}")
+    
+    async def close(self):
+        """Cleanup resources"""
+        if hasattr(self, 'model'):
+            del self.model
+        if hasattr(self, 'processor'):
+            del self.processor
+        torch.cuda.empty_cache() 

isa_model/inference/services/vision/__init__.py

@@ -0,0 +1,12 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+"""
+Vision服务包
+包含所有视觉相关服务模块
+"""
+
+# 导出ReplicateVisionService
+from isa_model.inference.services.vision.replicate_vision_service import ReplicateVisionService
+
+__all__ = ["ReplicateVisionService"] 

isa_model/inference/services/vision/helpers/image_utils.py

@@ -0,0 +1,58 @@
+from io import BytesIO
+from PIL import Image
+from typing import Union
+import base64
+from app.config.config_manager import config_manager
+
+logger = config_manager.get_logger(__name__)
+
+def compress_image(image_data: Union[bytes, BytesIO], max_size: int = 1024) -> bytes:
+    """压缩图片以减小大小
+    
+    Args:
+        image_data: 图片数据，可以是 bytes 或 BytesIO
+        max_size: 最大尺寸（像素）
+        
+    Returns:
+        bytes: 压缩后的图片数据
+    """
+    try:
+        # 如果输入是 bytes，转换为 BytesIO
+        if isinstance(image_data, bytes):
+            image_data = BytesIO(image_data)
+            
+        img = Image.open(image_data)
+        
+        # 转换为 RGB 模式（如果需要）
+        if img.mode in ('RGBA', 'P'):
+            img = img.convert('RGB')
+        
+        # 计算新尺寸，保持宽高比
+        ratio = max_size / max(img.size)
+        if ratio < 1:
+            new_size = tuple(int(dim * ratio) for dim in img.size)
+            img = img.resize(new_size, Image.Resampling.LANCZOS)
+        
+        # 保存压缩后的图片
+        output = BytesIO()
+        img.save(output, format='JPEG', quality=85, optimize=True)
+        return output.getvalue()
+        
+    except Exception as e:
+        logger.error(f"Error compressing image: {e}")
+        raise
+
+def encode_image_to_base64(image_data: bytes) -> str:
+    """将图片数据编码为 base64 字符串
+    
+    Args:
+        image_data: 图片二进制数据
+        
+    Returns:
+        str: base64 编码的字符串
+    """
+    try:
+        return base64.b64encode(image_data).decode('utf-8')
+    except Exception as e:
+        logger.error(f"Error encoding image to base64: {e}")
+        raise