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

autobyteus/agent_team/bootstrap_steps/agent_configuration_preparation_step.py

@@ -15,8 +15,8 @@ logger = logging.getLogger(__name__)
 class AgentConfigurationPreparationStep(BaseAgentTeamBootstrapStep):
     """
     Bootstrap step to prepare the final, immutable configuration for every
-    agent in the team. It injects team-specific context and applies the final
-    coordinator prompt. It no longer injects tools.
+    agent in the team. It injects team-specific context, applies team-level
+    settings like tool format overrides, and prepares the final coordinator prompt.
     """
     async def execute(self, context: 'AgentTeamContext', phase_manager: 'AgentTeamPhaseManager') -> bool:
         team_id = context.team_id
@@ -44,6 +44,13 @@ class AgentConfigurationPreparationStep(BaseAgentTeamBootstrapStep):
                 
                 final_config = node_definition.copy()
 
+                # --- Team-level Setting Propagation ---
+                # If the team config specifies a tool format, it overrides any agent-level setting.
+                if context.config.use_xml_tool_format is not None:
+                    final_config.use_xml_tool_format = context.config.use_xml_tool_format
+                    logger.debug(f"Team '{team_id}': Applied team-level use_xml_tool_format={final_config.use_xml_tool_format} to agent '{unique_name}'.")
+
+
                 # --- Shared Context Injection ---
                 # The shared context is injected into the initial_custom_data dictionary,
                 # which is then used by the AgentFactory to create the AgentRuntimeState.

autobyteus/agent_team/context/agent_team_config.py

@@ -20,6 +20,7 @@ class AgentTeamConfig:
     coordinator_node: TeamNodeConfig
     role: Optional[str] = None
     task_notification_mode: TaskNotificationMode = TaskNotificationMode.AGENT_MANUAL_NOTIFICATION
+    use_xml_tool_format: Optional[bool] = None
 
     def __post_init__(self):
         if not self.name or not isinstance(self.name, str):

autobyteus/llm/api/autobyteus_llm.py

@@ -4,6 +4,7 @@ from autobyteus.llm.models import LLMModel
 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_client.client import AutobyteusClient
 import logging
 import uuid
@@ -12,36 +13,35 @@ logger = logging.getLogger(__name__)
 
 class AutobyteusLLM(BaseLLM):
     def __init__(self, model: LLMModel, llm_config: LLMConfig):
-        # The host URL is now passed via the model object.
         if not model.host_url:
             raise ValueError("AutobyteusLLM requires a host_url to be set in its LLMModel object.")
 
         super().__init__(model=model, llm_config=llm_config)
         
-        # Instantiate the client with the specific host for this model.
         self.client = AutobyteusClient(server_url=self.model.host_url)
         self.conversation_id = str(uuid.uuid4())
         logger.info(f"AutobyteusLLM initialized for model '{self.model.model_identifier}' with conversation ID: {self.conversation_id}")
 
     async def _send_user_message_to_llm(
         self,
-        user_message: str,
-        image_urls: Optional[List[str]] = None,
+        user_message: LLMUserMessage,
         **kwargs
     ) -> CompleteResponse:
         self.add_user_message(user_message)
         try:
             response = await self.client.send_message(
                 conversation_id=self.conversation_id,
-                model_name=self.model.name, # Use `name` as it's the original model name for the API
-                user_message=user_message,
-                image_urls=image_urls
+                model_name=self.model.name,
+                user_message=user_message.content,
+                image_urls=user_message.image_urls,
+                audio_urls=user_message.audio_urls,
+                video_urls=user_message.video_urls
             )
             
             assistant_message = response['response']
             self.add_assistant_message(assistant_message)
             
-            token_usage_data = response.get('token_usage', {})
+            token_usage_data = response.get('token_usage') or {}
             token_usage = TokenUsage(
                 prompt_tokens=token_usage_data.get('prompt_tokens', 0),
                 completion_tokens=token_usage_data.get('completion_tokens', 0),
@@ -59,8 +59,7 @@ class AutobyteusLLM(BaseLLM):
 
     async def _stream_user_message_to_llm(
         self,
-        user_message: str,
-        image_urls: Optional[List[str]] = None,
+        user_message: LLMUserMessage,
         **kwargs
     ) -> AsyncGenerator[ChunkResponse, None]:
         self.add_user_message(user_message)
@@ -69,36 +68,38 @@ class AutobyteusLLM(BaseLLM):
         try:
             async for chunk in self.client.stream_message(
                 conversation_id=self.conversation_id,
-                model_name=self.model.name, # Use `name` for the API call
-                user_message=user_message,
-                image_urls=image_urls
+                model_name=self.model.name,
+                user_message=user_message.content,
+                image_urls=user_message.image_urls,
+                audio_urls=user_message.audio_urls,
+                video_urls=user_message.video_urls
             ):
                 if 'error' in chunk:
                     raise RuntimeError(chunk['error'])
                 
                 content = chunk.get('content', '')
-                complete_response += content
+                if content:
+                    complete_response += content
+
                 is_complete = chunk.get('is_complete', False)
-                
-                # If this is the final chunk, include token usage
+                token_usage = None
                 if is_complete:
-                    token_usage = None
-                    if chunk.get('token_usage'):
-                        token_usage = TokenUsage(
-                            prompt_tokens=chunk['token_usage'].get('prompt_tokens', 0),
-                            completion_tokens=chunk['token_usage'].get('completion_tokens', 0),
-                            total_tokens=chunk['token_usage'].get('total_tokens', 0)
-                        )
-                    yield ChunkResponse(
-                        content=content,
-                        is_complete=True,
-                        usage=token_usage
-                    )
-                else:
-                    yield ChunkResponse(
-                        content=content,
-                        is_complete=False
+                    token_usage_data = chunk.get('token_usage') or {}
+                    token_usage = TokenUsage(
+                        prompt_tokens=token_usage_data.get('prompt_tokens', 0),
+                        completion_tokens=token_usage_data.get('completion_tokens', 0),
+                        total_tokens=token_usage_data.get('total_tokens', 0)
                     )
+
+                yield ChunkResponse(
+                    content=content,
+                    reasoning=chunk.get('reasoning'),
+                    is_complete=is_complete,
+                    image_urls=chunk.get('image_urls', []),
+                    audio_urls=chunk.get('audio_urls', []),
+                    video_urls=chunk.get('video_urls', []),
+                    usage=token_usage
+                )
             
             self.add_assistant_message(complete_response)
         except Exception as e:
@@ -116,7 +117,6 @@ class AutobyteusLLM(BaseLLM):
             await self.client.close()
 
     async def _handle_error_cleanup(self):
-        """Handle cleanup operations after errors"""
         try:
             await self.cleanup()
         except Exception as cleanup_error:

autobyteus/llm/api/bedrock_llm.py

@@ -9,10 +9,10 @@ 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
 
 class BedrockLLM(BaseLLM):
     def __init__(self, model: LLMModel = None, llm_config: LLMConfig = None):
-        # Provide defaults if not specified
         if model is None:
             model = LLMModel.BEDROCK_CLAUDE_3_5_SONNET_API
         if llm_config is None:
@@ -43,14 +43,17 @@ class BedrockLLM(BaseLLM):
         except Exception as e:
             raise ValueError(f"Failed to initialize Bedrock 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)
         
+        # NOTE: This implementation does not yet support multimodal inputs for Bedrock.
+        # It will only send the text content.
+        
         request_body = json.dumps({
             "anthropic_version": "bedrock-2023-05-31",
             "max_tokens": 1000,
             "temperature": 0,
-            "messages": [msg.to_dict() for msg in self.messages],
+            "messages": [msg.to_dict() for msg in self.messages if msg.role != MessageRole.SYSTEM],
             "system": self.system_message if self.system_message else ""
         })
 
@@ -79,6 +82,11 @@ class BedrockLLM(BaseLLM):
             raise ValueError(f"Bedrock API error: {error_code} - {error_message}")
         except Exception as e:
             raise ValueError(f"Error in Bedrock API call: {str(e)}")
-    
+
+    async def _stream_user_message_to_llm(self, user_message: LLMUserMessage, **kwargs) -> AsyncGenerator[ChunkResponse, None]:
+        # Placeholder for future implementation
+        response = await self._send_user_message_to_llm(user_message, **kwargs)
+        yield ChunkResponse(content=response.content, is_complete=True, usage=response.usage)
+
     async def cleanup(self):
-        super().cleanup()
+        await super().cleanup()

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()