quantalogic 0.35.0__py3-none-any.whl → 0.50.0__py3-none-any.whl

quantalogic/generative_model.py

@@ -1,24 +1,20 @@
 """Generative model module for AI-powered text generation."""
 
+import asyncio
 from datetime import datetime
-from typing import Any, Dict, List
+from typing import Any, AsyncGenerator, Dict, List
 
-import litellm
 import openai
-from litellm import exceptions
 from loguru import logger
 from pydantic import BaseModel, Field, field_validator
 
 from quantalogic.event_emitter import EventEmitter  # Importing the EventEmitter class
 from quantalogic.get_model_info import get_max_input_tokens, get_max_output_tokens, get_max_tokens
-from quantalogic.llm import count_tokens, generate_completion, generate_image
+from quantalogic.quantlitellm import acompletion, aimage_generation, exceptions, token_counter
 
 MIN_RETRIES = 1
 
 
-litellm.suppress_debug_info = True  # Very important to suppress prints don't remove
-
-
 # Define the Message class for conversation handling
 class Message(BaseModel):
     """Represents a message in a conversation with a specific role and content."""
@@ -76,7 +72,7 @@ class ResponseStats(BaseModel):
 
 
 class GenerativeModel:
-    """Generative model for AI-powered text generation and image generation."""
+    """Generative model for AI-powered text and image generation with async support."""
 
     def __init__(
         self,
@@ -121,26 +117,25 @@ class GenerativeModel:
         exceptions.PermissionDeniedError,
     )
 
-    # Generate a response with conversation history and optional streaming
-    def generate_with_history(
+    async def async_generate_with_history(
         self,
         messages_history: list[Message],
         prompt: str,
         image_url: str | None = None,
         streaming: bool = False,
         stop_words: list[str] | None = None,
-    ) -> ResponseStats:
-        """Generate a response with conversation history and optional image.
+    ) -> ResponseStats | AsyncGenerator[str, None]:
+        """Asynchronously generate a response with conversation history and optional image.
 
         Args:
             messages_history: Previous conversation messages.
             prompt: Current user prompt.
             image_url: Optional image URL for visual queries.
             streaming: Whether to stream the response.
-            stop_words: Optional list of stop words for streaming
+            stop_words: Optional list of stop words for streaming.
 
         Returns:
-            Detailed response statistics or a generator in streaming mode.
+            ResponseStats if streaming=False, or an AsyncGenerator for streaming=True.
         """
         messages = [{"role": msg.role, "content": str(msg.content)} for msg in messages_history]
 
@@ -158,59 +153,79 @@ class GenerativeModel:
             messages.append({"role": "user", "content": str(prompt)})
 
         if streaming:
-            self.event_emitter.emit("stream_start")  # Emit stream start event
-            return self._stream_response(messages)  # Return generator
+            self.event_emitter.emit("stream_start")
+            return self._async_stream_response(messages, stop_words)
 
         try:
-            logger.debug(f"Generating response for prompt: {prompt}")
-
-            response = generate_completion(
-                temperature=self.temperature,
+            logger.debug(f"Async generating response for prompt: {prompt} with messages: {messages}")
+            response = await acompletion(
                 model=self.model,
                 messages=messages,
+                temperature=self.temperature,
                 num_retries=MIN_RETRIES,
                 stop=stop_words,
                 extra_headers={"X-Title": "quantalogic"},
             )
+            logger.debug(f"Raw response from {self.model}: {response}")
+
+            # Check for error in response
+            if hasattr(response, "error") and response.error:
+                error_msg = response.error.get("message", "Unknown error")
+                logger.warning(f"API returned error: {error_msg}")
+                raise openai.APIError(
+                    message=f"API error: {error_msg}",
+                    request={"model": self.model, "messages": messages},
+                    body={"error": response.error},
+                )
 
             token_usage = TokenUsage(
                 prompt_tokens=response.usage.prompt_tokens,
                 completion_tokens=response.usage.completion_tokens,
                 total_tokens=response.usage.total_tokens,
             )
+            # Get the content with a check for None
+            content = response.choices[0].message.content
+            if content is None:
+                logger.warning(f"Received None content from {self.model}. Raw response: {response}")
+                raise ValueError(f"Model {self.model} returned no content for the given input.")
 
             return ResponseStats(
-                response=response.choices[0].message.content,
+                response=content,
                 usage=token_usage,
                 model=self.model,
                 finish_reason=response.choices[0].finish_reason,
             )
-
         except Exception as e:
             self._handle_generation_exception(e)
+            # We should never reach here as _handle_generation_exception always raises
 
-    def _stream_response(self, messages, stop_words: list[str] | None = None):
-        """Private method to handle streaming responses."""
+    async def _async_stream_response(self, messages, stop_words: list[str] | None = None):
+        """Private method to handle asynchronous streaming responses."""
         try:
-            for chunk in generate_completion(
-                temperature=self.temperature,
+            response = await acompletion(
                 model=self.model,
                 messages=messages,
-                num_retries=MIN_RETRIES,
-                stream=True,  # Enable streaming,
+                temperature=self.temperature,
+                stream=True,
                 stop=stop_words,
-            ):
+                num_retries=MIN_RETRIES,
+            )
+            async for chunk in response:
                 if chunk.choices[0].delta.content is not None:
                     self.event_emitter.emit("stream_chunk", chunk.choices[0].delta.content)
-                    yield chunk.choices[0].delta.content  # Yield each chunk of content
-
-            self.event_emitter.emit("stream_end")  # Emit stream end event
+                    yield chunk.choices[0].delta.content
+            self.event_emitter.emit("stream_end")
         except Exception as e:
-            logger.error(f"Streaming error: {str(e)}")
+            logger.error(f"Async streaming error: {str(e)}")
             raise
 
-    def generate(self, prompt: str, image_url: str | None = None, streaming: bool = False) -> ResponseStats:
-        """Generate a response without conversation history.
+    async def async_generate(
+        self,
+        prompt: str,
+        image_url: str | None = None,
+        streaming: bool = False,
+    ) -> ResponseStats | AsyncGenerator[str, None]:
+        """Asynchronously generate a response without conversation history.
 
         Args:
             prompt: User prompt.
@@ -218,9 +233,9 @@ class GenerativeModel:
             streaming: Whether to stream the response.
 
         Returns:
-            Detailed response statistics or a generator in streaming mode.
+            ResponseStats if streaming=False, or an AsyncGenerator for streaming=True.
         """
-        return self.generate_with_history([], prompt, image_url, streaming)
+        return await self.async_generate_with_history([], prompt, image_url, streaming)
 
     def _handle_generation_exception(self, e):
         """Handle exceptions during generation."""
@@ -238,35 +253,39 @@ class GenerativeModel:
 
         if isinstance(e, self.AUTH_EXCEPTIONS):
             logger.debug("Authentication error occurred")
-            raise openai.AuthenticationError(f"Authentication failed with provider {error_details['provider']}") from e
+            raise openai.AuthenticationError(
+                message=f"Authentication failed with provider {error_details['provider']}",
+                request={"model": self.model, "temperature": self.temperature},
+                body={"error": {"message": str(e), "type": "authentication_error"}},
+            ) from e
 
         if isinstance(e, self.CONTEXT_EXCEPTIONS):
-            raise openai.InvalidRequestError(f"Context window exceeded or invalid request: {str(e)}") from e
+            raise openai.InvalidRequestError(
+                message=f"Context window exceeded or invalid request: {str(e)}",
+                request={"model": self.model, "temperature": self.temperature},
+                body={"error": {"message": str(e), "type": "invalid_request_error"}},
+            ) from e
 
         if isinstance(e, self.POLICY_EXCEPTIONS):
-            raise openai.APIError(f"Content policy violation: {str(e)}") from e
+            raise openai.APIError(
+                message=f"Content policy violation: {str(e)}",
+                request={"model": self.model, "temperature": self.temperature},
+                body={"error": {"message": str(e), "type": "policy_violation"}},
+            ) from e
 
         if isinstance(e, openai.OpenAIError):
             raise
 
-        raise openai.APIError(f"Unexpected error during generation: {str(e)}") from e
+        raise openai.APIError(
+            message=f"Unexpected error during generation: {str(e)}",
+            request={"model": self.model, "temperature": self.temperature},
+            body={"error": {"message": str(e), "type": "unexpected_error"}},
+        ) from e
 
     def get_max_tokens(self) -> int:
         """Get the maximum number of tokens that can be generated by the model."""
         return get_max_tokens(self.model)
 
-    def token_counter(self, messages: list[Message]) -> int:
-        """Count the number of tokens in a list of messages."""
-        logger.debug(f"Counting tokens for {len(messages)} messages using model {self.model}")
-        litellm_messages = [{"role": msg.role, "content": str(msg.content)} for msg in messages]
-        return count_tokens(model=self.model, messages=litellm_messages)
-
-    def token_counter_with_history(self, messages_history: list[Message], prompt: str) -> int:
-        """Count the number of tokens in a list of messages and a prompt."""
-        litellm_messages = [{"role": msg.role, "content": str(msg.content)} for msg in messages_history]
-        litellm_messages.append({"role": "user", "content": str(prompt)})
-        return count_tokens(model=self.model, messages=litellm_messages)
-
     def get_model_max_input_tokens(self) -> int | None:
         """Get the maximum number of input tokens for the model."""
         return get_max_input_tokens(self.model)
@@ -275,35 +294,26 @@ class GenerativeModel:
         """Get the maximum number of output tokens for the model."""
         return get_max_output_tokens(self.model)
 
-    def generate_image(self, prompt: str, params: Dict[str, Any]) -> ResponseStats:
-        """Generate an image using the specified model and parameters.
+    async def async_generate_image(self, prompt: str, params: Dict[str, Any]) -> ResponseStats:
+        """Asynchronously generate an image using the specified model and parameters.
 
         Args:
-            prompt: Text description of the image to generate
-            params: Dictionary of parameters for image generation including:
-                   - model: Name of the image generation model
-                   - size: Size of the generated image
-                   - quality: Quality level (DALL-E only)
-                   - style: Style preference (DALL-E only)
-                   - response_format: Format of the response (url/base64)
-                   - negative_prompt: What to avoid in the image (SD only)
-                   - cfg_scale: Classifier Free Guidance scale (SD only)
+            prompt: Text description of the image to generate.
+            params: Dictionary of parameters for image generation.
 
         Returns:
-            ResponseStats containing the image generation results
+            ResponseStats containing the image generation results.
 
         Raises:
-            Exception: If there's an error during image generation
+            Exception: If there's an error during image generation.
         """
         try:
-            logger.debug(f"Generating image with params: {params}")
-
-            # Ensure prompt is in params
-            generation_params = {**params}
-            generation_params["prompt"] = prompt
+            logger.debug(f"Async generating image with params: {params}")
+            generation_params = {**params, "prompt": prompt}
+            model = generation_params.pop("model")
 
-            # Call litellm's image generation function
-            response = generate_image(model=generation_params.pop("model"), **generation_params)
+            # Check if litellm provides an async image generation method; if not, adapt sync
+            response = await aimage_generation(model=model, **generation_params)
 
             # Convert response data to list of dictionaries with string values
             if hasattr(response, "data"):
@@ -329,15 +339,37 @@ class GenerativeModel:
             else:
                 created = None
 
-            # Convert response to our ResponseStats format
             return ResponseStats(
-                response="",  # Empty for image generation
+                response="",
                 usage=TokenUsage(prompt_tokens=0, completion_tokens=0, total_tokens=0),
                 model=str(params["model"]),
                 data=data,
                 created=created,
             )
-
         except Exception as e:
-            logger.error(f"Error in image generation: {str(e)}")
+            logger.error(f"Error in async image generation: {str(e)}")
             raise
+
+    def token_counter(self, messages: list[Message]) -> int:
+        """Count the number of tokens in a list of messages."""
+        logger.debug(f"Counting tokens for {len(messages)} messages using model {self.model}")
+        litellm_messages = [{"role": msg.role, "content": str(msg.content)} for msg in messages]
+        return token_counter(model=self.model, messages=litellm_messages)
+
+    def token_counter_with_history(self, messages_history: list[Message], prompt: str) -> int:
+        """Count the number of tokens in a list of messages and a prompt."""
+        litellm_messages = [{"role": msg.role, "content": str(msg.content)} for msg in messages_history]
+        litellm_messages.append({"role": "user", "content": str(prompt)})
+        return token_counter(model=self.model, messages=litellm_messages)
+
+    async def async_token_counter(self, messages: list[Message]) -> int:
+        """Asynchronously count the number of tokens in a list of messages."""
+        logger.debug(f"Async counting tokens for {len(messages)} messages using model {self.model}")
+        litellm_messages = [{"role": msg.role, "content": str(msg.content)} for msg in messages]
+        return await asyncio.to_thread(token_counter, model=self.model, messages=litellm_messages)
+
+    async def async_token_counter_with_history(self, messages_history: list[Message], prompt: str) -> int:
+        """Asynchronously count the number of tokens in a list of messages and a prompt."""
+        litellm_messages = [{"role": msg.role, "content": str(msg.content)} for msg in messages_history]
+        litellm_messages.append({"role": "user", "content": str(prompt)})
+        return await asyncio.to_thread(token_counter, model=self.model, messages=litellm_messages)

quantalogic/get_model_info.py

@@ -24,11 +24,11 @@ def get_max_output_tokens(model_name: str) -> int:
     """Get max output tokens with safe fallback"""
     validate_model_name(model_name)
 
-    if model_name.startswith('lm_studio/'):
+    if model_name.startswith("lm_studio/"):
         try:
             models = get_model_list()
             for model in models.data:
-                if model.id == model_name[len('lm_studio/'):]:
+                if model.id == model_name[len("lm_studio/") :]:
                     return model.max_context_length
         except Exception:
             loguru.logger.warning(f"Could not fetch LM Studio model info for {model_name}, using default")
@@ -38,7 +38,7 @@ def get_max_output_tokens(model_name: str) -> int:
 
     try:
         return litellm_get_model_max_output_tokens(model_name)
-    except Exception as e:
+    except Exception:
         loguru.logger.warning(f"Model {model_name} not found in LiteLLM registry, using default")
         return DEFAULT_MAX_OUTPUT_TOKENS
 
@@ -47,11 +47,11 @@ def get_max_input_tokens(model_name: str) -> int:
     """Get max input tokens with safe fallback"""
     validate_model_name(model_name)
 
-    if model_name.startswith('lm_studio/'):
+    if model_name.startswith("lm_studio/"):
         try:
             models = get_model_list()
             for model in models.data:
-                if model.id == model_name[len('lm_studio/'):]:
+                if model.id == model_name[len("lm_studio/") :]:
                     return model.max_context_length
         except Exception:
             loguru.logger.warning(f"Could not fetch LM Studio model info for {model_name}, using default")