autobyteus 1.2.1__py3-none-any.whl → 1.2.3__py3-none-any.whl

autobyteus/llm/api/zhipu_llm.py

@@ -1,16 +1,33 @@
 import logging
-from typing import Optional
+from typing import Optional, Dict, Any
 from autobyteus.llm.models import LLMModel
 from autobyteus.llm.utils.llm_config import LLMConfig
 from autobyteus.llm.api.openai_compatible_llm import OpenAICompatibleLLM
 
 logger = logging.getLogger(__name__)
 
+def _normalize_zhipu_extra_params(extra_params: Dict[str, Any]) -> Dict[str, Any]:
+    if not extra_params:
+        return {}
+
+    params = dict(extra_params)
+    thinking_type = params.pop("thinking_type", None)
+
+    if thinking_type is not None:
+        thinking = dict(params.get("thinking") or {})
+
+        if thinking_type is not None:
+            thinking["type"] = thinking_type
+
+        params["thinking"] = thinking
+
+    return params
+
 class ZhipuLLM(OpenAICompatibleLLM):
     def __init__(self, model: LLMModel = None, llm_config: LLMConfig = None):
         # Provide defaults if not specified
         if model is None:
-            model = LLMModel['glm-4.6']
+            model = LLMModel['glm-4.7']
         if llm_config is None:
             llm_config = LLMConfig()
             
@@ -20,6 +37,8 @@ class ZhipuLLM(OpenAICompatibleLLM):
             api_key_env_var="ZHIPU_API_KEY",
             base_url="https://open.bigmodel.cn/api/paas/v4/"
         )
+        if self.config and isinstance(self.config.extra_params, dict):
+            self.config.extra_params = _normalize_zhipu_extra_params(self.config.extra_params)
         logger.info(f"ZhipuLLM initialized with model: {self.model}")
 
     async def cleanup(self):

autobyteus/llm/autobyteus_provider.py

@@ -32,76 +32,86 @@ class AutobyteusModelProvider:
             
         return []
 
+    @staticmethod
+    def get_models() -> List[LLMModel]:
+        """
+        Fetches models from all configured Autobyteus hosts and returns them as LLMModel objects.
+        """
+        hosts = AutobyteusModelProvider._get_hosts()
+        if not hosts:
+            logger.info("No Autobyteus LLM server hosts configured. Skipping Autobyteus LLM model discovery.")
+            return []
+
+        all_models = []
+
+        for host_url in hosts:
+            if not AutobyteusModelProvider.is_valid_url(host_url):
+                logger.error(f"Invalid Autobyteus host URL: {host_url}, skipping.")
+                continue
+            
+            logger.info(f"Discovering Autobyteus models from host: {host_url}")
+            client = None
+            try:
+                # Instantiate client for this specific host
+                client = AutobyteusClient(server_url=host_url)
+                response = client.get_available_llm_models_sync()
+            except Exception as e:
+                logger.warning(f"Could not connect or fetch models from Autobyteus server at {host_url}: {e}")
+                continue
+            finally:
+                if client:
+                    client.sync_client.close()
+
+            if not AutobyteusModelProvider._validate_server_response(response):
+                continue
+
+            models = response.get('models', [])
+            for model_info in models:
+                try:
+                    validation_result = AutobyteusModelProvider._validate_model_info(model_info)
+                    if not validation_result["valid"]:
+                        logger.warning(validation_result["message"])
+                        continue
+                    
+                    llm_config = AutobyteusModelProvider._parse_llm_config(model_info["config"])
+                    if not llm_config:
+                        continue
+                    
+                    llm_model = LLMModel(
+                        name=model_info["name"],
+                        value=model_info["value"],
+                        provider=LLMProvider(model_info["provider"]),
+                        llm_class=AutobyteusLLM,
+                        canonical_name=model_info["canonical_name"],
+                        runtime=LLMRuntime.AUTOBYTEUS,
+                        host_url=host_url,
+                        default_config=llm_config
+                    )
+                    all_models.append(llm_model)
+                    
+                except Exception as e:
+                    logger.error(f"Failed to create LLMModel for '{model_info.get('name')}' from {host_url}: {e}")
+        
+        return all_models
+
     @staticmethod
     def discover_and_register():
         """Discover and register Autobyteus models from all configured hosts."""
         try:
             from autobyteus.llm.llm_factory import LLMFactory
 
-            hosts = AutobyteusModelProvider._get_hosts()
-            if not hosts:
-                logger.info("No Autobyteus LLM server hosts configured. Skipping Autobyteus LLM model discovery.")
-                return
+            discovered_models = AutobyteusModelProvider.get_models()
+            registered_count = 0
 
-            total_registered_count = 0
-
-            for host_url in hosts:
-                if not AutobyteusModelProvider.is_valid_url(host_url):
-                    logger.error(f"Invalid Autobyteus host URL: {host_url}, skipping.")
-                    continue
-                
-                logger.info(f"Discovering Autobyteus models from host: {host_url}")
-                client = None
+            for model in discovered_models:
                 try:
-                    # Instantiate client for this specific host
-                    client = AutobyteusClient(server_url=host_url)
-                    response = client.get_available_llm_models_sync()
+                    LLMFactory.register_model(model)
+                    registered_count += 1
                 except Exception as e:
-                    logger.warning(f"Could not connect or fetch models from Autobyteus server at {host_url}: {e}")
-                    continue
-                finally:
-                    if client:
-                        client.sync_client.close()
-
-                if not AutobyteusModelProvider._validate_server_response(response):
-                    continue
-
-                models = response.get('models', [])
-                host_registered_count = 0
-                for model_info in models:
-                    try:
-                        validation_result = AutobyteusModelProvider._validate_model_info(model_info)
-                        if not validation_result["valid"]:
-                            logger.warning(validation_result["message"])
-                            continue
-                        
-                        llm_config = AutobyteusModelProvider._parse_llm_config(model_info["config"])
-                        if not llm_config:
-                            continue
-                        
-                        llm_model = LLMModel(
-                            name=model_info["name"],
-                            value=model_info["value"],
-                            provider=LLMProvider(model_info["provider"]),
-                            llm_class=AutobyteusLLM,
-                            canonical_name=model_info["canonical_name"],
-                            runtime=LLMRuntime.AUTOBYTEUS,
-                            host_url=host_url,
-                            default_config=llm_config
-                        )
-                        
-                        LLMFactory.register_model(llm_model)
-                        host_registered_count += 1
-                        
-                    except Exception as e:
-                        logger.error(f"Failed to register Autobyteus model '{model_info.get('name')}' from {host_url}: {e}")
-                
-                if host_registered_count > 0:
-                    logger.info(f"Registered {host_registered_count} models from Autobyteus host {host_url}")
-                total_registered_count += host_registered_count
+                    logger.warning(f"Failed to register Autobyteus model '{model.name}': {e}")
             
-            if total_registered_count > 0:
-                 logger.info(f"Finished Autobyteus discovery. Total models registered: {total_registered_count}")
+            if registered_count > 0:
+                 logger.info(f"Finished Autobyteus discovery. Total models registered: {registered_count}")
 
         except Exception as e:
             logger.error(f"An unexpected error occurred during Autobyteus model discovery: {e}", exc_info=True)

autobyteus/llm/base_llm.py

@@ -26,12 +26,13 @@ class BaseLLM(ABC):
 
         self._token_usage_extension: TokenUsageTrackingExtension = self.register_extension(TokenUsageTrackingExtension)
 
-        self.messages: List[Message] = []
         self.system_message = self.config.system_message or self.DEFAULT_SYSTEM_MESSAGE
-        self.add_system_message(self.system_message)
 
     @property
     def latest_token_usage(self):
+        """Get latest token usage. Returns None if token tracking is disabled."""
+        if not self._token_usage_extension.is_enabled:
+            return None
         return self._token_usage_extension.latest_token_usage
 
     def register_extension(self, extension_class: Type[LLMExtension]) -> LLMExtension:
@@ -45,42 +46,19 @@ class BaseLLM(ABC):
     def get_extension(self, extension_class: Type[LLMExtension]) -> Optional[LLMExtension]:
         return self._extension_registry.get(extension_class)
 
-    def add_system_message(self, message: str):
-        self.messages.append(Message(MessageRole.SYSTEM, content=message))
-
-    def add_user_message(self, user_message: LLMUserMessage):
-        """
-        Adds a user message to history, converting from LLMUserMessage to Message.
-        """
-        msg = Message(
+    def _build_user_message(self, user_message: LLMUserMessage) -> Message:
+        return Message(
             role=MessageRole.USER,
             content=user_message.content,
             image_urls=user_message.image_urls,
             audio_urls=user_message.audio_urls,
-            video_urls=user_message.video_urls
-        )
-        self.messages.append(msg)
-        self._trigger_on_user_message_added(msg)
-
-    def add_assistant_message(self,
-                              content: Optional[str],
-                              reasoning_content: Optional[str] = None,
-                              image_urls: Optional[List[str]] = None,
-                              audio_urls: Optional[List[str]] = None,
-                              video_urls: Optional[List[str]] = None):
-        """
-        Adds a multimodal assistant message to the conversation history.
-        """
-        msg = Message(
-            role=MessageRole.ASSISTANT,
-            content=content,
-            reasoning_content=reasoning_content,
-            image_urls=image_urls,
-            audio_urls=audio_urls,
-            video_urls=video_urls
+            video_urls=user_message.video_urls,
         )
-        self.messages.append(msg)
-        self._trigger_on_assistant_message_added(msg)
+
+    def _build_system_message(self) -> Optional[Message]:
+        if not self.system_message:
+            return None
+        return Message(MessageRole.SYSTEM, content=self.system_message)
 
     def configure_system_prompt(self, new_system_prompt: str):
         if not new_system_prompt or not isinstance(new_system_prompt, str):
@@ -89,51 +67,40 @@ class BaseLLM(ABC):
 
         self.system_message = new_system_prompt
         self.config.system_message = new_system_prompt
-
-        system_message_found = False
-        for i, msg in enumerate(self.messages):
-            if msg.role == MessageRole.SYSTEM:
-                self.messages[i] = Message(MessageRole.SYSTEM, new_system_prompt)
-                system_message_found = True
-                logging.debug(f"Replaced existing system message at index {i}.")
-                break
-        
-        if not system_message_found:
-            self.messages.insert(0, Message(MessageRole.SYSTEM, new_system_prompt))
-            logging.debug("No existing system message found, inserted new one at the beginning.")
-        
         logging.info(f"LLM instance system prompt updated. New prompt length: {len(new_system_prompt)}")
 
-    def _trigger_on_user_message_added(self, message: Message):
+    async def _execute_before_hooks(self, messages: List[Message], rendered_payload: Optional[dict] = None, **kwargs) -> None:
         for extension in self._extension_registry.get_all():
-            extension.on_user_message_added(message)
+            await extension.before_invoke(messages, rendered_payload, **kwargs)
 
-    def _trigger_on_assistant_message_added(self, message: Message):
+    async def _execute_after_hooks(self, messages: List[Message], response: CompleteResponse = None, **kwargs) -> None:
         for extension in self._extension_registry.get_all():
-            extension.on_assistant_message_added(message)
-
-    async def _execute_before_hooks(self, user_message: LLMUserMessage, **kwargs) -> None:
-        for extension in self._extension_registry.get_all():
-            await extension.before_invoke(user_message, **kwargs)
-
-    async def _execute_after_hooks(self, user_message: LLMUserMessage, response: CompleteResponse = None, **kwargs) -> None:
-        for extension in self._extension_registry.get_all():
-            await extension.after_invoke(user_message, response, **kwargs)
-
-    async def send_user_message(self, user_message: LLMUserMessage, **kwargs) -> CompleteResponse:
-        await self._execute_before_hooks(user_message, **kwargs)
-        response = await self._send_user_message_to_llm(user_message, **kwargs)
-        await self._execute_after_hooks(user_message, response, **kwargs)
+            await extension.after_invoke(messages, response, **kwargs)
+
+    async def send_messages(
+        self,
+        messages: List[Message],
+        rendered_payload: Optional[dict] = None,
+        **kwargs,
+    ) -> CompleteResponse:
+        await self._execute_before_hooks(messages, rendered_payload, **kwargs)
+        response = await self._send_messages_to_llm(messages, **kwargs)
+        await self._execute_after_hooks(messages, response, **kwargs)
         return response
 
-    async def stream_user_message(self, user_message: LLMUserMessage, **kwargs) -> AsyncGenerator[ChunkResponse, None]:
-        await self._execute_before_hooks(user_message, **kwargs)
+    async def stream_messages(
+        self,
+        messages: List[Message],
+        rendered_payload: Optional[dict] = None,
+        **kwargs,
+    ) -> AsyncGenerator[ChunkResponse, None]:
+        await self._execute_before_hooks(messages, rendered_payload, **kwargs)
 
         accumulated_content = ""
         accumulated_reasoning = ""
         final_chunk = None
-        
-        async for chunk in self._stream_user_message_to_llm(user_message, **kwargs):
+
+        async for chunk in self._stream_messages_to_llm(messages, **kwargs):
             if chunk.content:
                 accumulated_content += chunk.content
             if chunk.reasoning:
@@ -146,34 +113,72 @@ class BaseLLM(ABC):
         complete_response = CompleteResponse(
             content=accumulated_content,
             reasoning=accumulated_reasoning if accumulated_reasoning else None,
-            usage=final_chunk.usage if final_chunk else None
+            usage=final_chunk.usage if final_chunk else None,
         )
-        
-        await self._execute_after_hooks(user_message, complete_response, **kwargs)
 
-    @abstractmethod
+        await self._execute_after_hooks(messages, complete_response, **kwargs)
+
+    async def send_user_message(self, user_message: LLMUserMessage, **kwargs) -> CompleteResponse:
+        messages: List[Message] = []
+        system_message = self._build_system_message()
+        if system_message:
+            messages.append(system_message)
+        messages.append(self._build_user_message(user_message))
+        return await self.send_messages(messages, **kwargs)
+
+    async def stream_user_message(
+        self, user_message: LLMUserMessage, **kwargs
+    ) -> AsyncGenerator[ChunkResponse, None]:
+        messages: List[Message] = []
+        system_message = self._build_system_message()
+        if system_message:
+            messages.append(system_message)
+        messages.append(self._build_user_message(user_message))
+        async for chunk in self.stream_messages(messages, **kwargs):
+            yield chunk
+
     async def _send_user_message_to_llm(self, user_message: LLMUserMessage, **kwargs) -> CompleteResponse:
+        messages: List[Message] = []
+        system_message = self._build_system_message()
+        if system_message:
+            messages.append(system_message)
+        messages.append(self._build_user_message(user_message))
+        return await self._send_messages_to_llm(messages, **kwargs)
+
+    async def _stream_user_message_to_llm(
+        self, user_message: LLMUserMessage, **kwargs
+    ) -> AsyncGenerator[ChunkResponse, None]:
+        messages: List[Message] = []
+        system_message = self._build_system_message()
+        if system_message:
+            messages.append(system_message)
+        messages.append(self._build_user_message(user_message))
+        async for chunk in self._stream_messages_to_llm(messages, **kwargs):
+            yield chunk
+
+    @abstractmethod
+    async def _send_messages_to_llm(self, messages: List[Message], **kwargs) -> CompleteResponse:
         """
-        Abstract method for sending a user message to an LLM. Must be implemented by subclasses.
-        
+        Abstract method for sending a list of messages to an LLM. Must be implemented by subclasses.
+
         Args:
-            user_message (LLMUserMessage): The user message object.
+            messages (List[Message]): The message list to send.
             **kwargs: Additional arguments for LLM-specific usage.
-            
+
         Returns:
             CompleteResponse: The complete response from the LLM.
         """
         pass
 
     @abstractmethod
-    async def _stream_user_message_to_llm(self, user_message: LLMUserMessage, **kwargs) -> AsyncGenerator[ChunkResponse, None]:
+    async def _stream_messages_to_llm(self, messages: List[Message], **kwargs) -> AsyncGenerator[ChunkResponse, None]:
         """
-        Abstract method for streaming a user message response from the LLM. Must be implemented by subclasses.
-        
+        Abstract method for streaming a response from an LLM. Must be implemented by subclasses.
+
         Args:
-            user_message (LLMUserMessage): The user message object.
+            messages (List[Message]): The message list to send.
             **kwargs: Additional arguments for LLM-specific usage.
-            
+
         Yields:
             AsyncGenerator[ChunkResponse, None]: Streaming chunks from the LLM response.
         """
@@ -183,4 +188,3 @@ class BaseLLM(ABC):
         for extension in self._extension_registry.get_all():
             await extension.cleanup()
         self._extension_registry.clear()
-        self.messages = []

autobyteus/llm/converters/__init__.py

@@ -0,0 +1,14 @@
+# file: autobyteus/autobyteus/llm/converters/__init__.py
+"""
+LLM provider-specific converters.
+
+These converters transform provider-specific data formats into
+normalized internal representations.
+"""
+
+from .openai_tool_call_converter import convert_openai_tool_calls
+from .gemini_tool_call_converter import convert_gemini_tool_calls
+from .anthropic_tool_call_converter import convert_anthropic_tool_call
+from .mistral_tool_call_converter import convert_mistral_tool_calls
+
+__all__ = ["convert_openai_tool_calls", "convert_gemini_tool_calls", "convert_anthropic_tool_call", "convert_mistral_tool_calls"]

autobyteus/llm/converters/anthropic_tool_call_converter.py

@@ -0,0 +1,37 @@
+from typing import List, Optional, Any
+from autobyteus.llm.utils.tool_call_delta import ToolCallDelta
+
+# Anthropic event types are typically simple objects or dicts from the SDK stream
+# The SDK yields `MessageStreamEvent` objects.
+
+def convert_anthropic_tool_call(event: Any) -> Optional[List[ToolCallDelta]]:
+    """
+    Convert an Anthropic stream event into ToolCallDelta objects.
+    
+    Handles:
+    - ContentBlockStartEvent (type='content_block_start') -> Tool Name + ID
+    - ContentBlockDeltaEvent (type='content_block_delta') -> Arguments Delta
+    """
+    
+    # Handle ContentBlockStartEvent (Start of tool use)
+    # event.content_block.type == 'tool_use'
+    if event.type == "content_block_start":
+        if event.content_block.type == "tool_use":
+            return [ToolCallDelta(
+                index=event.index, # Anthropic provides index
+                call_id=event.content_block.id,
+                name=event.content_block.name,
+                arguments_delta=None # No args yet
+            )]
+            
+    # Handle ContentBlockDeltaEvent (JSON args update)
+    elif event.type == "content_block_delta":
+        if event.delta.type == "input_json_delta":
+            return [ToolCallDelta(
+                index=event.index, # Anthropic provides index matching the start event
+                call_id=None,
+                name=None,
+                arguments_delta=event.delta.partial_json
+            )]
+
+    return None

autobyteus/llm/converters/gemini_tool_call_converter.py

@@ -0,0 +1,57 @@
+from typing import List, Optional, Any
+from autobyteus.llm.utils.tool_call_delta import ToolCallDelta
+from google.genai import types as genai_types
+import json
+
+def convert_gemini_tool_calls(part: Any) -> Optional[List[ToolCallDelta]]:
+    """
+    Convert a Gemini stream part (which might be a functionCall) into ToolCallDelta objects.
+    
+    Gemini 'functionCall' usually comes as a complete object in the parts list, 
+    but we treat it as a delta to fit the unified streaming model.
+    """
+    # Check if the part has a function_call
+    # The SDK structure for part usually has a 'function_call' attribute
+    
+    # Note: genai_types.Part might have function_call (FunctionCall)
+    
+    # Safety check for attribute existence
+    if not hasattr(part, "function_call") or not part.function_call:
+        return None
+        
+    fc = part.function_call
+    
+    # Gemini provides the full arguments as a dict in the 'args' attribute
+    # We serialize this back to JSON string for the delta, because our unified
+    # ToolCallDelta expects string deltas for arguments. 
+    # Since we get the full args at once, this is a "complete" delta.
+    
+    # Note: 'args' is a dictionary.
+    arguments_json = json.dumps(fc.args) if fc.args else "{}"
+    
+    # We use a fixed index 0 because Gemini typically emits one function call per candidate part?
+    # Actually Gemini can return multiple function calls in parallel tool use? 
+    # But usually they are in separate parts or candidates.
+    # For now, we assume implicit index 0 for the part itself. 
+    # If parallel tool calling is supported via multiple parts, the caller loop needs to handle indexing 
+    # if it's not provided by the SDK.
+    # However, ToolCallDelta requires an index. We'll default to 0. 
+    # If the caller iterates parts, it might need to assign indices.
+    # But wait, ToolCallDelta logic merges by index. If we always return index 0 for distinct parts, they merge.
+    # We might need to generate a unique index or rely on the caller to offset it.
+    # BUT, Gemini streaming chunks usually contain ONE part at a time.
+    # If we treat each function_call part as a NEW tool call, we need a way to distinguish parallel ones?
+    # Or maybe we just generate a unique call_id if one isn't present.
+    # Gemini function calls don't always have IDs in the stream?
+    # Let's check `fc.name` and `fc.args`.
+    
+    # For simplicity, if we get a full function call, we return one delta that has EVERYTHING.
+    # The unified handler will assume it's a new call if index is new... or if it's the first time 
+    # seeing this index.
+    
+    return [ToolCallDelta(
+        index=0, # Assuming single tool call per part context for now
+        call_id=None, # Gemini doesn't always provide call ID in stream? We'll rely on handler to generate one.
+        name=fc.name,
+        arguments_delta=arguments_json
+    )]

autobyteus/llm/converters/mistral_tool_call_converter.py

@@ -0,0 +1,37 @@
+from typing import List, Optional, Any
+from autobyteus.llm.utils.tool_call_delta import ToolCallDelta
+import logging
+
+logger = logging.getLogger(__name__)
+
+def convert_mistral_tool_calls(tool_calls: List[Any]) -> Optional[List[ToolCallDelta]]:
+    """
+    Convert Mistral streaming tool calls (similar to OpenAI) into ToolCallDelta objects.
+    """
+    if not tool_calls:
+        return None
+
+    deltas = []
+    for tool_call in tool_calls:
+        # Mistral SDK 'ToolCall' or dict in delta
+        # Usually has: index, id, function (name, arguments)
+        
+        index = tool_call.get('index', 0) if isinstance(tool_call, dict) else getattr(tool_call, 'index', 0)
+        call_id = tool_call.get('id', None) if isinstance(tool_call, dict) else getattr(tool_call, 'id', None)
+        
+        function = tool_call.get('function', None) if isinstance(tool_call, dict) else getattr(tool_call, 'function', None)
+        name = None
+        arguments = None
+        
+        if function:
+            name = function.get('name', None) if isinstance(function, dict) else getattr(function, 'name', None)
+            arguments = function.get('arguments', None) if isinstance(function, dict) else getattr(function, 'arguments', None)
+
+        deltas.append(ToolCallDelta(
+            index=index,
+            call_id=call_id,
+            name=name,
+            arguments_delta=arguments
+        ))
+
+    return deltas

autobyteus/llm/converters/openai_tool_call_converter.py

@@ -0,0 +1,38 @@
+# file: autobyteus/autobyteus/llm/converters/openai_tool_call_converter.py
+"""
+Converter for OpenAI SDK tool call deltas to normalized ToolCallDelta format.
+"""
+
+from typing import List, Optional
+from autobyteus.llm.utils.tool_call_delta import ToolCallDelta
+
+# Type checking import for OpenAI SDK types
+try:
+    from openai.types.chat.chat_completion_chunk import ChoiceDeltaToolCall
+except ImportError:
+    ChoiceDeltaToolCall = None  # type: ignore
+
+
+def convert_openai_tool_calls(delta_tool_calls: Optional[List]) -> Optional[List[ToolCallDelta]]:
+    """
+    Convert OpenAI SDK tool call deltas to normalized ToolCallDelta objects.
+
+    Args:
+        delta_tool_calls: List of ChoiceDeltaToolCall from OpenAI SDK, or None.
+
+    Returns:
+        List of normalized ToolCallDelta objects, or None if input is None/empty.
+    """
+    if not delta_tool_calls:
+        return None
+
+    result = []
+    for tc in delta_tool_calls:
+        # Extract fields from OpenAI's ChoiceDeltaToolCall
+        result.append(ToolCallDelta(
+            index=tc.index,
+            call_id=tc.id if tc.id else None,
+            name=tc.function.name if tc.function and tc.function.name else None,
+            arguments_delta=tc.function.arguments if tc.function and tc.function.arguments else None,
+        ))
+    return result

autobyteus/llm/extensions/base_extension.py

@@ -1,8 +1,8 @@
 from abc import ABC, abstractmethod
-from typing import List, Optional, TYPE_CHECKING
+from typing import List, Optional, TYPE_CHECKING, Any
+
 from autobyteus.llm.utils.messages import Message
 from autobyteus.llm.utils.response_types import CompleteResponse
-from autobyteus.llm.user_message import LLMUserMessage
 
 if TYPE_CHECKING:
     from autobyteus.llm.base_llm import BaseLLM
@@ -13,32 +13,26 @@ class LLMExtension(ABC):
 
     @abstractmethod
     async def before_invoke(
-        self, user_message: LLMUserMessage, **kwargs
+        self, messages: List[Message], rendered_payload: Optional[Any] = None, **kwargs
     ) -> None:
         """
-        Called before invoking the LLM with a user message.
+        Called before invoking the LLM with explicit messages.
         """
         pass
 
     @abstractmethod
     async def after_invoke(
-        self, user_message: LLMUserMessage, response: CompleteResponse = None, **kwargs
+        self, messages: List[Message], response: CompleteResponse = None, **kwargs
     ) -> None:
         """
         Called after receiving the response from the LLM.
         
         Args:
-            user_message: The original user message object.
+            messages: The explicit prompt messages used for invocation.
             response: Complete response including content and usage information.
             kwargs: Additional arguments.
         """
         pass
 
-    def on_user_message_added(self, message: Message):
-        pass
-
-    def on_assistant_message_added(self, message: Message):
-        pass
-
     async def cleanup(self) -> None:
         pass