isa-model 0.3.9__py3-none-any.whl → 0.4.0__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
@@ -72,13 +73,38 @@ class OpenAISTTService(BaseSTTService):
             if prompt:
                 transcription_params["prompt"] = prompt
             
-            # Handle file input
+            # Handle file input - support base64 strings, file paths, and file objects
             if isinstance(audio_file, str):
-                with open(audio_file, "rb") as f:
-                    transcription = await self.client.audio.transcriptions.create(
-                        file=f,
-                        **transcription_params
-                    )
+                # 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,

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/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"""

isa_model/inference/services/embedding/isa_embed_service.py

@@ -0,0 +1,305 @@
+"""
+ISA Embedding Service
+
+ISA reranking service using deployed Jina Reranker v2 via Modal
+"""
+
+import logging
+from typing import Dict, Any, List, Optional
+
+try:
+    import modal
+    MODAL_AVAILABLE = True
+except ImportError:
+    MODAL_AVAILABLE = False
+    modal = None
+
+from isa_model.inference.services.embedding.base_embed_service import BaseEmbedService
+
+logger = logging.getLogger(__name__)
+
+class ISAEmbedService(BaseEmbedService):
+    """
+    ISA Embedding Service - calls ISA deployed reranking models
+    
+    Supported features:
+    - Document reranking (Jina Reranker v2 via Modal)
+    - Future: embedding generation
+    - Future: semantic similarity computation
+    """
+    
+    def __init__(self, 
+                 rerank_modal_app_name: str = "isa-embed-rerank",
+                 timeout: int = 30):
+        """
+        Initialize ISA Embedding service
+        
+        Args:
+            rerank_modal_app_name: Modal reranking app name
+            timeout: Request timeout in seconds
+        """
+        # For now, skip BaseService initialization to avoid config validation
+        # TODO: Properly configure ISA provider in config system
+        self.provider_name = "isa"
+        self.model_name = "isa-jina-reranker-v2-service"
+        self.rerank_modal_app_name = rerank_modal_app_name
+        self.timeout = timeout
+        
+        # Initialize Modal client
+        if MODAL_AVAILABLE:
+            try:
+                # Get deployed Modal application
+                self.modal_app = modal.App.lookup(rerank_modal_app_name)
+                logger.info(f"Connected to Modal rerank app: {rerank_modal_app_name}")
+                
+                self.modal_service = True  # Mark service as available
+                logger.info("Modal rerank app connection established")
+                    
+            except Exception as e:
+                logger.warning(f"Failed to connect to Modal rerank app: {e}")
+                self.modal_app = None
+                self.modal_service = None
+        else:
+            logger.warning("Modal SDK not available")
+            self.modal_app = None
+            self.modal_service = None
+        
+        # Service statistics
+        self.request_count = 0
+        self.total_cost = 0.0
+        
+    async def rerank_documents(
+        self, 
+        query: str, 
+        documents: List[str], 
+        top_k: Optional[int] = None,
+        return_documents: bool = True
+    ) -> Dict[str, Any]:
+        """
+        Rerank documents using Jina Reranker v2
+        
+        Args:
+            query: Query string
+            documents: List of documents to rerank
+            top_k: Return top k results (None = all)
+            return_documents: Whether to include document content in results
+            
+        Returns:
+            Reranking results
+        """
+        try:
+            if not self.modal_app or not self.modal_service:
+                return {
+                    'success': False,
+                    'provider': 'ISA',
+                    'service': 'isa-embed-rerank',
+                    'error': 'Modal rerank app or service not available'
+                }
+            
+            # Call reranking service directly via Modal SDK
+            result = await self._call_rerank_service(query, documents, top_k, return_documents)
+            
+            if result and result.get('success', False):
+                self.request_count += 1
+                
+                # Record cost
+                if 'billing' in result:
+                    cost = result['billing'].get('estimated_cost_usd', 0)
+                    self.total_cost += cost
+                
+                # Format response to match expected structure
+                formatted_result = {
+                    'success': True,
+                    'provider': 'ISA',
+                    'service': 'isa-embed-rerank',
+                    'result': {
+                        'results': result.get('results', []),
+                        'processing_time': result.get('processing_time'),
+                        'billing': result.get('billing', {}),
+                        'query': result.get('query'),
+                        'num_documents': result.get('num_documents'),
+                        'returned_count': result.get('returned_count')
+                    },
+                    'metadata': {
+                        'model_used': result.get('model'),
+                        'provider': result.get('provider', 'ISA'),
+                        'billing': result.get('billing', {})
+                    }
+                }
+                return formatted_result
+            else:
+                return {
+                    'success': False,
+                    'provider': 'ISA',
+                    'service': 'isa-embed-rerank',
+                    'error': f'Rerank service returned error: {result.get("error", "Unknown error") if result else "No response"}',
+                    'details': result
+                }
+                
+        except Exception as e:
+            logger.error(f"ISA document reranking failed: {e}")
+            import traceback
+            traceback.print_exc()
+            return {
+                'success': False,
+                'provider': 'ISA',
+                'service': 'isa-embed-rerank',
+                'error': str(e)
+            }
+    
+    async def _call_rerank_service(
+        self, 
+        query: str, 
+        documents: List[str], 
+        top_k: Optional[int], 
+        return_documents: bool
+    ) -> Dict[str, Any]:
+        """
+        Call reranking service via Modal SDK
+        """
+        try:
+            import modal
+            
+            logger.info("Calling Jina Reranker v2 service via Modal SDK...")
+            
+            # Correct Modal SDK usage: call deployed class method
+            ISAEmbedRerankService = modal.Cls.from_name(
+                app_name=self.rerank_modal_app_name,
+                name="ISAEmbedRerankService"
+            )
+            
+            # Create instance and call method
+            instance = ISAEmbedRerankService()
+            result = instance.rerank_documents.remote(
+                query=query,
+                documents=documents,
+                top_k=top_k,
+                return_documents=return_documents
+            )
+            
+            logger.info("Modal rerank SDK call successful")
+            return result
+                        
+        except Exception as e:
+            logger.error(f"Modal rerank SDK call failed: {e}")
+            return {
+                'success': False,
+                'error': f'Modal rerank SDK error: {str(e)}'
+            }
+    
+    # ==================== Embedding methods (future implementation) ====================
+    
+    async def create_text_embedding(self, text: str) -> List[float]:
+        """Create single text embedding - not yet implemented"""
+        raise NotImplementedError("Text embedding not yet implemented in ISA service")
+    
+    async def create_text_embeddings(self, texts: List[str]) -> List[List[float]]:
+        """Create multiple text embeddings - not yet implemented"""
+        raise NotImplementedError("Text embeddings not yet implemented in ISA service")
+    
+    async def create_chunks(self, text: str, metadata: Optional[Dict] = None) -> List[Dict]:
+        """Create text chunks with embeddings - not yet implemented"""
+        raise NotImplementedError("Text chunking not yet implemented in ISA service")
+    
+    async def compute_similarity(self, embedding1: List[float], embedding2: List[float]) -> float:
+        """Compute embedding similarity - not yet implemented"""
+        raise NotImplementedError("Similarity computation not yet implemented in ISA service")
+    
+    async def find_similar_texts(
+        self, 
+        query_embedding: List[float], 
+        candidate_embeddings: List[List[float]], 
+        top_k: int = 5
+    ) -> List[Dict[str, Any]]:
+        """Find similar texts - not yet implemented"""
+        raise NotImplementedError("Similar text search not yet implemented in ISA service")
+    
+    def get_embedding_dimension(self) -> int:
+        """Get embedding dimension - not applicable for rerank-only service"""
+        raise NotImplementedError("Embedding dimension not available for rerank-only service")
+    
+    def get_max_input_length(self) -> int:
+        """Get maximum input length"""
+        return 1024  # Jina Reranker v2 max length
+    
+    # ==================== Service management methods ====================
+    
+    async def health_check(self) -> Dict[str, Any]:
+        """Check ISA reranking service health"""
+        try:
+            # Simple health check: call reranking service
+            test_result = await self.rerank_documents(
+                query="test",
+                documents=["test document"],
+                top_k=1,
+                return_documents=False
+            )
+            
+            return {
+                'success': True,
+                'provider': 'ISA',
+                'service': 'isa-embed-rerank',
+                'status': 'healthy' if test_result.get('success') else 'error',
+                'rerank_service': test_result.get('success', False),
+                'usage_stats': {
+                    'total_requests': self.request_count,
+                    'total_cost_usd': round(self.total_cost, 6)
+                }
+            }
+                
+        except Exception as e:
+            return {
+                'success': False,
+                'provider': 'ISA',
+                'service': 'isa-embed-rerank',
+                'status': 'error',
+                'error': str(e)
+            }
+    
+    async def get_usage_stats(self) -> Dict[str, Any]:
+        """Get usage statistics"""
+        try:
+            modal_stats = {}
+            
+            # Try to get Modal service statistics
+            if self.modal_app:
+                try:
+                    # Can extend to get Modal service stats
+                    pass
+                except Exception as e:
+                    logger.warning(f"Failed to get Modal stats: {e}")
+            
+            return {
+                'provider': 'ISA',
+                'service': 'isa-embed-rerank',
+                'client_stats': {
+                    'total_requests': self.request_count,
+                    'total_cost_usd': round(self.total_cost, 6)
+                },
+                'modal_stats': modal_stats,
+                'combined_cost': round(self.total_cost, 6)
+            }
+            
+        except Exception as e:
+            return {
+                'provider': 'ISA', 
+                'service': 'isa-embed-rerank',
+                'error': str(e)
+            }
+    
+    def get_supported_tasks(self) -> List[str]:
+        """Get supported task list"""
+        return [
+            'rerank',           # Document reranking
+            'rerank_documents', # Document reranking (alias)
+            'document_ranking'  # Document ranking (alias)
+        ]
+    
+    def get_supported_formats(self) -> List[str]:
+        """Get supported formats"""
+        return ['text']  # Text only
+    
+    async def close(self):
+        """Cleanup resources"""
+        # Modal client doesn't need explicit closure
+        pass

isa_model/inference/services/embedding/openai_embed_service.py

@@ -129,11 +129,9 @@ class OpenAIEmbedService(BaseEmbedService):
             logger.error(f"Error creating text embeddings: {e}")
             raise
     
-    async def create_chunks(self, text: str, metadata: Optional[Dict] = None) -> List[Dict]:
+    async def create_chunks(self, text: str, metadata: Optional[Dict] = None, chunk_size: int = 400, overlap: int = 50, **kwargs) -> 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
+        # Use provided chunk_size and overlap, or defaults optimized for OpenAI models
         
         words = text.split()
         if not words: