abstractcore 2.5.3__py3-none-any.whl → 2.6.0__py3-none-any.whl

abstractcore/providers/lmstudio_provider.py

@@ -5,7 +5,7 @@ LM Studio provider implementation (OpenAI-compatible API).
 import httpx
 import json
 import time
-from typing import List, Dict, Any, Optional, Union, Iterator, Type
+from typing import List, Dict, Any, Optional, Union, Iterator, AsyncIterator, Type
 
 try:
     from pydantic import BaseModel
@@ -15,7 +15,7 @@ except ImportError:
     BaseModel = None
 from .base import BaseProvider
 from ..core.types import GenerateResponse
-from ..exceptions import ProviderAPIError, ModelNotFoundError, format_model_error
+from ..exceptions import ProviderAPIError, ModelNotFoundError, format_model_error, format_provider_error
 from ..tools import UniversalToolHandler, execute_tools
 from ..events import EventType
 
@@ -47,9 +47,21 @@ class LMStudioProvider(BaseProvider):
             except Exception:
                 raise RuntimeError(f"Failed to create HTTP client for LMStudio: {e}")
 
+        self._async_client = None  # Lazy-loaded async client
+
         # Validate model exists in LMStudio
         self._validate_model()
 
+    @property
+    def async_client(self):
+        """Lazy-load async HTTP client for native async operations."""
+        if self._async_client is None:
+            timeout_value = getattr(self, '_timeout', None)
+            if timeout_value is not None and timeout_value <= 0:
+                timeout_value = None
+            self._async_client = httpx.AsyncClient(timeout=timeout_value)
+        return self._async_client
+
     def _validate_model(self):
         """Validate that the model exists in LMStudio"""
         try:
@@ -87,6 +99,17 @@ class LMStudioProvider(BaseProvider):
             if hasattr(self, 'client') and self.client is not None:
                 self.client.close()
 
+            # Close async client if it was created
+            if self._async_client is not None:
+                import asyncio
+                try:
+                    loop = asyncio.get_running_loop()
+                    loop.create_task(self._async_client.aclose())
+                except RuntimeError:
+                    # No running loop
+                    import asyncio
+                    asyncio.run(self._async_client.aclose())
+
         except Exception as e:
             # Log but don't raise - unload should be best-effort
             if hasattr(self, 'logger'):
@@ -359,6 +382,227 @@ class LMStudioProvider(BaseProvider):
                 finish_reason="error"
             )
 
+    async def _agenerate_internal(self,
+                                   prompt: str,
+                                   messages: Optional[List[Dict[str, str]]] = None,
+                                   system_prompt: Optional[str] = None,
+                                   tools: Optional[List[Dict[str, Any]]] = None,
+                                   media: Optional[List['MediaContent']] = None,
+                                   stream: bool = False,
+                                   response_model: Optional[Type[BaseModel]] = None,
+                                   execute_tools: Optional[bool] = None,
+                                   tool_call_tags: Optional[str] = None,
+                                   **kwargs) -> Union[GenerateResponse, AsyncIterator[GenerateResponse]]:
+        """Native async implementation using httpx.AsyncClient - 3-10x faster for batch operations."""
+
+        # Build messages for chat completions with tool support (same logic as sync)
+        chat_messages = []
+
+        # Add tools to system prompt if provided
+        enhanced_system_prompt = system_prompt
+        if tools and self.tool_handler.supports_prompted:
+            tool_prompt = self.tool_handler.format_tools_prompt(tools)
+            if enhanced_system_prompt:
+                enhanced_system_prompt += f"\n\n{tool_prompt}"
+            else:
+                enhanced_system_prompt = tool_prompt
+
+        # Add system message if provided
+        if enhanced_system_prompt:
+            chat_messages.append({
+                "role": "system",
+                "content": enhanced_system_prompt
+            })
+
+        # Add conversation history
+        if messages:
+            chat_messages.extend(messages)
+
+        # Handle media content
+        if media:
+            user_message_text = prompt.strip() if prompt else ""
+            if not user_message_text and chat_messages:
+                for msg in reversed(chat_messages):
+                    if msg.get("role") == "user" and msg.get("content"):
+                        user_message_text = msg["content"]
+                        break
+            try:
+                processed_media = self._process_media_content(media)
+                media_handler = self._get_media_handler_for_model(self.model)
+                multimodal_message = media_handler.create_multimodal_message(user_message_text, processed_media)
+
+                if isinstance(multimodal_message, str):
+                    if chat_messages and chat_messages[-1].get("role") == "user":
+                        chat_messages[-1]["content"] = multimodal_message
+                    else:
+                        chat_messages.append({"role": "user", "content": multimodal_message})
+                else:
+                    if chat_messages and chat_messages[-1].get("role") == "user":
+                        chat_messages[-1] = multimodal_message
+                    else:
+                        chat_messages.append(multimodal_message)
+            except ImportError:
+                self.logger.warning("Media processing not available. Install with: pip install abstractcore[media]")
+                if user_message_text:
+                    chat_messages.append({"role": "user", "content": user_message_text})
+            except Exception as e:
+                self.logger.warning(f"Failed to process media content: {e}")
+                if user_message_text:
+                    chat_messages.append({"role": "user", "content": user_message_text})
+
+        # Add prompt as separate message if provided
+        elif prompt and prompt.strip():
+            chat_messages.append({"role": "user", "content": prompt})
+
+        # Build request payload
+        generation_kwargs = self._prepare_generation_kwargs(**kwargs)
+        max_output_tokens = self._get_provider_max_tokens_param(generation_kwargs)
+
+        payload = {
+            "model": self.model,
+            "messages": chat_messages,
+            "stream": stream,
+            "temperature": kwargs.get("temperature", self.temperature),
+            "max_tokens": max_output_tokens,
+            "top_p": kwargs.get("top_p", 0.9),
+        }
+
+        # Add additional parameters
+        if "frequency_penalty" in kwargs:
+            payload["frequency_penalty"] = kwargs["frequency_penalty"]
+        if "presence_penalty" in kwargs:
+            payload["presence_penalty"] = kwargs["presence_penalty"]
+        if "repetition_penalty" in kwargs:
+            payload["repetition_penalty"] = kwargs["repetition_penalty"]
+
+        # Add seed if provided
+        seed_value = kwargs.get("seed", self.seed)
+        if seed_value is not None:
+            payload["seed"] = seed_value
+
+        # Add structured output support
+        if response_model and PYDANTIC_AVAILABLE:
+            json_schema = response_model.model_json_schema()
+            payload["response_format"] = {
+                "type": "json_schema",
+                "json_schema": {
+                    "name": response_model.__name__,
+                    "schema": json_schema
+                }
+            }
+
+        if stream:
+            return self._async_stream_generate(payload)
+        else:
+            response = await self._async_single_generate(payload)
+
+            # Execute tools if enabled
+            if self.execute_tools and tools and self.tool_handler.supports_prompted and response.content:
+                response = self._handle_prompted_tool_execution(response, tools, execute_tools)
+
+            return response
+
+    async def _async_single_generate(self, payload: Dict[str, Any]) -> GenerateResponse:
+        """Native async single response generation."""
+        try:
+            # Track generation time
+            start_time = time.time()
+            response = await self.async_client.post(
+                f"{self.base_url}/chat/completions",
+                json=payload,
+                headers={"Content-Type": "application/json"}
+            )
+            response.raise_for_status()
+            gen_time = round((time.time() - start_time) * 1000, 1)
+
+            result = response.json()
+
+            # Extract response from OpenAI format
+            if "choices" in result and len(result["choices"]) > 0:
+                choice = result["choices"][0]
+                content = choice.get("message", {}).get("content", "")
+                finish_reason = choice.get("finish_reason", "stop")
+            else:
+                content = "No response generated"
+                finish_reason = "error"
+
+            # Extract usage info
+            usage = result.get("usage", {})
+
+            return GenerateResponse(
+                content=content,
+                model=self.model,
+                finish_reason=finish_reason,
+                raw_response=result,
+                usage={
+                    "input_tokens": usage.get("prompt_tokens", 0),
+                    "output_tokens": usage.get("completion_tokens", 0),
+                    "total_tokens": usage.get("total_tokens", 0),
+                    "prompt_tokens": usage.get("prompt_tokens", 0),
+                    "completion_tokens": usage.get("completion_tokens", 0)
+                },
+                gen_time=gen_time
+            )
+
+        except Exception as e:
+            error_str = str(e).lower()
+            if ('404' in error_str or 'not found' in error_str or 'model' in error_str) and ('not found' in error_str):
+                try:
+                    available_models = self.list_available_models(base_url=self.base_url)
+                    error_message = format_model_error("LMStudio", self.model, available_models)
+                    raise ModelNotFoundError(error_message)
+                except Exception:
+                    raise ModelNotFoundError(f"Model '{self.model}' not found in LMStudio")
+            else:
+                raise ProviderAPIError(f"LMStudio API error: {str(e)}")
+
+    async def _async_stream_generate(self, payload: Dict[str, Any]) -> AsyncIterator[GenerateResponse]:
+        """Native async streaming response generation."""
+        try:
+            async with self.async_client.stream(
+                "POST",
+                f"{self.base_url}/chat/completions",
+                json=payload,
+                headers={"Content-Type": "application/json"}
+            ) as response:
+                response.raise_for_status()
+
+                async for line in response.aiter_lines():
+                    if line:
+                        line = line.strip()
+
+                        if line.startswith("data: "):
+                            data = line[6:]  # Remove "data: " prefix
+
+                            if data == "[DONE]":
+                                break
+
+                            try:
+                                chunk = json.loads(data)
+
+                                if "choices" in chunk and len(chunk["choices"]) > 0:
+                                    choice = chunk["choices"][0]
+                                    delta = choice.get("delta", {})
+                                    content = delta.get("content", "")
+                                    finish_reason = choice.get("finish_reason")
+
+                                    yield GenerateResponse(
+                                        content=content,
+                                        model=self.model,
+                                        finish_reason=finish_reason,
+                                        raw_response=chunk
+                                    )
+
+                            except json.JSONDecodeError:
+                                continue
+
+        except Exception as e:
+            yield GenerateResponse(
+                content=f"Error: {str(e)}",
+                model=self.model,
+                finish_reason="error"
+            )
+
     def get_capabilities(self) -> List[str]:
         """Get LM Studio capabilities"""
         return ["streaming", "chat", "tools"]

abstractcore/providers/ollama_provider.py

@@ -5,7 +5,7 @@ Ollama provider implementation.
 import json
 import httpx
 import time
-from typing import List, Dict, Any, Optional, Union, Iterator, Type
+from typing import List, Dict, Any, Optional, Union, Iterator, AsyncIterator, Type
 
 try:
     from pydantic import BaseModel
@@ -15,7 +15,7 @@ except ImportError:
     BaseModel = None
 from .base import BaseProvider
 from ..core.types import GenerateResponse
-from ..exceptions import ProviderAPIError, ModelNotFoundError, format_model_error
+from ..exceptions import ProviderAPIError, ModelNotFoundError, format_model_error, format_provider_error
 from ..tools import UniversalToolHandler, ToolDefinition, execute_tools
 from ..events import EventType
 
@@ -29,10 +29,21 @@ class OllamaProvider(BaseProvider):
 
         self.base_url = base_url.rstrip('/')
         self.client = httpx.Client(timeout=self._timeout)
+        self._async_client = None  # Lazy-loaded async client
 
         # Initialize tool handler
         self.tool_handler = UniversalToolHandler(model)
 
+    @property
+    def async_client(self):
+        """Lazy-load async HTTP client for native async operations."""
+        if self._async_client is None:
+            self._async_client = httpx.AsyncClient(
+                base_url=self.base_url,
+                timeout=self._timeout
+            )
+        return self._async_client
+
     def unload(self) -> None:
         """
         Unload the model from Ollama server memory.
@@ -59,6 +70,17 @@ class OllamaProvider(BaseProvider):
             if hasattr(self, 'client') and self.client is not None:
                 self.client.close()
 
+            # Close async client if it was created
+            if self._async_client is not None:
+                import asyncio
+                try:
+                    loop = asyncio.get_running_loop()
+                    loop.create_task(self._async_client.aclose())
+                except RuntimeError:
+                    # No running loop, close synchronously
+                    import asyncio
+                    asyncio.run(self._async_client.aclose())
+
         except Exception as e:
             # Log but don't raise - unload should be best-effort
             if hasattr(self, 'logger'):
@@ -377,6 +399,226 @@ class OllamaProvider(BaseProvider):
                 finish_reason="error"
             )
 
+    async def _agenerate_internal(self,
+                                   prompt: str,
+                                   messages: Optional[List[Dict]],
+                                   system_prompt: Optional[str],
+                                   tools: Optional[List],
+                                   media: Optional[List],
+                                   stream: bool,
+                                   **kwargs):
+        """Native async implementation using httpx.AsyncClient - 3-10x faster for batch operations."""
+        # Handle tools for prompted models
+        effective_system_prompt = system_prompt
+        if tools and self.tool_handler.supports_prompted:
+            tool_prompt = self.tool_handler.format_tools_prompt(tools)
+            if effective_system_prompt:
+                effective_system_prompt = f"{effective_system_prompt}\n\n{tool_prompt}"
+            else:
+                effective_system_prompt = tool_prompt
+
+        # Build request payload (same logic as sync)
+        generation_kwargs = self._prepare_generation_kwargs(**kwargs)
+        max_output_tokens = self._get_provider_max_tokens_param(generation_kwargs)
+        response_model = kwargs.get('response_model')
+
+        payload = {
+            "model": self.model,
+            "stream": stream,
+            "options": {
+                "temperature": kwargs.get("temperature", self.temperature),
+                "num_predict": max_output_tokens,
+            }
+        }
+
+        seed_value = kwargs.get("seed", self.seed)
+        if seed_value is not None:
+            payload["options"]["seed"] = seed_value
+
+        # Add structured output support
+        if response_model and PYDANTIC_AVAILABLE:
+            json_schema = response_model.model_json_schema()
+            payload["format"] = json_schema
+
+        # Use chat format
+        use_chat_format = tools is not None or messages is not None or True
+
+        if use_chat_format:
+            payload["messages"] = []
+
+            if effective_system_prompt:
+                payload["messages"].append({
+                    "role": "system",
+                    "content": effective_system_prompt
+                })
+
+            if messages:
+                converted_messages = self._convert_messages_for_ollama(messages)
+                payload["messages"].extend(converted_messages)
+
+            if media:
+                user_message_text = prompt.strip() if prompt else ""
+                try:
+                    from ..media.handlers import LocalMediaHandler
+                    media_handler = LocalMediaHandler("ollama", self.model_capabilities, model_name=self.model)
+                    multimodal_message = media_handler.create_multimodal_message(user_message_text, media)
+
+                    if isinstance(multimodal_message, str):
+                        payload["messages"].append({"role": "user", "content": multimodal_message})
+                    else:
+                        payload["messages"].append(multimodal_message)
+                except Exception as e:
+                    if hasattr(self, 'logger'):
+                        self.logger.warning(f"Failed to process media: {e}")
+                    if user_message_text:
+                        payload["messages"].append({"role": "user", "content": user_message_text})
+
+            elif prompt and prompt.strip():
+                payload["messages"].append({"role": "user", "content": prompt})
+
+            endpoint = "/api/chat"
+        else:
+            full_prompt = prompt
+            if effective_system_prompt:
+                full_prompt = f"{effective_system_prompt}\n\n{prompt}"
+            payload["prompt"] = full_prompt
+            endpoint = "/api/generate"
+
+        if stream:
+            return self._async_stream_generate(endpoint, payload, tools, kwargs.get('tool_call_tags'))
+        else:
+            return await self._async_single_generate(endpoint, payload, tools, kwargs.get('media_metadata'))
+
+    async def _async_single_generate(self, endpoint: str, payload: Dict[str, Any],
+                                      tools: Optional[List[Dict[str, Any]]] = None,
+                                      media_metadata: Optional[List[Dict[str, Any]]] = None) -> GenerateResponse:
+        """Native async single response generation."""
+        try:
+            start_time = time.time()
+            response = await self.async_client.post(endpoint, json=payload)
+            response.raise_for_status()
+            gen_time = round((time.time() - start_time) * 1000, 1)
+
+            result = response.json()
+
+            if endpoint == "/api/chat":
+                content = result.get("message", {}).get("content", "")
+            else:
+                content = result.get("response", "")
+
+            generate_response = GenerateResponse(
+                content=content,
+                model=self.model,
+                finish_reason="stop",
+                raw_response=result,
+                usage={
+                    "input_tokens": result.get("prompt_eval_count", 0),
+                    "output_tokens": result.get("eval_count", 0),
+                    "total_tokens": result.get("prompt_eval_count", 0) + result.get("eval_count", 0),
+                    "prompt_tokens": result.get("prompt_eval_count", 0),
+                    "completion_tokens": result.get("eval_count", 0)
+                },
+                gen_time=gen_time
+            )
+
+            if media_metadata:
+                if not generate_response.metadata:
+                    generate_response.metadata = {}
+                generate_response.metadata['media_metadata'] = media_metadata
+
+            if self.execute_tools and tools and self.tool_handler.supports_prompted and content:
+                return self._handle_tool_execution(generate_response, tools)
+
+            return generate_response
+
+        except Exception as e:
+            error_str = str(e).lower()
+            if ('404' in error_str or 'not found' in error_str):
+                available_models = self.list_available_models(base_url=self.base_url)
+                error_message = format_model_error("Ollama", self.model, available_models)
+                raise ModelNotFoundError(error_message)
+            else:
+                return GenerateResponse(
+                    content=f"Error: {str(e)}",
+                    model=self.model,
+                    finish_reason="error"
+                )
+
+    async def _async_stream_generate(self, endpoint: str, payload: Dict[str, Any],
+                                      tools: Optional[List[Dict[str, Any]]] = None,
+                                      tool_call_tags: Optional[str] = None):
+        """Native async streaming response generation."""
+        try:
+            async with self.async_client.stream("POST", endpoint, json=payload) as response:
+                response.raise_for_status()
+
+                full_content = ""
+                rewriter = None
+                buffer = ""
+                if tool_call_tags:
+                    try:
+                        from ..tools.tag_rewriter import create_tag_rewriter
+                        rewriter = create_tag_rewriter(tool_call_tags)
+                    except ImportError:
+                        pass
+
+                async for line in response.aiter_lines():
+                    if line:
+                        try:
+                            chunk = json.loads(line)
+
+                            if endpoint == "/api/chat":
+                                content = chunk.get("message", {}).get("content", "")
+                            else:
+                                content = chunk.get("response", "")
+
+                            done = chunk.get("done", False)
+                            full_content += content
+
+                            if rewriter and content:
+                                rewritten_content, buffer = rewriter.rewrite_streaming_chunk(content, buffer)
+                                content = rewritten_content
+
+                            chunk_response = GenerateResponse(
+                                content=content,
+                                model=self.model,
+                                finish_reason="stop" if done else None,
+                                raw_response=chunk
+                            )
+
+                            yield chunk_response
+
+                            if done:
+                                break
+
+                        except json.JSONDecodeError:
+                            continue
+
+                # Execute tools if enabled
+                if self.execute_tools and tools and self.tool_handler.supports_prompted and full_content:
+                    complete_response = GenerateResponse(
+                        content=full_content,
+                        model=self.model,
+                        finish_reason="stop"
+                    )
+
+                    final_response = self._handle_tool_execution(complete_response, tools)
+
+                    if final_response.content != full_content:
+                        tool_results_content = final_response.content[len(full_content):]
+                        yield GenerateResponse(
+                            content=tool_results_content,
+                            model=self.model,
+                            finish_reason="stop"
+                        )
+
+        except Exception as e:
+            yield GenerateResponse(
+                content=f"Error: {str(e)}",
+                model=self.model,
+                finish_reason="error"
+            )
+
     def _handle_tool_execution(self, response: GenerateResponse, tools: List[Dict[str, Any]]) -> GenerateResponse:
         """Handle tool execution for prompted models"""
         # Parse tool calls from response