isa-model 0.4.0__py3-none-any.whl → 0.4.4__py3-none-any.whl

isa_model/client.py

@@ -72,10 +72,14 @@ print(result["result"])
 
 import logging
 import asyncio
+import time
+import uuid
 from typing import Any, Dict, Optional, List, Union
 from pathlib import Path
+from datetime import datetime, timezone
 
 from isa_model.inference.ai_factory import AIFactory
+from isa_model.core.logging import get_inference_logger, generate_request_id
 
 try:
     from isa_model.core.services.intelligent_model_selector import IntelligentModelSelector, get_model_selector
@@ -213,6 +217,9 @@ class ISAModelClient:
         # Cache for frequently used services
         self._service_cache: Dict[str, Any] = {}
         
+        # Initialize inference logger
+        self.inference_logger = get_inference_logger()
+        
         logger.info("ISA Model Client initialized")
     
     async def _get_http_session(self):
@@ -327,6 +334,9 @@ class ISAModelClient:
         provider: Optional[str] = None,
         stream: Optional[bool] = None,
         show_reasoning: Optional[bool] = False,
+        output_format: Optional[str] = None,
+        json_schema: Optional[Dict] = None,
+        repair_attempts: Optional[int] = 3,
         **kwargs
     ) -> Dict[str, Any]:
         """
@@ -409,6 +419,9 @@ class ISAModelClient:
                     model_hint=model,
                     provider_hint=provider,
                     show_reasoning=show_reasoning,  # Explicitly pass show_reasoning
+                    output_format=output_format,
+                    json_schema=json_schema,
+                    repair_attempts=repair_attempts,
                     **kwargs
                 )
             else:
@@ -420,6 +433,9 @@ class ISAModelClient:
                     model_hint=model,
                     provider_hint=provider,
                     stream=False,  # Force non-streaming
+                    output_format=output_format,
+                    json_schema=json_schema,
+                    repair_attempts=repair_attempts,
                     **kwargs
                 )
                 
@@ -488,7 +504,7 @@ class ISAModelClient:
             )
             
             # Step 2: Get appropriate service
-            service = await self._get_service(
+            service, _ = await self._get_service(
                 service_type=service_type,
                 model_name=selected_model["model_id"],
                 provider=selected_model["provider"],
@@ -508,7 +524,9 @@ class ISAModelClient:
             content_chunks = []
             async for token in service.astream(input_data):
                 content_chunks.append(token)
-                yield token
+                # Only yield string tokens for streaming (filter out dict/objects)
+                if isinstance(token, str):
+                    yield token
             
             # Step 6: After streaming is complete, calculate billing info and optionally return metadata
             try:
@@ -533,7 +551,7 @@ class ISAModelClient:
                         "billing": billing_info,
                         "streaming": True,
                         "tokens_streamed": len(content_chunks),
-                        "content_length": len("".join(content_chunks))
+                        "content_length": len("".join(str(chunk) if isinstance(chunk, str) else "" for chunk in content_chunks))
                     }
                     yield ('metadata', metadata)
                 
@@ -554,7 +572,7 @@ class ISAModelClient:
                         },
                         "streaming": True,
                         "tokens_streamed": len(content_chunks),
-                        "content_length": len("".join(content_chunks))
+                        "content_length": len("".join(str(chunk) if isinstance(chunk, str) else "" for chunk in content_chunks))
                     }
                     yield ('metadata', fallback_metadata)
                 
@@ -562,6 +580,111 @@ class ISAModelClient:
             logger.error(f"Streaming invoke failed: {e}")
             raise
     
+    def _is_rate_limit_error(self, error: Exception) -> bool:
+        """Check if an error is due to rate limiting"""
+        error_str = str(error).lower()
+        
+        # Check for common rate limit indicators
+        rate_limit_indicators = [
+            'rate limit',
+            'rate_limit', 
+            'ratelimit',
+            'too many requests',
+            'quota exceeded',
+            'limit exceeded',
+            'throttled',
+            '429'
+        ]
+        
+        return any(indicator in error_str for indicator in rate_limit_indicators)
+    
+    async def _invoke_with_fallback(
+        self,
+        service_type: str,
+        task: str,
+        input_data: Any,
+        selected_model: Dict[str, Any],
+        **kwargs
+    ) -> Any:
+        """Invoke service with automatic fallback on rate limit"""
+        try:
+            # First attempt with selected model
+            return await self._invoke_service_direct(service_type, task, input_data, selected_model, **kwargs)
+        except Exception as e:
+            # Check if this is a rate limit error
+            if self._is_rate_limit_error(e):
+                logger.warning(f"Rate limit detected for {selected_model['provider']}: {e}")
+                
+                # Try to get fallback model using intelligent model selector
+                if INTELLIGENT_SELECTOR_AVAILABLE and self.model_selector:
+                    try:
+                        fallback_selection = self.model_selector.get_rate_limit_fallback(
+                            service_type, 
+                            selected_model['provider']
+                        )
+                        
+                        if fallback_selection.get('success') and fallback_selection.get('is_fallback'):
+                            fallback_model = fallback_selection['selected_model']
+                            logger.info(f"Switching to fallback: {fallback_model['provider']}/{fallback_model['model_id']}")
+                            
+                            # Retry with fallback model
+                            return await self._invoke_service_direct(service_type, task, input_data, fallback_model, **kwargs)
+                    except Exception as fallback_error:
+                        logger.error(f"Fallback also failed: {fallback_error}")
+                        raise e  # Raise original rate limit error
+            
+            # Re-raise the original error if not rate limit or fallback failed
+            raise
+    
+    async def _invoke_service_direct(
+        self,
+        service_type: str,
+        task: str,
+        input_data: Any,
+        model_config: Dict[str, Any],
+        **kwargs
+    ) -> Any:
+        """Direct service invocation without fallback logic"""
+        # Get appropriate service
+        factory = AIFactory.get_instance()
+        
+        # Create service with the specified model
+        if service_type == "text":
+            service = factory.get_llm(model_config["model_id"], model_config["provider"])
+        elif service_type == "vision":
+            service = factory.get_vision(model_config["model_id"], model_config["provider"])
+        elif service_type == "audio":
+            service = factory.get_audio(model_config["model_id"], model_config["provider"])
+        elif service_type == "image":
+            service = factory.get_image(model_config["model_id"], model_config["provider"])
+        elif service_type == "embedding":
+            service = factory.get_embed(model_config["model_id"], model_config["provider"])
+        else:
+            raise ValueError(f"Unsupported service type: {service_type}")
+        
+        # Invoke the service
+        if service_type == "text":
+            show_reasoning = kwargs.pop('show_reasoning', False)
+            
+            # Check if service supports show_reasoning parameter (mainly OpenAI services)
+            if model_config["provider"] == "openai":
+                result = await service.invoke(
+                    input_data=input_data,
+                    task=task,
+                    show_reasoning=show_reasoning,
+                    **kwargs
+                )
+            else:
+                # For other providers like yyds, don't pass show_reasoning
+                result = await service.invoke(
+                    input_data=input_data,
+                    task=task,
+                    **kwargs
+                )
+            return result
+        else:
+            return await service.invoke(input_data=input_data, task=task, **kwargs)
+
     async def _select_model(
         self,
         input_data: Any,
@@ -661,6 +784,7 @@ class ISAModelClient:
             "audio": {
                 "tts": {"model_id": "tts-1", "provider": "openai"},
                 "stt": {"model_id": "whisper-1", "provider": "openai"},
+                "realtime": {"model_id": "gpt-4o-realtime-preview-2024-10-01", "provider": "openai"},
                 "default": {"model_id": "whisper-1", "provider": "openai"}
             },
             "text": {
@@ -680,9 +804,14 @@ class ISAModelClient:
         
         # Handle audio service type with task-specific models
         if service_type == "audio":
-            if "speech" in task or "tts" in task:
+            # Realtime audio tasks
+            if any(realtime_task in task for realtime_task in ["realtime", "audio_chat", "text_chat", "create_session", "connect", "send_audio", "send_text", "listen"]):
+                default = defaults["audio"]["realtime"]
+            # Traditional TTS tasks
+            elif "speech" in task or "tts" in task or task in ["synthesize", "text_to_speech", "generate_speech"]:
                 default = defaults["audio"]["tts"]
-            elif "transcribe" in task or "stt" in task:
+            # Traditional STT tasks
+            elif "transcribe" in task or "stt" in task or task in ["speech_to_text", "transcription"]:
                 default = defaults["audio"]["stt"] 
             else:
                 default = defaults["audio"]["default"]
@@ -714,14 +843,16 @@ class ISAModelClient:
         provider: str,
         task: str,
         use_cache: bool = True
-    ) -> Any:
-        """Get appropriate service instance"""
+    ) -> tuple[Any, str]:
+        """Get appropriate service instance and return actual model used"""
         
-        cache_key = f"{service_type}_{provider}_{model_name}"
+        cache_key = f"{service_type}_{provider}_{model_name}_{task}"
+        actual_model_used = model_name  # Track the actual model used
         
         # Check cache first (if caching is enabled)
         if use_cache and cache_key in self._service_cache:
-            return self._service_cache[cache_key]
+            cached_service, cached_model = self._service_cache[cache_key]
+            return cached_service, cached_model
         
         try:
             # Validate service type
@@ -730,24 +861,46 @@ class ISAModelClient:
             # Route to appropriate AIFactory method
             if service_type == "vision":
                 service = self.ai_factory.get_vision(model_name, provider)
+                actual_model_used = model_name
             elif service_type == "audio":
-                if "speech" in task or "tts" in task:
-                    service = self.ai_factory.get_tts(model_name, provider)
-                elif "transcribe" in task or "stt" in task:
-                    service = self.ai_factory.get_stt(model_name, provider)
+                # Realtime audio tasks
+                if any(realtime_task in task for realtime_task in ["realtime", "audio_chat", "text_chat", "create_session", "connect", "send_audio", "send_text", "listen"]):
+                    # Use realtime model
+                    realtime_model = "gpt-4o-realtime-preview-2024-10-01" if model_name == "tts-1" or model_name == "whisper-1" else model_name
+                    service = self.ai_factory.get_realtime(realtime_model, provider)
+                    actual_model_used = realtime_model
+                # Traditional TTS tasks
+                elif "speech" in task or "tts" in task or task in ["synthesize", "text_to_speech", "generate_speech"]:
+                    # Use TTS model
+                    tts_model = "tts-1" if model_name == "whisper-1" else model_name
+                    service = self.ai_factory.get_tts(tts_model, provider)
+                    actual_model_used = tts_model
+                # Traditional STT tasks
+                elif "transcribe" in task or "stt" in task or task in ["speech_to_text", "transcription"]:
+                    # Use STT model
+                    stt_model = "whisper-1" if model_name == "tts-1" else model_name
+                    service = self.ai_factory.get_stt(stt_model, provider)
+                    actual_model_used = stt_model
+                # Default to STT for backward compatibility
                 else:
-                    service = self.ai_factory.get_stt(model_name, provider)
+                    # Use STT model by default
+                    stt_model = "whisper-1" if model_name == "tts-1" else model_name
+                    service = self.ai_factory.get_stt(stt_model, provider)
+                    actual_model_used = stt_model
             elif service_type == "text":
                 service = self.ai_factory.get_llm(model_name, provider)
+                actual_model_used = model_name
             elif service_type == "image":
                 service = self.ai_factory.get_img("t2i", model_name, provider)
+                actual_model_used = model_name
             elif service_type == "embedding":
                 service = self.ai_factory.get_embed(model_name, provider)
+                actual_model_used = model_name
             
-            # Cache the service (if caching is enabled)
+            # Cache the service and actual model (if caching is enabled)
             if use_cache:
-                self._service_cache[cache_key] = service
-            return service
+                self._service_cache[cache_key] = (service, actual_model_used)
+            return service, actual_model_used
             
         except Exception as e:
             logger.error(f"Failed to get service {service_type}/{provider}/{model_name}: {e}")
@@ -785,12 +938,26 @@ class ISAModelClient:
                 )
             
             elif service_type == "audio":
-                if unified_task in ["synthesize", "text_to_speech", "tts"]:
+                # Realtime audio tasks
+                if any(realtime_task in unified_task for realtime_task in ["realtime", "audio_chat", "text_chat", "create_session", "connect", "send_audio", "send_text", "listen"]):
+                    # For realtime text_chat and audio_chat, pass text parameter
+                    if unified_task in ["text_chat", "audio_chat"]:
+                        if isinstance(input_data, str):
+                            kwargs['text'] = input_data
+                        elif isinstance(input_data, bytes):
+                            kwargs['audio_data'] = input_data
+                    return await service.invoke(
+                        task=unified_task,
+                        **kwargs
+                    )
+                # Traditional TTS tasks
+                elif unified_task in ["synthesize", "text_to_speech", "tts", "generate_speech"]:
                     return await service.invoke(
                         text=input_data,
                         task=unified_task,
                         **kwargs
                     )
+                # Traditional STT tasks
                 else:
                     return await service.invoke(
                         audio_input=input_data,
@@ -801,22 +968,58 @@ class ISAModelClient:
             elif service_type == "text":
                 # Extract show_reasoning from kwargs if present
                 show_reasoning = kwargs.pop('show_reasoning', False)
-                result = await service.invoke(
-                    input_data=input_data,
-                    task=unified_task,
-                    show_reasoning=show_reasoning,
-                    **kwargs
-                )
+                
+                # Check if service provider supports show_reasoning
+                # Only OpenAI services support this parameter
+                if hasattr(service, 'provider_name') and service.provider_name == 'openai':
+                    result = await service.invoke(
+                        input_data=input_data,
+                        task=unified_task,
+                        show_reasoning=show_reasoning,
+                        **kwargs
+                    )
+                else:
+                    # For other providers like yyds, don't pass show_reasoning
+                    result = await service.invoke(
+                        input_data=input_data,
+                        task=unified_task,
+                        **kwargs
+                    )
                 
                 logger.debug(f"Service result type: {type(result)}")
                 logger.debug(f"Service result: {result}")
                 
-                if isinstance(result, dict) and 'message' in result:
+                # Check if this is a formatted result from invoke method
+                if isinstance(result, dict) and 'formatted' in result:
+                    # This is a formatted result from the new invoke method
+                    logger.debug(f"Returning formatted result: {result}")
+                    return result
+                elif isinstance(result, dict) and 'message' in result:
+                    # This is a traditional message result
                     message = result['message']
                     logger.debug(f"Extracted message type: {type(message)}")
-                    logger.debug(f"Extracted message: {message}")
-                    return message
+                    logger.debug(f"Extracted message length: {len(str(message)) if message else 0}")
+                    
+                    # Handle AIMessage objects from LangChain
+                    if hasattr(message, 'content'):
+                        # Check if there are tool_calls
+                        if hasattr(message, 'tool_calls') and message.tool_calls:
+                            logger.debug(f"AIMessage contains tool_calls: {len(message.tool_calls)}")
+                            # Return a dict with both content and tool_calls
+                            return {
+                                "content": message.content if message.content else "",
+                                "tool_calls": message.tool_calls
+                            }
+                        else:
+                            content = message.content
+                            logger.debug(f"Extracted content from AIMessage: {len(content) if content else 0} chars")
+                            return content
+                    else:
+                        # Direct string message
+                        logger.debug(f"Returning direct message: {len(str(message)) if message else 0} chars")
+                        return message
                 else:
+                    logger.debug(f"Returning result directly: {result}")
                     return result
             
             elif service_type == "image":
@@ -886,7 +1089,7 @@ class ISAModelClient:
             
             for service_type, provider, model in test_services:
                 try:
-                    await self._get_service(service_type, model, provider, "test")
+                    service, _ = await self._get_service(service_type, model, provider, "test")
                     health_status["services"][f"{service_type}_{provider}"] = "healthy"
                 except Exception as e:
                     health_status["services"][f"{service_type}_{provider}"] = f"error: {str(e)}"
@@ -916,9 +1119,18 @@ class ISAModelClient:
         service_type: str,
         model_hint: Optional[str] = None,
         provider_hint: Optional[str] = None,
+        output_format: Optional[str] = None,
+        json_schema: Optional[Dict] = None,
+        repair_attempts: Optional[int] = 3,
         **kwargs
     ) -> Dict[str, Any]:
         """Service invoke that returns streaming response with async generator"""
+        
+        # Generate unique request ID for logging
+        request_id = generate_request_id()
+        start_time = datetime.now(timezone.utc)
+        execution_start_time = time.time()
+        
         try:
             # Step 1: Select best model for this task
             selected_model = await self._select_model(
@@ -930,18 +1142,20 @@ class ISAModelClient:
             )
             
             # Step 2: Get appropriate service
-            service = await self._get_service(
+            service, actual_model_used = await self._get_service(
                 service_type=service_type,
                 model_name=selected_model["model_id"],
                 provider=selected_model["provider"],
                 task=task,
                 use_cache=False  # Don't cache for streaming to avoid state issues
             )
+            # Update selected model with actual model used
+            selected_model["model_id"] = actual_model_used
             
             # Step 3: Handle tools for LLM services (bind tools if provided)
             tools = kwargs.pop("tools", None)
             if service_type == "text" and tools:
-                service = await self._get_service(
+                service, _ = await self._get_service(
                     service_type=service_type,
                     model_name=selected_model["model_id"],
                     provider=selected_model["provider"],
@@ -964,7 +1178,8 @@ class ISAModelClient:
                 if service_type == "text" and hasattr(service, 'astream'):
                     show_reasoning = kwargs.get('show_reasoning', False)
                     logger.debug(f"Stream generator: show_reasoning={show_reasoning}")
-                    if 'show_reasoning' in kwargs:
+                    # Only pass show_reasoning to OpenAI providers
+                    if 'show_reasoning' in kwargs and hasattr(service, 'provider_name') and service.provider_name == 'openai':
                         async for token in service.astream(input_data, show_reasoning=show_reasoning):
                             yield token
                     else:
@@ -999,9 +1214,18 @@ class ISAModelClient:
         model_hint: Optional[str] = None,
         provider_hint: Optional[str] = None,
         stream: Optional[bool] = None,
+        output_format: Optional[str] = None,
+        json_schema: Optional[Dict] = None,
+        repair_attempts: Optional[int] = 3,
         **kwargs
     ) -> Dict[str, Any]:
         """Direct service invoke - passes LangChain objects and tools directly to services"""
+        
+        # Generate unique request ID for logging
+        request_id = generate_request_id()
+        start_time = datetime.now(timezone.utc)
+        execution_start_time = time.time()
+        
         try:
             # Step 1: Select best model for this task
             selected_model = await self._select_model(
@@ -1012,18 +1236,35 @@ class ISAModelClient:
                 provider_hint=provider_hint
             )
             
+            # Step 1.5: Log inference start
+            self.inference_logger.log_inference_start(
+                request_id=request_id,
+                service_type=service_type,
+                task=task,
+                provider=selected_model["provider"],
+                model_name=selected_model["model_id"],
+                input_data=input_data if self.inference_logger.log_detailed_requests else None,
+                is_streaming=stream or False,
+                custom_metadata={
+                    "selection_reason": selected_model.get("reason", "Default selection"),
+                    "has_tools": "tools" in kwargs
+                }
+            )
+            
             # Step 2: Get appropriate service
-            service = await self._get_service(
+            service, actual_model_used = await self._get_service(
                 service_type=service_type,
                 model_name=selected_model["model_id"],
                 provider=selected_model["provider"],
                 task=task
             )
+            # Update selected model with actual model used
+            selected_model["model_id"] = actual_model_used
             
             # Step 3: Handle tools for LLM services (bind tools if provided)
             tools = kwargs.pop("tools", None)
             if service_type == "text" and tools:
-                service = await self._get_service(
+                service, _ = await self._get_service(
                     service_type=service_type,
                     model_name=selected_model["model_id"],
                     provider=selected_model["provider"],
@@ -1039,32 +1280,214 @@ class ISAModelClient:
                     service.streaming = stream
             
             # Step 5: Execute task with unified interface
-            result = await self._execute_task(
-                service=service,
-                input_data=input_data,
-                task=task,
-                service_type=service_type,
-                **kwargs
-            )
+            # Pass JSON formatting parameters to the service
+            task_kwargs = kwargs.copy()
+            if service_type == "text":
+                if output_format:
+                    task_kwargs["output_format"] = output_format
+                if json_schema:
+                    task_kwargs["json_schema"] = json_schema
+                if repair_attempts is not None:
+                    task_kwargs["repair_attempts"] = repair_attempts
+            
+            # Try to execute with rate limit detection
+            try:
+                result = await self._execute_task(
+                    service=service,
+                    input_data=input_data,
+                    task=task,
+                    service_type=service_type,
+                    **task_kwargs
+                )
+            except Exception as e:
+                # Check if this is a rate limit error and we can fallback
+                if self._is_rate_limit_error(e) and service_type == "text":
+                    # Ensure model selector is initialized
+                    if not self.model_selector:
+                        self.model_selector = await get_model_selector(self.config)
+                    
+                    # Get fallback model selection
+                    fallback_selection = self.model_selector.get_rate_limit_fallback(
+                        service_type=service_type,
+                        original_provider=selected_model["provider"]
+                    )
+                    
+                    if fallback_selection.get('success'):
+                        fallback_model = fallback_selection.get('selected_model', {})
+                        logger.info(f"Rate limit hit, switching to fallback: {fallback_model}")
+                        
+                        # Get fallback service
+                        fallback_service, fallback_model_used = await self._get_service(
+                            service_type=service_type,
+                            model_name=fallback_model["model_id"],
+                            provider=fallback_model["provider"],
+                            task=task
+                        )
+                        
+                        # Update selected model for metadata
+                        selected_model = fallback_model
+                        selected_model["model_id"] = fallback_model_used
+                        selected_model["reason"] = "Rate limit fallback"
+                        
+                        # Retry with fallback service
+                        result = await self._execute_task(
+                            service=fallback_service,
+                            input_data=input_data,
+                            task=task,
+                            service_type=service_type,
+                            **task_kwargs
+                        )
+                    else:
+                        # No fallback available, re-raise original error
+                        raise
+                else:
+                    # Not a rate limit error or no fallback, re-raise
+                    raise
             
             # Step 6: Wait for billing tracking to complete, then get billing information
             await asyncio.sleep(0.01)  # Small delay to ensure billing tracking completes
             billing_info = self._get_billing_info(service, selected_model["model_id"])
             
+            # Step 6.5: Calculate execution time and log completion
+            execution_time_ms = int((time.time() - execution_start_time) * 1000)
+            
+            # Log inference completion
+            self.inference_logger.log_inference_complete(
+                request_id=request_id,
+                status="completed",
+                execution_time_ms=execution_time_ms,
+                input_tokens=billing_info.get("input_tokens"),
+                output_tokens=billing_info.get("output_tokens"),
+                estimated_cost_usd=billing_info.get("cost_usd"),
+                output_data=result if self.inference_logger.log_detailed_requests else None,
+                custom_metadata={
+                    "billing_operation": billing_info.get("operation"),
+                    "timestamp": billing_info.get("timestamp")
+                }
+            )
+            
+            # Log detailed token usage if available
+            if billing_info.get("input_tokens") and billing_info.get("output_tokens"):
+                self.inference_logger.log_token_usage(
+                    request_id=request_id,
+                    provider=selected_model["provider"],
+                    model_name=selected_model["model_id"],
+                    prompt_tokens=billing_info.get("input_tokens"),
+                    completion_tokens=billing_info.get("output_tokens"),
+                    prompt_cost_usd=billing_info.get("cost_usd", 0) * 0.6 if billing_info.get("cost_usd") else None,  # Rough estimate
+                    completion_cost_usd=billing_info.get("cost_usd", 0) * 0.4 if billing_info.get("cost_usd") else None
+                )
+            
+            # Handle formatting - check if result is already formatted
+            formatted_result = result
+            if service_type == "text" and output_format:
+                # Check if result is already formatted by the service
+                if isinstance(result, dict) and result.get("formatted"):
+                    # Result is already formatted by the service
+                    formatted_result = result.get("result", result)
+                    billing_info["formatting"] = {
+                        "output_format": output_format,
+                        "format_success": True,
+                        "format_method": "service_level",
+                        "format_errors": result.get("format_errors", []),
+                        "repaired": False,
+                        "pre_formatted": True
+                    }
+                else:
+                    # Apply formatting at client level (fallback)
+                    try:
+                        service, _ = await self._get_service(
+                            service_type=service_type,
+                            model_name=selected_model["model_id"],
+                            provider=selected_model["provider"],
+                            task=task
+                        )
+                        if hasattr(service, 'format_structured_output'):
+                            formatting_result = service.format_structured_output(
+                                response=result,
+                                output_format=output_format,
+                                schema=json_schema,
+                                repair_attempts=repair_attempts or 3
+                            )
+                            # Update result and add formatting metadata
+                            if formatting_result.get("success") and formatting_result.get("data") is not None:
+                                # Extract the actual formatted data
+                                formatted_data = formatting_result["data"]
+                                
+                                # For JSON output, ensure we return clean data
+                                if output_format == "json" and isinstance(formatted_data, dict):
+                                    formatted_result = formatted_data
+                                else:
+                                    formatted_result = formatted_data
+                            else:
+                                # Keep original result if formatting failed
+                                formatted_result = result
+                            
+                            # Add formatting info to metadata
+                            billing_info["formatting"] = {
+                                "output_format": output_format,
+                                "format_success": formatting_result.get("success", False),
+                                "format_method": formatting_result.get("method"),
+                                "format_errors": formatting_result.get("errors", []),
+                                "repaired": formatting_result.get("repaired", False),
+                                "pre_formatted": False
+                            }
+                            
+                    except Exception as format_error:
+                        logger.warning(f"Failed to apply output formatting: {format_error}")
+                        # Continue with unformatted result
+                        formatted_result = result
+                        billing_info["formatting"] = {
+                            "output_format": output_format,
+                            "format_success": False,
+                            "format_error": str(format_error)
+                        }
+            
             # Return unified response
-            return {
+            response = {
                 "success": True,
-                "result": result,
+                "result": formatted_result,
                 "metadata": {
+                    "request_id": request_id,  # Include request ID for tracking
                     "model_used": selected_model["model_id"],
                     "provider": selected_model["provider"], 
                     "task": task,
                     "service_type": service_type,
                     "selection_reason": selected_model.get("reason", "Default selection"),
+                    "execution_time_ms": execution_time_ms,
                     "billing": billing_info
                 }
             }
+            
+            return response
         except Exception as e:
+            # Calculate execution time even for errors
+            execution_time_ms = int((time.time() - execution_start_time) * 1000)
+            
+            # Log inference error
+            error_type = type(e).__name__
+            error_message = str(e)
+            
+            self.inference_logger.log_inference_complete(
+                request_id=request_id,
+                status="failed",
+                execution_time_ms=execution_time_ms,
+                error_message=error_message,
+                error_code=error_type,
+                custom_metadata={
+                    "error_location": "client._invoke_service"
+                }
+            )
+            
+            # Also log to the error table
+            self.inference_logger.log_error(
+                request_id=request_id,
+                error_type=error_type,
+                error_message=error_message,
+                provider=model_hint or "unknown",
+                model_name=provider_hint or "unknown"
+            )
+            
             logger.error(f"Service invoke failed: {e}")
             raise