autobyteus 1.1.5__py3-none-any.whl → 1.1.7__py3-none-any.whl

autobyteus/llm/api/claude_llm.py

@@ -8,14 +8,14 @@ from autobyteus.llm.utils.llm_config import LLMConfig
 from autobyteus.llm.utils.messages import MessageRole, Message
 from autobyteus.llm.utils.token_usage import TokenUsage
 from autobyteus.llm.utils.response_types import CompleteResponse, ChunkResponse
+from autobyteus.llm.user_message import LLMUserMessage
 
 logger = logging.getLogger(__name__)
 
 class ClaudeLLM(BaseLLM):
     def __init__(self, model: LLMModel = None, llm_config: LLMConfig = None):
-        # Provide defaults if not specified
         if model is None:
-            model = LLMModel.CLAUDE_3_5_SONNET_API
+            model = LLMModel['claude-4-sonnet']
         if llm_config is None:
             llm_config = LLMConfig()
             
@@ -37,22 +37,22 @@ class ClaudeLLM(BaseLLM):
             raise ValueError(f"Failed to initialize Anthropic client: {str(e)}")
     
     def _get_non_system_messages(self) -> List[Dict]:
-        """
-        Returns all messages excluding system messages for Anthropic API compatibility.
-        """
+        # NOTE: This will need to be updated to handle multimodal messages for Claude
         return [msg.to_dict() for msg in self.messages if msg.role != MessageRole.SYSTEM]
     
     def _create_token_usage(self, input_tokens: int, output_tokens: int) -> TokenUsage:
-        """Convert Anthropic usage data to TokenUsage format."""
         return TokenUsage(
             prompt_tokens=input_tokens,
             completion_tokens=output_tokens,
             total_tokens=input_tokens + output_tokens
         )
     
-    async def _send_user_message_to_llm(self, user_message: str, image_urls: Optional[List[str]] = None, **kwargs) -> CompleteResponse:
+    async def _send_user_message_to_llm(self, user_message: LLMUserMessage, **kwargs) -> CompleteResponse:
         self.add_user_message(user_message)
 
+        # NOTE: This implementation does not yet support multimodal inputs for Claude.
+        # It will only send the text content.
+
         try:
             response = self.client.messages.create(
                 model=self.model.value,
@@ -81,12 +81,15 @@ class ClaudeLLM(BaseLLM):
             raise ValueError(f"Error in Claude API call: {str(e)}")
     
     async def _stream_user_message_to_llm(
-        self, user_message: str, image_urls: Optional[List[str]] = None, **kwargs
+        self, user_message: LLMUserMessage, **kwargs
     ) -> AsyncGenerator[ChunkResponse, None]:
         self.add_user_message(user_message)
         complete_response = ""
         final_message = None
 
+        # NOTE: This implementation does not yet support multimodal inputs for Claude.
+        # It will only send the text content.
+
         try:
             with self.client.messages.stream(
                 model=self.model.value,
@@ -96,30 +99,13 @@ class ClaudeLLM(BaseLLM):
                 messages=self._get_non_system_messages(),
             ) as stream:
                 for event in stream:
-                    logger.debug(f"Event Received: {event.type}")
-                    
-                    if event.type == "message_start":
-                        logger.debug(f"Message Start: {event.message}")
-                    
-                    elif event.type == "content_block_start":
-                        logger.debug(f"Content Block Start at index {event.index}: {event.content_block}")
-                    
-                    elif event.type == "content_block_delta" and event.delta.type == "text_delta":
-                        logger.debug(f"Text Delta: {event.delta.text}")
+                    if event.type == "content_block_delta" and event.delta.type == "text_delta":
                         complete_response += event.delta.text
                         yield ChunkResponse(
                             content=event.delta.text,
                             is_complete=False
                         )
                     
-                    elif event.type == "message_delta":
-                        logger.debug(f"Message Delta: Stop Reason - {event.delta.stop_reason}, "
-                                   f"Stop Sequence - {event.delta.stop_sequence}")
-                    
-                    elif event.type == "content_block_stop":
-                        logger.debug(f"Content Block Stop at index {event.index}: {event.content_block}")
-
-                # Get final message for token usage
                 final_message = stream.get_final_message()
                 if final_message:
                     token_usage = self._create_token_usage(
@@ -140,4 +126,4 @@ class ClaudeLLM(BaseLLM):
             raise ValueError(f"Error in Claude API streaming: {str(e)}")
     
     async def cleanup(self):
-        super().cleanup()
+        await super().cleanup()

autobyteus/llm/api/gemini_llm.py

@@ -1,6 +1,6 @@
 import logging
-from typing import Dict, Optional, List, AsyncGenerator
-import google.generativeai as genai
+from typing import Dict, List, AsyncGenerator, Any
+import google.generativeai as genai  # CHANGED: Using the older 'google.generativeai' library
 import os
 from autobyteus.llm.models import LLMModel
 from autobyteus.llm.base_llm import BaseLLM
@@ -8,68 +8,93 @@ from autobyteus.llm.utils.llm_config import LLMConfig
 from autobyteus.llm.utils.messages import MessageRole, Message
 from autobyteus.llm.utils.token_usage import TokenUsage
 from autobyteus.llm.utils.response_types import CompleteResponse, ChunkResponse
+from autobyteus.llm.user_message import LLMUserMessage
 
 logger = logging.getLogger(__name__)
 
+def _format_gemini_history(messages: List[Message]) -> List[Dict[str, Any]]:
+    """
+    Formats internal message history for the Gemini API.
+    This function remains compatible with the older library.
+    """
+    history = []
+    # System message is handled separately in the model initialization
+    for msg in messages:
+        if msg.role in [MessageRole.USER, MessageRole.ASSISTANT]:
+            role = 'model' if msg.role == MessageRole.ASSISTANT else 'user'
+            history.append({"role": role, "parts": [{"text": msg.content}]})
+    return history
+
 class GeminiLLM(BaseLLM):
     def __init__(self, model: LLMModel = None, llm_config: LLMConfig = None):
-        self.generation_config = {
-            "temperature": 0,
-            "top_p": 0.95,
-            "top_k": 64,
-            "max_output_tokens": 8192,
-            "response_mime_type": "text/plain",
-        }
-        
-        # Provide defaults if not specified
         if model is None:
-            model = LLMModel.GEMINI_1_5_FLASH_API
+            model = LLMModel['gemini-2.5-flash'] # Note: Ensure model name is compatible, e.g., 'gemini-1.5-flash-latest'
         if llm_config is None:
             llm_config = LLMConfig()
-            
+
         super().__init__(model=model, llm_config=llm_config)
-        self.client = self.initialize()
-        self.chat_session = None
+        
+        # CHANGED: Initialization flow. Configure API key and then instantiate the model.
+        self.initialize()
+        
+        system_instruction = self.system_message if self.system_message else None
+        
+        self.model = genai.GenerativeModel(
+            model_name=self.model.value,
+            system_instruction=system_instruction
+        )
 
-    @classmethod
-    def initialize(cls):
+    @staticmethod
+    def initialize():
+        """
+        CHANGED: This method now configures the genai library with the API key
+        instead of creating a client instance.
+        """
         api_key = os.environ.get("GEMINI_API_KEY")
         if not api_key:
             logger.error("GEMINI_API_KEY environment variable is not set.")
-            raise ValueError(
-                "GEMINI_API_KEY environment variable is not set. "
-                "Please set this variable in your environment."
-            )
+            raise ValueError("GEMINI_API_KEY environment variable is not set.")
         try:
             genai.configure(api_key=api_key)
-            return genai
         except Exception as e:
-            logger.error(f"Failed to initialize Gemini client: {str(e)}")
-            raise ValueError(f"Failed to initialize Gemini client: {str(e)}")
+            logger.error(f"Failed to configure Gemini client: {str(e)}")
+            raise ValueError(f"Failed to configure Gemini client: {str(e)}")
 
-    def _ensure_chat_session(self):
-        if not self.chat_session:
-            model = self.client.GenerativeModel(
-                model_name=self.model.value,
-                generation_config=self.generation_config
-            )
-            history = []
-            for msg in self.messages:
-                history.append({"role": msg.role.value, "parts": [msg.content]})
-            self.chat_session = model.start_chat(history=history)
+    def _get_generation_config(self) -> Dict[str, Any]:
+        """
+        CHANGED: Builds the generation config as a dictionary.
+        'thinking_config' is not available in the old library.
+        'system_instruction' is passed during model initialization.
+        """
+        # Basic configuration, you can expand this with temperature, top_p, etc.
+        # from self.llm_config if needed.
+        config = {
+            "response_mime_type": "text/plain",
+            # Example: "temperature": self.llm_config.temperature
+        }
+        return config
 
-    async def _send_user_message_to_llm(self, user_message: str, image_urls: Optional[List[str]] = None, **kwargs) -> CompleteResponse:
+    async def _send_user_message_to_llm(self, user_message: LLMUserMessage, **kwargs) -> CompleteResponse:
         self.add_user_message(user_message)
+        
         try:
-            self._ensure_chat_session()
-            response = self.chat_session.send_message(user_message)
+            history = _format_gemini_history(self.messages)
+            generation_config = self._get_generation_config()
+
+            # CHANGED: API call now uses the model instance directly.
+            response = await self.model.generate_content_async(
+                contents=history,
+                generation_config=generation_config,
+            )
+            
             assistant_message = response.text
             self.add_assistant_message(assistant_message)
             
+            # CHANGED: Token usage is extracted from 'usage_metadata'.
             token_usage = TokenUsage(
-                prompt_tokens=0,
-                completion_tokens=0,
-                total_tokens=0
+                prompt_tokens=response.usage_metadata.prompt_token_count,
+                completion_tokens=response.usage_metadata.candidates_token_count,
+                total_tokens=response.usage_metadata.total_token_count
             )
             
             return CompleteResponse(
@@ -80,6 +105,47 @@ class GeminiLLM(BaseLLM):
             logger.error(f"Error in Gemini API call: {str(e)}")
             raise ValueError(f"Error in Gemini API call: {str(e)}")
     
+    async def _stream_user_message_to_llm(self, user_message: LLMUserMessage, **kwargs) -> AsyncGenerator[ChunkResponse, None]:
+        self.add_user_message(user_message)
+        complete_response = ""
+        
+        try:
+            history = _format_gemini_history(self.messages)
+            generation_config = self._get_generation_config()
+
+            # CHANGED: API call for streaming is now part of generate_content_async.
+            response_stream = await self.model.generate_content_async(
+                contents=history,
+                generation_config=generation_config,
+                stream=True
+            )
+
+            async for chunk in response_stream:
+                chunk_text = chunk.text
+                complete_response += chunk_text
+                yield ChunkResponse(
+                    content=chunk_text,
+                    is_complete=False
+                )
+
+            self.add_assistant_message(complete_response)
+
+            # NOTE: The old library's async stream does not easily expose token usage.
+            # Keeping it at 0, consistent with your original implementation.
+            token_usage = TokenUsage(
+                prompt_tokens=0,
+                completion_tokens=0,
+                total_tokens=0
+            )
+
+            yield ChunkResponse(
+                content="",
+                is_complete=True,
+                usage=token_usage
+            )
+        except Exception as e:
+            logger.error(f"Error in Gemini API streaming call: {str(e)}")
+            raise ValueError(f"Error in Gemini API streaming call: {str(e)}")
+
     async def cleanup(self):
-        self.chat_session = None
-        super().cleanup()
+        await super().cleanup()

autobyteus/llm/api/groq_llm.py

@@ -7,6 +7,7 @@ from autobyteus.llm.utils.llm_config import LLMConfig
 from autobyteus.llm.utils.messages import MessageRole, Message
 from autobyteus.llm.utils.token_usage import TokenUsage
 from autobyteus.llm.utils.response_types import CompleteResponse, ChunkResponse
+from autobyteus.llm.user_message import LLMUserMessage
 
 logger = logging.getLogger(__name__)
 
@@ -36,7 +37,7 @@ class GroqLLM(BaseLLM):
         except Exception as e:
             raise ValueError(f"Failed to initialize Groq client: {str(e)}")
     
-    async def _send_user_message_to_llm(self, user_message: str, image_urls: Optional[List[str]] = None, **kwargs) -> CompleteResponse:
+    async def _send_user_message_to_llm(self, user_message: LLMUserMessage, **kwargs) -> CompleteResponse:
         self.add_user_message(user_message)
         try:
             # Placeholder for sending message to Groq API
@@ -58,7 +59,7 @@ class GroqLLM(BaseLLM):
             raise ValueError(f"Error in Groq API call: {str(e)}")
     
     async def _stream_user_message_to_llm(
-        self, user_message: str, image_urls: Optional[List[str]] = None, **kwargs
+        self, user_message: LLMUserMessage, **kwargs
     ) -> AsyncGenerator[ChunkResponse, None]:
         self.add_user_message(user_message)
         complete_response = ""
@@ -90,4 +91,4 @@ class GroqLLM(BaseLLM):
             raise ValueError(f"Error in Groq API streaming: {str(e)}")
     
     async def cleanup(self):
-        super().cleanup()
+        await super().cleanup()

autobyteus/llm/api/mistral_llm.py

@@ -1,45 +1,91 @@
-from typing import Dict, Optional, List, AsyncGenerator
+from typing import Dict, Optional, List, Any, AsyncGenerator, Union
 import os
 import logging
+import httpx
+import asyncio
 from autobyteus.llm.models import LLMModel
 from autobyteus.llm.base_llm import BaseLLM
 from mistralai import Mistral
-from autobyteus.llm.utils.messages import MessageRole, Message
+from autobyteus.llm.utils.messages import Message, MessageRole
 from autobyteus.llm.utils.llm_config import LLMConfig
 from autobyteus.llm.utils.token_usage import TokenUsage
 from autobyteus.llm.utils.response_types import CompleteResponse, ChunkResponse
+from autobyteus.llm.user_message import LLMUserMessage
+from autobyteus.llm.utils.media_payload_formatter import image_source_to_base64, get_mime_type, is_valid_image_path
 
-# Configure logger
 logger = logging.getLogger(__name__)
 
+async def _format_mistral_messages(messages: List[Message]) -> List[Dict[str, Any]]:
+    """Formats a list of internal Message objects into a list of dictionaries for the Mistral API."""
+    mistral_messages = []
+    for msg in messages:
+        # Skip empty messages from non-system roles as Mistral API may reject them
+        if not msg.content and not msg.image_urls and msg.role != MessageRole.SYSTEM:
+            continue
+
+        content: Union[str, List[Dict[str, Any]]]
+
+        if msg.image_urls:
+            content_parts: List[Dict[str, Any]] = []
+            if msg.content:
+                content_parts.append({"type": "text", "text": msg.content})
+
+            image_tasks = [image_source_to_base64(url) for url in msg.image_urls]
+            try:
+                base64_images = await asyncio.gather(*image_tasks)
+                for i, b64_image in enumerate(base64_images):
+                    original_url = msg.image_urls[i]
+                    mime_type = get_mime_type(original_url) if is_valid_image_path(original_url) else "image/jpeg"
+                    data_uri = f"data:{mime_type};base64,{b64_image}"
+                    
+                    # Mistral's format for image parts
+                    content_parts.append({
+                        "type": "image_url",
+                        "image_url": {
+                            "url": data_uri
+                        }
+                    })
+            except Exception as e:
+                logger.error(f"Error processing images for Mistral: {e}")
+            
+            if msg.audio_urls:
+                logger.warning("MistralLLM does not yet support audio; skipping.")
+            if msg.video_urls:
+                logger.warning("MistralLLM does not yet support video; skipping.")
+            
+            content = content_parts
+        else:
+            content = msg.content or ""
+        
+        mistral_messages.append({"role": msg.role.value, "content": content})
+            
+    return mistral_messages
+
+
 class MistralLLM(BaseLLM):
     def __init__(self, model: LLMModel = None, llm_config: LLMConfig = None):
-        # Provide defaults if not specified
         if model is None:
-            model = LLMModel.mistral_large
+            model = LLMModel['mistral-large']
         if llm_config is None:
             llm_config = LLMConfig()
             
         super().__init__(model=model, llm_config=llm_config)
-        self.client = self.initialize()
+        self.http_client = httpx.AsyncClient()
+        self.client: Mistral = self._initialize()
         logger.info(f"MistralLLM initialized with model: {self.model}")
 
-    @classmethod
-    def initialize(cls):
+    def _initialize(self) -> Mistral:
         mistral_api_key = os.environ.get("MISTRAL_API_KEY")
         if not mistral_api_key:
             logger.error("MISTRAL_API_KEY environment variable is not set")
-            raise ValueError(
-                "MISTRAL_API_KEY environment variable is not set. "
-                "Please set this variable in your environment."
-            )
+            raise ValueError("MISTRAL_API_KEY environment variable is not set.")
         try:
-            return Mistral(api_key=mistral_api_key)
+            return Mistral(api_key=mistral_api_key, client=self.http_client)
         except Exception as e:
             logger.error(f"Failed to initialize Mistral client: {str(e)}")
             raise ValueError(f"Failed to initialize Mistral client: {str(e)}")
 
-    def _create_token_usage(self, usage_data: Dict) -> TokenUsage:
+    def _create_token_usage(self, usage_data: Any) -> TokenUsage:
         """Convert Mistral usage data to TokenUsage format."""
         return TokenUsage(
             prompt_tokens=usage_data.prompt_tokens,
@@ -48,26 +94,26 @@ class MistralLLM(BaseLLM):
         )
 
     async def _send_user_message_to_llm(
-        self, user_message: str, image_urls: Optional[List[str]] = None, **kwargs
+        self, user_message: LLMUserMessage, **kwargs
     ) -> CompleteResponse:
         self.add_user_message(user_message)
-
+        
         try:
-            mistral_messages = [msg.to_mistral_message() for msg in self.messages]
+            mistral_messages = await _format_mistral_messages(self.messages)
             
-            chat_response = self.client.chat.complete(
+            chat_response = await self.client.chat.complete_async(
                 model=self.model.value,
                 messages=mistral_messages,
+                temperature=self.config.temperature,
+                max_tokens=self.config.max_tokens,
+                top_p=self.config.top_p,
             )
 
-            assistant_message = chat_response.choices.message.content
+            assistant_message = chat_response.choices[0].message.content
             self.add_assistant_message(assistant_message)
 
-            # Create token usage if available
-            token_usage = None
-            if hasattr(chat_response, 'usage') and chat_response.usage:
-                token_usage = self._create_token_usage(chat_response.usage)
-                logger.debug(f"Token usage recorded: {token_usage}")
+            token_usage = self._create_token_usage(chat_response.usage)
+            logger.debug(f"Token usage recorded: {token_usage}")
 
             return CompleteResponse(
                 content=assistant_message,
@@ -78,48 +124,48 @@ class MistralLLM(BaseLLM):
             raise ValueError(f"Error in Mistral API call: {str(e)}")
     
     async def _stream_user_message_to_llm(
-        self, user_message: str, image_urls: Optional[List[str]] = None, **kwargs
+        self, user_message: LLMUserMessage, **kwargs
     ) -> AsyncGenerator[ChunkResponse, None]:
         self.add_user_message(user_message)
         
+        accumulated_message = ""
+        final_usage = None
+        
         try:
-            mistral_messages = [msg.to_mistral_message() for msg in self.messages]
-            
-            stream = await self.client.chat.stream_async(
+            mistral_messages = await _format_mistral_messages(self.messages)
+
+            stream = self.client.chat.stream_async(
                 model=self.model.value,
                 messages=mistral_messages,
+                temperature=self.config.temperature,
+                max_tokens=self.config.max_tokens,
+                top_p=self.config.top_p,
             )
 
-            accumulated_message = ""
-            
             async for chunk in stream:
-                if chunk.data.choices.delta.content is not None:
-                    token = chunk.data.choices.delta.content
+                if chunk.choices and chunk.choices[0].delta.content is not None:
+                    token = chunk.choices[0].delta.content
                     accumulated_message += token
                     
-                    # For intermediate chunks, yield without usage
-                    yield ChunkResponse(
-                        content=token,
-                        is_complete=False
-                    )
-
-                # Check if this is the last chunk with usage data
-                if hasattr(chunk.data, 'usage') and chunk.data.usage is not None:
-                    token_usage = self._create_token_usage(chunk.data.usage)
-                    yield ChunkResponse(
-                        content="",
-                        is_complete=True,
-                        usage=token_usage
-                    )
-
-            # After streaming is complete, store the full message
+                    yield ChunkResponse(content=token, is_complete=False)
+
+                if hasattr(chunk, 'usage') and chunk.usage:
+                    final_usage = self._create_token_usage(chunk.usage)
+
+            # Yield the final chunk with usage data
+            yield ChunkResponse(
+                content="",
+                is_complete=True,
+                usage=final_usage
+            )
+            
             self.add_assistant_message(accumulated_message)
         except Exception as e:
             logger.error(f"Error in Mistral API streaming call: {str(e)}")
             raise ValueError(f"Error in Mistral API streaming call: {str(e)}")
     
     async def cleanup(self):
-        # Clean up any resources if needed
         logger.debug("Cleaning up MistralLLM instance")
-        self.messages = []
-        super().cleanup()
+        if self.http_client and not self.http_client.is_closed:
+            await self.http_client.aclose()
+        await super().cleanup()

autobyteus/llm/api/nvidia_llm.py

@@ -8,6 +8,7 @@ from autobyteus.llm.utils.llm_config import LLMConfig
 from autobyteus.llm.utils.messages import MessageRole, Message
 from autobyteus.llm.utils.token_usage import TokenUsage
 from autobyteus.llm.utils.response_types import CompleteResponse, ChunkResponse
+from autobyteus.llm.user_message import LLMUserMessage
 
 logger = logging.getLogger(__name__)
 
@@ -38,11 +39,11 @@ class NvidiaLLM(BaseLLM):
         except Exception as e:
             raise ValueError(f"Failed to initialize Nvidia client: {str(e)}")
     
-    async def _send_user_message_to_llm(self, user_message: str, image_urls: Optional[List[str]] = None, **kwargs) -> CompleteResponse:
+    async def _send_user_message_to_llm(self, user_message: LLMUserMessage, **kwargs) -> CompleteResponse:
         self.add_user_message(user_message)
         try:
             completion = self.client.chat.completions.create(
-                model=self.model,
+                model=self.model.value,
                 messages=[msg.to_dict() for msg in self.messages],
                 temperature=0,
                 top_p=1,
@@ -65,12 +66,12 @@ class NvidiaLLM(BaseLLM):
         except Exception as e:
             raise ValueError(f"Error in Nvidia API call: {str(e)}")
     
-    async def stream_response(self, user_message: str) -> AsyncGenerator[ChunkResponse, None]:
+    async def _stream_user_message_to_llm(self, user_message: LLMUserMessage, **kwargs) -> AsyncGenerator[ChunkResponse, None]:
         self.add_user_message(user_message)
         complete_response = ""
         try:
             completion = self.client.chat.completions.create(
-                model=self.model,
+                model=self.model.value,
                 messages=[msg.to_dict() for msg in self.messages],
                 temperature=0,
                 top_p=1,
@@ -104,4 +105,4 @@ class NvidiaLLM(BaseLLM):
             raise ValueError(f"Error in Nvidia API streaming call: {str(e)}")
     
     async def cleanup(self):
-        super().cleanup()
+        await super().cleanup()