isa-model 0.3.4__py3-none-any.whl → 0.3.6__py3-none-any.whl

isa_model/inference/services/llm/ollama_llm_service.py

@@ -3,19 +3,21 @@ import httpx
 import json
 from typing import Dict, Any, List, Union, AsyncGenerator, Optional, Callable
 from isa_model.inference.services.llm.base_llm_service import BaseLLMService
-from isa_model.inference.providers.base_provider import BaseProvider
 
 logger = logging.getLogger(__name__)
 
 class OllamaLLMService(BaseLLMService):
     """Ollama LLM service with unified invoke interface and proper adapter support"""
     
-    def __init__(self, provider: 'BaseProvider', model_name: str = "llama3.2:3b-instruct-fp16"):
-        super().__init__(provider, model_name)
+    def __init__(self, provider_name: str, model_name: str = "llama3.2:3b-instruct-fp16", **kwargs):
+        super().__init__(provider_name, model_name, **kwargs)
+        
+        # Get configuration from centralized config manager
+        provider_config = self.get_provider_config()
         
         # Create HTTP client for Ollama API
-        base_url = self.config.get("base_url", "http://localhost:11434")
-        timeout = self.config.get("timeout", 60)
+        base_url = provider_config.get("base_url", "http://localhost:11434")
+        timeout = provider_config.get("timeout", 60)
         
         self.client = httpx.AsyncClient(
             base_url=base_url,
@@ -31,13 +33,14 @@ class OllamaLLMService(BaseLLMService):
     def _ensure_client(self):
         """Ensure the HTTP client is available and not closed"""
         if not hasattr(self, 'client') or not self.client or self.client.is_closed:
-            base_url = self.config.get("base_url", "http://localhost:11434")
-            timeout = self.config.get("timeout", 60)
+            provider_config = self.get_provider_config()
+            base_url = provider_config.get("base_url", "http://localhost:11434")
+            timeout = provider_config.get("timeout", 60)
             self.client = httpx.AsyncClient(base_url=base_url, timeout=timeout)
     
     def _create_bound_copy(self) -> 'OllamaLLMService':
         """Create a copy of this service for tool binding"""
-        bound_service = OllamaLLMService(self.provider, self.model_name)
+        bound_service = OllamaLLMService(self.provider_name, self.model_name)
         bound_service._bound_tools = self._bound_tools.copy()
         return bound_service
     
@@ -70,14 +73,15 @@ class OllamaLLMService(BaseLLMService):
             messages = self._prepare_messages(input_data)
             
             # Prepare request parameters
+            provider_config = self.get_provider_config()
             payload = {
                 "model": self.model_name,
                 "messages": messages,
                 "stream": self.streaming,
                 "options": {
-                    "temperature": self.config.get("temperature", 0.7),
-                    "top_p": self.config.get("top_p", 0.9),
-                    "num_predict": self.config.get("max_tokens", 2048)
+                    "temperature": provider_config.get("temperature", 0.7),
+                    "top_p": provider_config.get("top_p", 0.9),
+                    "num_predict": provider_config.get("max_tokens", 2048)
                 }
             }
             
@@ -86,9 +90,15 @@ class OllamaLLMService(BaseLLMService):
             if tool_schemas:
                 payload["tools"] = tool_schemas
             
-            # Handle streaming
+            # Handle streaming vs non-streaming
             if self.streaming:
-                return self._stream_response(payload)
+                # TRUE STREAMING MODE - collect all chunks from the stream
+                content_chunks = []
+                async for token in self.astream(input_data):
+                    content_chunks.append(token)
+                content = "".join(content_chunks)
+                
+                return self._format_response(content, input_data)
             
             # Regular request
             response = await self.client.post("/api/chat", json=payload)
@@ -98,6 +108,7 @@ class OllamaLLMService(BaseLLMService):
             # Update token usage if available
             if "eval_count" in result:
                 self._update_token_usage(result)
+                await self._track_ollama_billing(result)
             
             # Handle tool calls if present - let adapter process the complete message
             message = result["message"]
@@ -190,6 +201,44 @@ class OllamaLLMService(BaseLLMService):
         # Get final response from the model
         return await self.ainvoke(messages)
     
+    async def _track_streaming_usage(self, messages: List[Dict[str, str]], content: str):
+        """Track usage for streaming requests (estimated)"""
+        # Create a mock usage object for tracking
+        class MockUsage:
+            def __init__(self):
+                self.prompt_tokens = len(str(messages)) // 4  # Rough estimate
+                self.completion_tokens = len(content) // 4   # Rough estimate  
+                self.total_tokens = self.prompt_tokens + self.completion_tokens
+        
+        usage = MockUsage()
+        self._update_token_usage_from_mock(usage)
+        
+        # Track billing
+        await self._track_llm_usage(
+            operation="chat_stream",
+            input_tokens=usage.prompt_tokens,
+            output_tokens=usage.completion_tokens,
+            metadata={
+                "model": self.model_name,
+                "provider": "ollama",
+                "streaming": True
+            }
+        )
+    
+    def _update_token_usage_from_mock(self, usage):
+        """Update token usage statistics from mock usage object"""
+        self.last_token_usage = {
+            "prompt_tokens": usage.prompt_tokens,
+            "completion_tokens": usage.completion_tokens,
+            "total_tokens": usage.total_tokens
+        }
+        
+        # Update total usage
+        self.total_token_usage["prompt_tokens"] += self.last_token_usage["prompt_tokens"]
+        self.total_token_usage["completion_tokens"] += self.last_token_usage["completion_tokens"]
+        self.total_token_usage["total_tokens"] += self.last_token_usage["total_tokens"]
+        self.total_token_usage["requests_count"] += 1
+
     def _update_token_usage(self, result: Dict[str, Any]):
         """Update token usage statistics"""
         self.last_token_usage = {
@@ -204,6 +253,21 @@ class OllamaLLMService(BaseLLMService):
         self.total_token_usage["total_tokens"] += self.last_token_usage["total_tokens"]
         self.total_token_usage["requests_count"] += 1
     
+    async def _track_ollama_billing(self, result: Dict[str, Any]):
+        """Track billing information for Ollama requests"""
+        prompt_tokens = result.get("prompt_eval_count", 0)
+        completion_tokens = result.get("eval_count", 0)
+        
+        await self._track_llm_usage(
+            operation="chat",
+            input_tokens=prompt_tokens,
+            output_tokens=completion_tokens,
+            metadata={
+                "model": self.model_name,
+                "provider": "ollama"
+            }
+        )
+    
     def get_token_usage(self) -> Dict[str, Any]:
         """Get total token usage statistics"""
         return self.total_token_usage
@@ -214,9 +278,10 @@ class OllamaLLMService(BaseLLMService):
     
     def get_model_info(self) -> Dict[str, Any]:
         """Get information about the current model"""
+        provider_config = self.get_provider_config()
         return {
             "name": self.model_name,
-            "max_tokens": self.config.get("max_tokens", 2048),
+            "max_tokens": provider_config.get("max_tokens", 2048),
             "supports_streaming": True,
             "supports_functions": True,
             "provider": "ollama"
@@ -230,4 +295,74 @@ class OllamaLLMService(BaseLLMService):
                 if not self.client.is_closed:
                     await self.client.aclose()
             except Exception as e:
-                logger.warning(f"Error closing Ollama client: {e}")
+                logger.warning(f"Error closing Ollama client: {e}")
+
+    async def astream(self, input_data: Union[str, List[Dict[str, str]], Any]) -> AsyncGenerator[str, None]:
+        """
+        True streaming method that yields tokens one by one as they arrive
+        
+        Args:
+            input_data: Can be:
+                - str: Simple text prompt
+                - list: Message history like [{"role": "user", "content": "hello"}]
+                - Any: LangChain message objects or other formats
+            
+        Yields:
+            Individual tokens as they arrive from the model
+        """
+        try:
+            # Ensure client is available
+            self._ensure_client()
+            
+            # Use adapter manager to prepare messages
+            messages = self._prepare_messages(input_data)
+            
+            # Prepare request parameters for streaming
+            provider_config = self.get_provider_config()
+            payload = {
+                "model": self.model_name,
+                "messages": messages,
+                "stream": True,  # Force streaming for astream
+                "options": {
+                    "temperature": provider_config.get("temperature", 0.7),
+                    "top_p": provider_config.get("top_p", 0.9),
+                    "num_predict": provider_config.get("max_tokens", 2048)
+                }
+            }
+            
+            # Add tools if bound using adapter manager
+            tool_schemas = await self._prepare_tools_for_request()
+            if tool_schemas:
+                payload["tools"] = tool_schemas
+            
+            # Stream tokens one by one
+            content_chunks = []
+            try:
+                async with self.client.stream("POST", "/api/chat", json=payload) as response:
+                    response.raise_for_status()
+                    async for line in response.aiter_lines():
+                        if line.strip():
+                            try:
+                                chunk = json.loads(line)
+                                if "message" in chunk and "content" in chunk["message"]:
+                                    content = chunk["message"]["content"]
+                                    if content:
+                                        content_chunks.append(content)
+                                        yield content
+                            except json.JSONDecodeError:
+                                continue
+                
+                # Track usage after streaming is complete (estimated)
+                full_content = "".join(content_chunks)
+                await self._track_streaming_usage(messages, full_content)
+                
+            except Exception as e:
+                logger.error(f"Error in streaming: {e}")
+                raise
+            
+        except httpx.RequestError as e:
+            logger.error(f"HTTP request error in astream: {e}")
+            raise
+        except Exception as e:
+            logger.error(f"Error in astream: {e}")
+            raise

isa_model/inference/services/llm/openai_llm_service.py

@@ -7,19 +7,18 @@ from typing import Dict, Any, List, Union, AsyncGenerator, Optional, Callable
 from openai import AsyncOpenAI
 
 from isa_model.inference.services.llm.base_llm_service import BaseLLMService
-from isa_model.inference.providers.base_provider import BaseProvider
-from isa_model.inference.billing_tracker import ServiceType
+from ....core.types import ServiceType
 
 logger = logging.getLogger(__name__)
 
 class OpenAILLMService(BaseLLMService):
     """OpenAI LLM service implementation with unified invoke interface"""
     
-    def __init__(self, provider: 'BaseProvider', model_name: str = "gpt-4.1-nano"):
-        super().__init__(provider, model_name)
+    def __init__(self, model_name: str = "gpt-4o-mini", provider_name: str = "openai", **kwargs):
+        super().__init__(provider_name, model_name, **kwargs)
         
-        # Get full configuration from provider (including sensitive data)
-        provider_config = provider.get_full_config()
+        # Get configuration from centralized config manager
+        provider_config = self.get_provider_config()
         
         # Initialize AsyncOpenAI client with provider configuration
         try:
@@ -28,7 +27,7 @@ class OpenAILLMService(BaseLLMService):
             
             self.client = AsyncOpenAI(
                 api_key=provider_config["api_key"],
-                base_url=provider_config.get("base_url", "https://api.openai.com/v1"),
+                base_url=provider_config.get("api_base_url", "https://api.openai.com/v1"),
                 organization=provider_config.get("organization")
             )
             
@@ -44,7 +43,7 @@ class OpenAILLMService(BaseLLMService):
     
     def _create_bound_copy(self) -> 'OpenAILLMService':
         """Create a copy of this service for tool binding"""
-        bound_service = OpenAILLMService(self.provider, self.model_name)
+        bound_service = OpenAILLMService(self.model_name, self.provider_name)
         bound_service._bound_tools = self._bound_tools.copy()
         return bound_service
     
@@ -67,6 +66,58 @@ class OpenAILLMService(BaseLLMService):
         
         return bound_service
     
+    async def astream(self, input_data: Union[str, List[Dict[str, str]], Any]) -> AsyncGenerator[str, None]:
+        """
+        True streaming method - yields tokens one by one as they arrive
+        
+        Args:
+            input_data: Same as ainvoke
+            
+        Yields:
+            Individual tokens as they arrive from the API
+        """
+        try:
+            # Use adapter manager to prepare messages
+            messages = self._prepare_messages(input_data)
+            
+            # Prepare request kwargs
+            provider_config = self.get_provider_config()
+            kwargs = {
+                "model": self.model_name,
+                "messages": messages,
+                "temperature": provider_config.get("temperature", 0.7),
+                "max_tokens": provider_config.get("max_tokens", 1024),
+                "stream": True
+            }
+            
+            # Add tools if bound using adapter manager
+            tool_schemas = await self._prepare_tools_for_request()
+            if tool_schemas:
+                kwargs["tools"] = tool_schemas
+                kwargs["tool_choice"] = "auto"
+            
+            # Stream tokens one by one
+            content_chunks = []
+            try:
+                stream = await self.client.chat.completions.create(**kwargs)
+                async for chunk in stream:
+                    content = chunk.choices[0].delta.content
+                    if content:
+                        content_chunks.append(content)
+                        yield content
+                
+                # Track usage after streaming is complete
+                full_content = "".join(content_chunks)
+                self._track_streaming_usage(messages, full_content)
+                
+            except Exception as e:
+                logger.error(f"Error in streaming: {e}")
+                raise
+            
+        except Exception as e:
+            logger.error(f"Error in astream: {e}")
+            raise
+    
     async def ainvoke(self, input_data: Union[str, List[Dict[str, str]], Any]) -> Union[str, Any]:
         """Unified invoke method for all input types"""
         try:
@@ -74,11 +125,12 @@ class OpenAILLMService(BaseLLMService):
             messages = self._prepare_messages(input_data)
             
             # Prepare request kwargs
+            provider_config = self.get_provider_config()
             kwargs = {
                 "model": self.model_name,
                 "messages": messages,
-                "temperature": self.config.get("temperature", 0.7),
-                "max_tokens": self.config.get("max_tokens", 1024)
+                "temperature": provider_config.get("temperature", 0.7),
+                "max_tokens": provider_config.get("max_tokens", 1024)
             }
             
             # Add tools if bound using adapter manager
@@ -89,23 +141,12 @@ class OpenAILLMService(BaseLLMService):
             
             # Handle streaming vs non-streaming
             if self.streaming:
-                # Streaming mode - collect all chunks
+                # TRUE STREAMING MODE - collect all chunks from the stream
                 content_chunks = []
-                async for chunk in await self._stream_response(kwargs):
-                    content_chunks.append(chunk)
+                async for token in self.astream(input_data):
+                    content_chunks.append(token)
                 content = "".join(content_chunks)
                 
-                # Create a mock usage object for tracking
-                class MockUsage:
-                    def __init__(self):
-                        self.prompt_tokens = len(str(messages)) // 4  # Rough estimate
-                        self.completion_tokens = len(content) // 4   # Rough estimate  
-                        self.total_tokens = self.prompt_tokens + self.completion_tokens
-                
-                usage = MockUsage()
-                self._update_token_usage(usage)
-                self._track_billing(usage)
-                
                 return self._format_response(content, input_data)
             else:
                 # Non-streaming mode
@@ -115,7 +156,7 @@ class OpenAILLMService(BaseLLMService):
                 # Update usage tracking
                 if response.usage:
                     self._update_token_usage(response.usage)
-                    self._track_billing(response.usage)
+                    await self._track_billing(response.usage)
                 
                 # Handle tool calls if present - let adapter process the complete message
                 if message.tool_calls:
@@ -129,9 +170,28 @@ class OpenAILLMService(BaseLLMService):
             logger.error(f"Error in ainvoke: {e}")
             raise
     
+    def _track_streaming_usage(self, messages: List[Dict[str, str]], content: str):
+        """Track usage for streaming requests (estimated)"""
+        # Create a mock usage object for tracking
+        class MockUsage:
+            def __init__(self):
+                self.prompt_tokens = len(str(messages)) // 4  # Rough estimate
+                self.completion_tokens = len(content) // 4   # Rough estimate  
+                self.total_tokens = self.prompt_tokens + self.completion_tokens
+        
+        usage = MockUsage()
+        self._update_token_usage(usage)
+        # Fire and forget async tracking
+        import asyncio
+        try:
+            loop = asyncio.get_event_loop()
+            loop.create_task(self._track_billing(usage))
+        except:
+            # If no event loop, skip tracking
+            pass
     
     async def _stream_response(self, kwargs: Dict[str, Any]) -> AsyncGenerator[str, None]:
-        """Handle streaming responses"""
+        """Handle streaming responses - DEPRECATED: Use astream() instead"""
         kwargs["stream"] = True
         
         async def stream_generator():
@@ -162,16 +222,17 @@ class OpenAILLMService(BaseLLMService):
         self.total_token_usage["total_tokens"] += self.last_token_usage["total_tokens"]
         self.total_token_usage["requests_count"] += 1
     
-    def _track_billing(self, usage):
+    async def _track_billing(self, usage):
         """Track billing information"""
-        self._track_usage(
+        provider_config = self.get_provider_config()
+        await self._track_usage(
             service_type=ServiceType.LLM,
             operation="chat",
             input_tokens=usage.prompt_tokens,
             output_tokens=usage.completion_tokens,
             metadata={
-                "temperature": self.config.get("temperature", 0.7), 
-                "max_tokens": self.config.get("max_tokens", 1024)
+                "temperature": provider_config.get("temperature", 0.7), 
+                "max_tokens": provider_config.get("max_tokens", 1024)
             }
         )
     
@@ -185,15 +246,57 @@ class OpenAILLMService(BaseLLMService):
     
     def get_model_info(self) -> Dict[str, Any]:
         """Get information about the current model"""
+        provider_config = self.get_provider_config()
         return {
             "name": self.model_name,
-            "max_tokens": self.config.get("max_tokens", 1024),
+            "max_tokens": provider_config.get("max_tokens", 1024),
             "supports_streaming": True,
             "supports_functions": True,
             "provider": "openai"
         }
     
         
+    async def chat(
+        self,
+        input_data: Union[str, List[Dict[str, str]], Any],
+        max_tokens: Optional[int] = None
+    ) -> Dict[str, Any]:
+        """
+        Chat method that wraps ainvoke for compatibility with base class
+        
+        Args:
+            input_data: Input messages
+            max_tokens: Maximum tokens to generate
+            
+        Returns:
+            Dict containing chat response
+        """
+        try:
+            # Call ainvoke and get the response
+            response = await self.ainvoke(input_data)
+            
+            # Return in expected format
+            return {
+                "text": response if isinstance(response, str) else str(response),
+                "success": True,
+                "metadata": {
+                    "model": self.model_name,
+                    "provider": self.provider_name,
+                    "max_tokens": max_tokens or self.max_tokens
+                }
+            }
+        except Exception as e:
+            logger.error(f"Chat method failed: {e}")
+            return {
+                "text": "",
+                "success": False,
+                "error": str(e),
+                "metadata": {
+                    "model": self.model_name,
+                    "provider": self.provider_name
+                }
+            }
+    
     async def close(self):
         """Close the backend client"""
         await self.client.close()