isa-model 0.3.91__py3-none-any.whl → 0.4.3__py3-none-any.whl

isa_model/inference/services/audio/openai_stt_service.py

@@ -1,4 +1,5 @@
 import logging
+import os
 import aiohttp
 from typing import Dict, Any, List, Union, Optional, BinaryIO
 from openai import AsyncOpenAI
@@ -46,7 +47,7 @@ class OpenAISTTService(BaseSTTService):
         wait=wait_exponential(multiplier=1, min=4, max=10),
         reraise=True
     )
-    async def transcribe(self, audio_file: Union[str, BinaryIO], language: Optional[str] = None, prompt: Optional[str] = None) -> Dict[str, Any]:
+    async def transcribe(self, audio_file: Union[str, BinaryIO, bytes], language: Optional[str] = None, prompt: Optional[str] = None, **kwargs) -> Dict[str, Any]:
         """
         Transcribe audio file to text using OpenAI's Whisper model.
         
@@ -72,13 +73,54 @@ class OpenAISTTService(BaseSTTService):
             if prompt:
                 transcription_params["prompt"] = prompt
             
-            # Handle file input
-            if isinstance(audio_file, str):
-                with open(audio_file, "rb") as f:
-                    transcription = await self.client.audio.transcriptions.create(
-                        file=f,
-                        **transcription_params
-                    )
+            # Handle file input - support bytes, base64 strings, file paths, and file objects
+            if isinstance(audio_file, bytes):
+                # Handle bytes data directly
+                logger.info(f"Processing bytes audio data ({len(audio_file)} bytes)")
+                from io import BytesIO
+                audio_buffer = BytesIO(audio_file)
+                
+                # Use filename from kwargs if provided, otherwise default to .mp3
+                filename = kwargs.get('filename', 'audio.mp3')
+                if filename and not filename.endswith(('.mp3', '.wav', '.m4a', '.flac', '.ogg', '.webm', '.mp4')):
+                    filename += '.mp3'  # Add extension if missing
+                audio_buffer.name = filename
+                logger.info(f"Using filename: {filename}")
+                transcription = await self.client.audio.transcriptions.create(
+                    file=audio_buffer,
+                    **transcription_params
+                )
+            elif isinstance(audio_file, str):
+                # Check if it's a base64 string or file path
+                if len(audio_file) > 100 and not os.path.exists(audio_file):
+                    # Likely a base64 string
+                    try:
+                        import base64
+                        from io import BytesIO
+                        logger.info(f"Attempting to decode base64 audio data (length: {len(audio_file)})")
+                        audio_data = base64.b64decode(audio_file)
+                        audio_buffer = BytesIO(audio_data)
+                        audio_buffer.name = "audio.wav"  # OpenAI needs a filename hint
+                        logger.info(f"Successfully decoded base64 to {len(audio_data)} bytes")
+                        transcription = await self.client.audio.transcriptions.create(
+                            file=audio_buffer,
+                            **transcription_params
+                        )
+                    except Exception as e:
+                        # If base64 decoding fails, treat as file path
+                        logger.error(f"Base64 decoding failed: {e}, treating as file path")
+                        with open(audio_file, "rb") as f:
+                            transcription = await self.client.audio.transcriptions.create(
+                                file=f,
+                                **transcription_params
+                            )
+                else:
+                    # Regular file path
+                    with open(audio_file, "rb") as f:
+                        transcription = await self.client.audio.transcriptions.create(
+                            file=f,
+                            **transcription_params
+                        )
             else:
                 transcription = await self.client.audio.transcriptions.create(
                     file=audio_file,
@@ -121,7 +163,7 @@ class OpenAISTTService(BaseSTTService):
         wait=wait_exponential(multiplier=1, min=4, max=10),
         reraise=True
     )
-    async def translate(self, audio_file: Union[str, BinaryIO]) -> Dict[str, Any]:
+    async def translate(self, audio_file: Union[str, BinaryIO, bytes]) -> Dict[str, Any]:
         """
         Translate audio file to English text using OpenAI's Whisper model.
         
@@ -185,7 +227,7 @@ class OpenAISTTService(BaseSTTService):
             logger.error(f"Translation failed: {e}")
             raise
 
-    async def transcribe_batch(self, audio_files: List[Union[str, BinaryIO]], language: Optional[str] = None, prompt: Optional[str] = None) -> List[Dict[str, Any]]:
+    async def transcribe_batch(self, audio_files: List[Union[str, BinaryIO, bytes]], language: Optional[str] = None, prompt: Optional[str] = None) -> List[Dict[str, Any]]:
         """
         Transcribe multiple audio files in batch.
         
@@ -212,7 +254,7 @@ class OpenAISTTService(BaseSTTService):
         
         return results
 
-    async def detect_language(self, audio_file: Union[str, BinaryIO]) -> Dict[str, Any]:
+    async def detect_language(self, audio_file: Union[str, BinaryIO, bytes]) -> Dict[str, Any]:
         """
         Detect the language of an audio file.
         

isa_model/inference/services/base_service.py

@@ -66,7 +66,8 @@ class BaseService(ABC):
                     output_tokens=output_tokens
                 )
             
-            # Track usage through model manager
+            # Track usage through both systems (legacy and new)
+            # Legacy detailed tracking (will be phased out)
             self.model_manager.billing_tracker.track_model_usage(
                 model_id=self.model_name,
                 operation_type="inference",
@@ -80,6 +81,21 @@ class BaseService(ABC):
                 cost_usd=cost_usd,
                 metadata=metadata
             )
+            
+            # New aggregated statistics tracking
+            self.model_manager.statistics_tracker.track_usage(
+                model_id=self.model_name,
+                provider=self.provider_name,
+                service_type=service_type if isinstance(service_type, str) else service_type.value,
+                operation_type="inference",
+                operation=operation,
+                input_tokens=input_tokens,
+                output_tokens=output_tokens,
+                input_units=input_units,
+                output_units=output_units,
+                cost_usd=cost_usd or 0.0,
+                metadata=metadata
+            )
         except Exception as e:
             # Don't let billing tracking break the service
             import logging

isa_model/inference/services/custom_model_manager.py

@@ -0,0 +1,277 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+"""
+Custom Model Manager - Handles registration and management of custom trained models
+Provides integration for models trained through ISA Model training pipeline
+"""
+
+import logging
+import json
+import os
+from typing import Dict, List, Any, Optional
+from datetime import datetime
+from dataclasses import dataclass, asdict
+
+logger = logging.getLogger(__name__)
+
+@dataclass
+class CustomModelInfo:
+    """Information about a custom model"""
+    model_id: str
+    model_name: str
+    model_type: str  # 'text', 'vision', 'audio', etc.
+    provider: str
+    base_model: str  # The base model this was fine-tuned from
+    training_date: str
+    model_path: str  # Local path or HuggingFace repo
+    metadata: Dict[str, Any]
+    capabilities: List[str]
+    performance_metrics: Optional[Dict[str, float]] = None
+    deployment_config: Optional[Dict[str, Any]] = None
+    
+    def to_dict(self) -> Dict[str, Any]:
+        return asdict(self)
+
+class CustomModelManager:
+    """
+    Manages custom trained models in the ISA Model ecosystem
+    Handles registration, discovery, and integration of custom models
+    """
+    
+    def __init__(self, models_registry_path: str = None):
+        self.models_registry_path = models_registry_path or os.path.join(
+            os.path.expanduser("~"), ".isa_model", "custom_models.json"
+        )
+        self._models: Dict[str, CustomModelInfo] = {}
+        self._load_models_registry()
+    
+    def _load_models_registry(self):
+        """Load custom models registry from file"""
+        if os.path.exists(self.models_registry_path):
+            try:
+                with open(self.models_registry_path, 'r', encoding='utf-8') as f:
+                    models_data = json.load(f)
+                
+                for model_data in models_data.get('models', []):
+                    model_info = CustomModelInfo(**model_data)
+                    self._models[model_info.model_id] = model_info
+                
+                logger.info(f"Loaded {len(self._models)} custom models from registry")
+            except Exception as e:
+                logger.warning(f"Failed to load models registry: {e}")
+                self._models = {}
+        else:
+            # Create default registry with some ISA models
+            self._create_default_registry()
+    
+    def _create_default_registry(self):
+        """Create default registry with ISA models"""
+        default_models = [
+            CustomModelInfo(
+                model_id="isa-llm-service",
+                model_name="ISA LLM Service",
+                model_type="text",
+                provider="isa",
+                base_model="DialoGPT-small",
+                training_date="2024-12-19",
+                model_path="modal://isa-llm-inference",
+                metadata={
+                    "description": "ISA custom LLM service with fallback support",
+                    "parameters": "124M",
+                    "context_length": 1024,
+                    "languages": ["en", "zh"]
+                },
+                capabilities=["chat", "text_generation", "conversation"],
+                performance_metrics={
+                    "perplexity": 3.2,
+                    "bleu_score": 0.75,
+                    "response_time_ms": 850
+                },
+                deployment_config={
+                    "platform": "modal",
+                    "gpu_type": "A10G",
+                    "memory_gb": 16,
+                    "concurrent_requests": 5
+                }
+            ),
+            CustomModelInfo(
+                model_id="xenodennis/dialoGPT-small-20241219-v1",
+                model_name="ISA Fine-tuned DialoGPT",
+                model_type="text",
+                provider="huggingface",
+                base_model="microsoft/DialoGPT-small",
+                training_date="2024-12-19",
+                model_path="xenodennis/dialoGPT-small-20241219-v1",
+                metadata={
+                    "description": "DialoGPT model fine-tuned with ISA training pipeline",
+                    "parameters": "124M",
+                    "trainable_parameters": "294K (LoRA)",
+                    "training_steps": 1000,
+                    "languages": ["en", "zh"]
+                },
+                capabilities=["chat", "text_generation", "dialogue"],
+                performance_metrics={
+                    "final_loss": 2.1234,
+                    "eval_loss": 2.3456,
+                    "training_time_minutes": 15
+                }
+            ),
+            CustomModelInfo(
+                model_id="isa-custom-embeddings",
+                model_name="ISA Custom Embeddings",
+                model_type="embedding",
+                provider="isa",
+                base_model="sentence-transformers/all-MiniLM-L6-v2",
+                training_date="2024-12-19",
+                model_path="local://models/isa-embeddings",
+                metadata={
+                    "description": "Custom embeddings trained on ISA domain data",
+                    "dimensions": 384,
+                    "max_sequence_length": 512
+                },
+                capabilities=["embed", "similarity", "clustering"]
+            )
+        ]
+        
+        for model in default_models:
+            self._models[model.model_id] = model
+        
+        self._save_models_registry()
+        logger.info(f"Created default registry with {len(default_models)} models")
+    
+    def _save_models_registry(self):
+        """Save models registry to file"""
+        try:
+            os.makedirs(os.path.dirname(self.models_registry_path), exist_ok=True)
+            
+            registry_data = {
+                "version": "1.0",
+                "last_updated": datetime.now().isoformat(),
+                "models": [model.to_dict() for model in self._models.values()]
+            }
+            
+            with open(self.models_registry_path, 'w', encoding='utf-8') as f:
+                json.dump(registry_data, f, indent=2, ensure_ascii=False)
+            
+            logger.debug(f"Saved models registry to {self.models_registry_path}")
+        except Exception as e:
+            logger.error(f"Failed to save models registry: {e}")
+    
+    def register_model(self, model_info: CustomModelInfo) -> bool:
+        """Register a new custom model"""
+        try:
+            self._models[model_info.model_id] = model_info
+            self._save_models_registry()
+            logger.info(f"Registered custom model: {model_info.model_id}")
+            return True
+        except Exception as e:
+            logger.error(f"Failed to register model {model_info.model_id}: {e}")
+            return False
+    
+    def unregister_model(self, model_id: str) -> bool:
+        """Unregister a custom model"""
+        if model_id in self._models:
+            del self._models[model_id]
+            self._save_models_registry()
+            logger.info(f"Unregistered custom model: {model_id}")
+            return True
+        return False
+    
+    def get_model(self, model_id: str) -> Optional[CustomModelInfo]:
+        """Get custom model information"""
+        return self._models.get(model_id)
+    
+    def list_models(self, model_type: str = None, provider: str = None) -> List[CustomModelInfo]:
+        """List custom models with optional filtering"""
+        models = list(self._models.values())
+        
+        if model_type:
+            models = [m for m in models if m.model_type == model_type]
+        
+        if provider:
+            models = [m for m in models if m.provider == provider]
+        
+        return models
+    
+    def get_models_for_api(self) -> List[Dict[str, Any]]:
+        """Get models in API format for model listing"""
+        api_models = []
+        
+        for model in self._models.values():
+            api_model = {
+                "model_id": model.model_id,
+                "service_type": model.model_type,
+                "provider": model.provider,
+                "description": model.metadata.get("description", ""),
+                "capabilities": model.capabilities,
+                "custom": True,
+                "base_model": model.base_model,
+                "training_date": model.training_date
+            }
+            
+            # Add performance metrics if available
+            if model.performance_metrics:
+                api_model["performance"] = model.performance_metrics
+            
+            api_models.append(api_model)
+        
+        return api_models
+    
+    def search_models(self, query: str) -> List[CustomModelInfo]:
+        """Search custom models by query"""
+        query_lower = query.lower()
+        matching_models = []
+        
+        for model in self._models.values():
+            # Search in model_id, name, description, and capabilities
+            searchable_text = f"{model.model_id} {model.model_name} {model.metadata.get('description', '')} {' '.join(model.capabilities)}".lower()
+            
+            if query_lower in searchable_text:
+                matching_models.append(model)
+        
+        return matching_models
+    
+    def get_deployment_config(self, model_id: str) -> Optional[Dict[str, Any]]:
+        """Get deployment configuration for a model"""
+        model = self.get_model(model_id)
+        return model.deployment_config if model else None
+    
+    def update_performance_metrics(self, model_id: str, metrics: Dict[str, float]) -> bool:
+        """Update performance metrics for a model"""
+        model = self.get_model(model_id)
+        if model:
+            model.performance_metrics = metrics
+            self._save_models_registry()
+            return True
+        return False
+    
+    def get_provider_models(self, provider: str) -> List[CustomModelInfo]:
+        """Get all models for a specific provider"""
+        return [model for model in self._models.values() if model.provider == provider]
+    
+    def get_stats(self) -> Dict[str, Any]:
+        """Get statistics about custom models"""
+        models_by_type = {}
+        models_by_provider = {}
+        
+        for model in self._models.values():
+            models_by_type[model.model_type] = models_by_type.get(model.model_type, 0) + 1
+            models_by_provider[model.provider] = models_by_provider.get(model.provider, 0) + 1
+        
+        return {
+            "total_models": len(self._models),
+            "models_by_type": models_by_type,
+            "models_by_provider": models_by_provider,
+            "registry_path": self.models_registry_path
+        }
+
+# Global instance
+_custom_model_manager = None
+
+def get_custom_model_manager() -> CustomModelManager:
+    """Get the global custom model manager instance"""
+    global _custom_model_manager
+    if _custom_model_manager is None:
+        _custom_model_manager = CustomModelManager()
+    return _custom_model_manager

isa_model/inference/services/embedding/__init__.py

@@ -0,0 +1,13 @@
+"""
+Embedding Services - Text and Document Embedding Services
+"""
+
+from .base_embed_service import BaseEmbedService
+from .openai_embed_service import OpenAIEmbedService
+from .ollama_embed_service import OllamaEmbedService
+
+__all__ = [
+    'BaseEmbedService',
+    'OpenAIEmbedService',
+    'OllamaEmbedService'
+]

isa_model/inference/services/embedding/base_embed_service.py

@@ -36,17 +36,29 @@ class BaseEmbedService(BaseService):
             if not isinstance(input_data, list):
                 input_data = [input_data]
             return await self.create_text_embeddings(input_data)
-        elif task == "chunk_and_embed":
+        elif task in ["chunk", "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"))
+                raise ValueError("chunk task requires single text input")
+            return await self.create_chunks(input_data, **kwargs)
         elif task == "similarity":
+            # Support both text-based and embedding-based similarity
+            candidates = kwargs.get("candidates")
             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}
+            
+            if candidates:
+                # Text-based similarity - compute embeddings first
+                if isinstance(input_data, list):
+                    raise ValueError("similarity task with candidates requires single query text")
+                # Remove candidates from kwargs to avoid duplicate parameter
+                similarity_kwargs = {k: v for k, v in kwargs.items() if k != 'candidates'}
+                return await self._text_similarity_search(input_data, candidates, **similarity_kwargs)
+            elif embedding1 and embedding2:
+                # Direct embedding similarity
+                similarity = await self.compute_similarity(embedding1, embedding2)
+                return {"similarity": similarity}
+            else:
+                raise ValueError("similarity task requires either 'candidates' parameter or both 'embedding1' and 'embedding2' parameters")
         elif task == "find_similar":
             query_embedding = kwargs.get("query_embedding")
             candidate_embeddings = kwargs.get("candidate_embeddings")
@@ -57,6 +69,21 @@ class BaseEmbedService(BaseService):
                 candidate_embeddings, 
                 kwargs.get("top_k", 5)
             )
+        
+        # ==================== 重排序类任务 ====================
+        elif task in ["rerank", "rerank_documents", "document_ranking"]:
+            query = kwargs.get("query") or input_data
+            documents = kwargs.get("documents")
+            if not documents:
+                raise ValueError("rerank task requires documents parameter")
+            if isinstance(query, list):
+                raise ValueError("rerank task requires single query string")
+            return await self.rerank_documents(
+                query=query,
+                documents=documents,
+                top_k=kwargs.get("top_k"),
+                return_documents=kwargs.get("return_documents", True)
+            )
         else:
             raise NotImplementedError(f"{self.__class__.__name__} does not support task: {task}")
     
@@ -67,7 +94,51 @@ class BaseEmbedService(BaseService):
         Returns:
             List of supported task names
         """
-        return ["embed", "embed_batch", "chunk_and_embed", "similarity", "find_similar"]
+        return ["embed", "embed_batch", "chunk", "chunk_and_embed", "similarity", "find_similar", "rerank", "rerank_documents", "document_ranking"]
+    
+    async def _text_similarity_search(self, query_text: str, candidates: List[str], **kwargs) -> Dict[str, Any]:
+        """
+        Helper method for text-based similarity search
+        
+        Args:
+            query_text: Query text
+            candidates: List of candidate texts
+            **kwargs: Additional parameters (top_k, threshold, etc.)
+            
+        Returns:
+            Dictionary containing similar documents with scores
+        """
+        # Get embeddings for query and candidates
+        query_embedding = await self.create_text_embedding(query_text)
+        candidate_embeddings = await self.create_text_embeddings(candidates)
+        
+        # Find similar texts
+        similar_results = await self.find_similar_texts(
+            query_embedding, 
+            candidate_embeddings, 
+            kwargs.get("top_k", len(candidates))
+        )
+        
+        # Apply threshold if specified
+        threshold = kwargs.get("threshold")
+        if threshold is not None:
+            similar_results = [r for r in similar_results if r["similarity"] >= threshold]
+        
+        # Convert to expected format with text content
+        similar_documents = []
+        for result in similar_results:
+            similar_documents.append({
+                "text": candidates[result["index"]],
+                "similarity": result["similarity"],
+                "index": result["index"]
+            })
+        
+        return {
+            "similar_documents": similar_documents,
+            "query": query_text,
+            "total_candidates": len(candidates),
+            "returned_count": len(similar_documents)
+        }
     
     @abstractmethod
     async def create_text_embedding(self, text: str) -> List[float]:
@@ -170,6 +241,38 @@ class BaseEmbedService(BaseService):
         """
         pass
     
+    async def rerank_documents(
+        self, 
+        query: str, 
+        documents: List[str], 
+        top_k: Optional[int] = None,
+        return_documents: bool = True
+    ) -> Dict[str, Any]:
+        """
+        Rerank documents based on relevance to query
+        
+        Default implementation returns NotImplementedError.
+        Override in subclasses that support reranking.
+        
+        Args:
+            query: Search query string
+            documents: List of documents to rerank
+            top_k: Number of top results to return (None = all)
+            return_documents: Whether to include document text in results
+            
+        Returns:
+            Dictionary containing:
+            - success: Boolean success status
+            - results: List of ranked documents with scores
+            - metadata: Additional information (model, timing, etc.)
+        """
+        return {
+            'success': False,
+            'error': f'Reranking not supported by {self.__class__.__name__}',
+            'provider': getattr(self, 'provider_name', 'unknown'),
+            'service': getattr(self, 'model_name', 'unknown')
+        }
+    
     @abstractmethod
     async def close(self):
         """Cleanup resources"""