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

autobyteus/llm/api/autobyteus_llm.py

@@ -4,8 +4,9 @@ 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.utils.messages import Message
 from autobyteus.clients import AutobyteusClient
+from autobyteus.llm.prompt_renderers.autobyteus_prompt_renderer import AutobyteusPromptRenderer
 import logging
 import uuid
 
@@ -20,26 +21,34 @@ class AutobyteusLLM(BaseLLM):
         
         self.client = AutobyteusClient(server_url=self.model.host_url)
         self.conversation_id = str(uuid.uuid4())
+        self._renderer = AutobyteusPromptRenderer()
         logger.info(f"AutobyteusLLM initialized for model '{self.model.model_identifier}' with conversation ID: {self.conversation_id}")
 
-    async def _send_user_message_to_llm(
+    async def _send_messages_to_llm(
         self,
-        user_message: LLMUserMessage,
+        messages: List[Message],
         **kwargs
     ) -> CompleteResponse:
-        self.add_user_message(user_message)
+        rendered = await self._renderer.render(messages)
+        if not rendered:
+            raise ValueError("AutobyteusLLM requires at least one user message.")
+        payload = rendered[0]
         try:
             response = await self.client.send_message(
                 conversation_id=self.conversation_id,
                 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
+                user_message=payload.get("content", ""),
+                image_urls=payload.get("image_urls", []),
+                audio_urls=payload.get("audio_urls", []),
+                video_urls=payload.get("video_urls", []),
             )
             
-            assistant_message = response['response']
-            self.add_assistant_message(assistant_message)
+            assistant_message = (
+                response.get("response")
+                or response.get("content")
+                or response.get("message")
+                or ""
+            )
             
             token_usage_data = response.get('token_usage') or {}
             token_usage = TokenUsage(
@@ -54,25 +63,27 @@ class AutobyteusLLM(BaseLLM):
             )
         except Exception as e:
             logger.error(f"Error processing message: {str(e)}")
-            await self._handle_error_cleanup()
             raise
 
-    async def _stream_user_message_to_llm(
+    async def _stream_messages_to_llm(
         self,
-        user_message: LLMUserMessage,
+        messages: List[Message],
         **kwargs
     ) -> AsyncGenerator[ChunkResponse, None]:
-        self.add_user_message(user_message)
+        rendered = await self._renderer.render(messages)
+        if not rendered:
+            raise ValueError("AutobyteusLLM requires at least one user message.")
+        payload = rendered[0]
         complete_response = ""
         
         try:
             async for chunk in self.client.stream_message(
                 conversation_id=self.conversation_id,
                 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
+                user_message=payload.get("content", ""),
+                image_urls=payload.get("image_urls", []),
+                audio_urls=payload.get("audio_urls", []),
+                video_urls=payload.get("video_urls", []),
             ):
                 if 'error' in chunk:
                     raise RuntimeError(chunk['error'])
@@ -100,24 +111,17 @@ class AutobyteusLLM(BaseLLM):
                     video_urls=chunk.get('video_urls', []),
                     usage=token_usage
                 )
-            
-            self.add_assistant_message(complete_response)
         except Exception as e:
             logger.error(f"Error streaming message: {str(e)}")
-            await self._handle_error_cleanup()
             raise
 
     async def cleanup(self):
+        """
+        Clean up the remote conversation. The owning agent controls the HTTP
+        client lifecycle.
+        """
         try:
             await self.client.cleanup(self.conversation_id)
             await super().cleanup()
         except Exception as e:
             logger.error(f"Error during cleanup: {str(e)}")
-        finally:
-            await self.client.close()
-
-    async def _handle_error_cleanup(self):
-        try:
-            await self.cleanup()
-        except Exception as cleanup_error:
-            logger.error(f"Error during error cleanup: {str(cleanup_error)}")

autobyteus/llm/api/claude_llm.py

@@ -1,27 +1,84 @@
-from typing import Dict, Optional, List, AsyncGenerator, Tuple
 import anthropic
 import os
 import logging
+from typing import Dict, Optional, List, AsyncGenerator, Tuple
+
 from autobyteus.llm.models import LLMModel
 from autobyteus.llm.base_llm import BaseLLM
 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
+from autobyteus.llm.converters import convert_anthropic_tool_call
+from autobyteus.llm.prompt_renderers.anthropic_prompt_renderer import AnthropicPromptRenderer
 
 logger = logging.getLogger(__name__)
 
+def _build_thinking_param(extra_params: Dict) -> Optional[Dict]:
+    enabled = extra_params.get("thinking_enabled", False)
+    if not isinstance(enabled, bool) or not enabled:
+        return None
+
+    budget = extra_params.get("thinking_budget_tokens", 1024)
+    try:
+        budget_int = int(budget)
+    except (TypeError, ValueError):
+        budget_int = 1024
+
+    return {"type": "enabled", "budget_tokens": budget_int}
+
+
+def _split_claude_content_blocks(blocks: List) -> Tuple[str, str]:
+    """Split Claude content blocks into visible text and thinking summaries."""
+    content_segments: List[str] = []
+    thinking_segments: List[str] = []
+
+    for block in blocks or []:
+        block_type = getattr(block, "type", None)
+        if block_type is None and isinstance(block, dict):
+            block_type = block.get("type")
+
+        if block_type == "text":
+            text = getattr(block, "text", None)
+            if text is None and isinstance(block, dict):
+                text = block.get("text")
+            if text:
+                content_segments.append(text)
+        elif block_type == "thinking":
+            thinking = getattr(block, "thinking", None)
+            if thinking is None and isinstance(block, dict):
+                thinking = block.get("thinking")
+            if thinking:
+                thinking_segments.append(thinking)
+        elif block_type == "redacted_thinking":
+            redacted = getattr(block, "redacted_thinking", None)
+            if redacted is None and isinstance(block, dict):
+                redacted = block.get("redacted_thinking")
+            if redacted:
+                thinking_segments.append(redacted)
+
+    return "".join(content_segments), "".join(thinking_segments)
+
+
+def _split_system_message(messages: List[Message]) -> Tuple[Optional[str], List[Message]]:
+    system_parts = [msg.content for msg in messages if msg.role == MessageRole.SYSTEM and msg.content]
+    system_prompt = "\n".join(system_parts) if system_parts else None
+    remaining = [msg for msg in messages if msg.role != MessageRole.SYSTEM]
+    return system_prompt, remaining
+
+
 class ClaudeLLM(BaseLLM):
     def __init__(self, model: LLMModel = None, llm_config: LLMConfig = None):
         if model is None:
-            model = LLMModel['claude-4-sonnet']
+            model = LLMModel['claude-4.5-sonnet']
         if llm_config is None:
             llm_config = LLMConfig()
             
         super().__init__(model=model, llm_config=llm_config)
         self.client = self.initialize()
-        self.max_tokens = 8000
+        self._renderer = AnthropicPromptRenderer()
+        # Claude Sonnet 4.5 currently allows up to ~8k output tokens; let config override.
+        self.max_tokens = llm_config.max_tokens if llm_config.max_tokens is not None else 8192
     
     @classmethod
     def initialize(cls):
@@ -36,10 +93,6 @@ class ClaudeLLM(BaseLLM):
         except Exception as e:
             raise ValueError(f"Failed to initialize Anthropic client: {str(e)}")
     
-    def _get_non_system_messages(self) -> List[Dict]:
-        # 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:
         return TokenUsage(
             prompt_tokens=input_tokens,
@@ -47,23 +100,37 @@ class ClaudeLLM(BaseLLM):
             total_tokens=input_tokens + output_tokens
         )
     
-    async def _send_user_message_to_llm(self, user_message: LLMUserMessage, **kwargs) -> CompleteResponse:
-        self.add_user_message(user_message)
+    async def _send_messages_to_llm(self, messages: List[Message], **kwargs) -> CompleteResponse:
+        try:
+            system_prompt, non_system = _split_system_message(messages)
+            formatted_messages = await self._renderer.render(non_system)
+            thinking_param = _build_thinking_param(self.config.extra_params)
 
-        # NOTE: This implementation does not yet support multimodal inputs for Claude.
-        # It will only send the text content.
+            request_kwargs = {
+                "model": self.model.value,
+                "max_tokens": self.max_tokens,
+                "messages": formatted_messages,
+            }
+            if system_prompt:
+                request_kwargs["system"] = system_prompt
+            if thinking_param:
+                # Extended thinking is not compatible with temperature modifications
+                request_kwargs["thinking"] = thinking_param
+            else:
+                request_kwargs["temperature"] = 0
 
-        try:
             response = self.client.messages.create(
-                model=self.model.value,
-                max_tokens=self.max_tokens,
-                temperature=0,
-                system=self.system_message,
-                messages=self._get_non_system_messages()
+                **request_kwargs
             )
 
-            assistant_message = response.content[0].text
-            self.add_assistant_message(assistant_message)
+            assistant_message = getattr(response, "text", "") or ""
+            reasoning_summary = None
+            if response.content:
+                parsed_text, parsed_thinking = _split_claude_content_blocks(response.content)
+                if parsed_text:
+                    assistant_message = parsed_text
+                if parsed_thinking:
+                    reasoning_summary = parsed_thinking
 
             token_usage = self._create_token_usage(
                 response.usage.input_tokens,
@@ -74,35 +141,72 @@ class ClaudeLLM(BaseLLM):
             
             return CompleteResponse(
                 content=assistant_message,
+                reasoning=reasoning_summary,
                 usage=token_usage
             )
         except anthropic.APIError as e:
             logger.error(f"Error in Claude API call: {str(e)}")
             raise ValueError(f"Error in Claude API call: {str(e)}")
     
-    async def _stream_user_message_to_llm(
-        self, user_message: LLMUserMessage, **kwargs
+    async def _stream_messages_to_llm(
+        self, messages: List[Message], **kwargs
     ) -> AsyncGenerator[ChunkResponse, None]:
-        self.add_user_message(user_message)
         complete_response = ""
+        complete_reasoning = ""
         final_message = None
 
-        # NOTE: This implementation does not yet support multimodal inputs for Claude.
-        # It will only send the text content.
+        # Extract tools if provided
+        tools = kwargs.get("tools")
 
         try:
-            with self.client.messages.stream(
-                model=self.model.value,
-                max_tokens=self.max_tokens,
-                temperature=0,
-                system=self.system_message,
-                messages=self._get_non_system_messages(),
-            ) as stream:
+            # Prepare arguments for stream
+            system_prompt, non_system = _split_system_message(messages)
+            formatted_messages = await self._renderer.render(non_system)
+            thinking_param = _build_thinking_param(self.config.extra_params)
+            stream_kwargs = {
+                "model": self.model.value,
+                "max_tokens": self.max_tokens,
+                "messages": formatted_messages,
+            }
+            if system_prompt:
+                stream_kwargs["system"] = system_prompt
+            if thinking_param:
+                # Extended thinking is not compatible with temperature modifications
+                stream_kwargs["thinking"] = thinking_param
+            else:
+                stream_kwargs["temperature"] = 0
+            
+            if tools:
+                stream_kwargs["tools"] = tools
+
+            with self.client.messages.stream(**stream_kwargs) as stream:
                 for event in stream:
-                    if event.type == "content_block_delta" and event.delta.type == "text_delta":
-                        complete_response += event.delta.text
+                    
+                    # Handle text content
+                    if event.type == "content_block_delta":
+                        delta_type = getattr(event.delta, "type", None)
+                        if delta_type == "text_delta":
+                            complete_response += event.delta.text
+                            yield ChunkResponse(
+                                content=event.delta.text,
+                                is_complete=False
+                            )
+                        elif delta_type == "thinking_delta":
+                            thinking_delta = getattr(event.delta, "thinking", None)
+                            if thinking_delta:
+                                complete_reasoning += thinking_delta
+                                yield ChunkResponse(
+                                    content="",
+                                    reasoning=thinking_delta,
+                                    is_complete=False
+                                )
+                    
+                    # Handle tool calls using common converter
+                    tool_calls = convert_anthropic_tool_call(event)
+                    if tool_calls:
                         yield ChunkResponse(
-                            content=event.delta.text,
+                            content="",
+                            tool_calls=tool_calls,
                             is_complete=False
                         )
                     
@@ -120,7 +224,9 @@ class ClaudeLLM(BaseLLM):
                         usage=token_usage
                     )
 
-            self.add_assistant_message(complete_response)
+            # Only add assistant message if there's actual content.
+            # Tool-call-only responses should not add empty messages, as Claude API
+            # rejects subsequent requests with "all messages must have non-empty content".
         except anthropic.APIError as e:
             logger.error(f"Error in Claude API streaming: {str(e)}")
             raise ValueError(f"Error in Claude API streaming: {str(e)}")

autobyteus/llm/api/gemini_llm.py

@@ -1,29 +1,41 @@
 import logging
-from typing import Dict, Optional, List, AsyncGenerator, Any
-from google import genai
+import asyncio
+from typing import Dict, List, AsyncGenerator, Any, Optional
 from google.genai import types as genai_types
-import os
 from autobyteus.llm.models import LLMModel
 from autobyteus.llm.base_llm import BaseLLM
 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
+from autobyteus.utils.gemini_helper import initialize_gemini_client_with_runtime
+from autobyteus.utils.gemini_model_mapping import resolve_model_for_runtime
+from autobyteus.llm.converters import convert_gemini_tool_calls
+from autobyteus.llm.prompt_renderers.gemini_prompt_renderer import GeminiPromptRenderer
 
 logger = logging.getLogger(__name__)
 
-def _format_gemini_history(messages: List[Message]) -> List[Dict[str, Any]]:
-    """Formats internal message history for the Gemini API."""
-    history = []
-    # System message is handled separately in the new API
-    for msg in messages:
-        if msg.role in [MessageRole.USER, MessageRole.ASSISTANT]:
-            # NOTE: This history conversion will need to be updated for multimodal messages
-            role = 'model' if msg.role == MessageRole.ASSISTANT else 'user'
-            # The `parts` must be a list of dictionaries (Part objects), not a list of strings.
-            history.append({"role": role, "parts": [{"text": msg.content}]})
-    return history
+def _split_gemini_parts(parts: List[Any]) -> tuple[str, str]:
+    """Split Gemini content parts into visible text and thought summaries."""
+    content_segments: List[str] = []
+    thought_segments: List[str] = []
+    for part in parts or []:
+        text = getattr(part, "text", None)
+        if not text:
+            continue
+        if getattr(part, "thought", False):
+            thought_segments.append(text)
+        else:
+            content_segments.append(text)
+    return "".join(content_segments), "".join(thought_segments)
+
+
+def _split_system_message(messages: List[Message]) -> tuple[Optional[str], List[Message]]:
+    system_parts = [msg.content for msg in messages if msg.role == MessageRole.SYSTEM and msg.content]
+    system_prompt = "\n".join(system_parts) if system_parts else None
+    remaining = [msg for msg in messages if msg.role != MessageRole.SYSTEM]
+    return system_prompt, remaining
+
 
 class GeminiLLM(BaseLLM):
     def __init__(self, model: LLMModel = None, llm_config: LLMConfig = None):
@@ -32,59 +44,84 @@ class GeminiLLM(BaseLLM):
         }
         
         if model is None:
-            model = LLMModel['gemini-2.5-flash']
+            # Default to the latest low-latency Gemini LLM.
+            model = LLMModel['gemini-3-flash-preview']
         if llm_config is None:
             llm_config = LLMConfig()
             
         super().__init__(model=model, llm_config=llm_config)
-        self.client = self.initialize()
-        self.async_client = self.client.aio
-
-    @classmethod
-    def initialize(cls) -> genai.client.Client:
-        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.")
+        
         try:
-            return genai.Client()
+            self.client, self.runtime_info = initialize_gemini_client_with_runtime()
+            self.async_client = self.client.aio
+            self._renderer = GeminiPromptRenderer()
         except Exception as e:
-            logger.error(f"Failed to initialize Gemini client: {str(e)}")
-            raise ValueError(f"Failed to initialize Gemini client: {str(e)}")
+            # Re-raise or handle initialization errors specifically for the LLM context if needed
+            logger.error(f"Failed to initialize Gemini LLM: {str(e)}")
+            raise
 
-    def _get_generation_config(self) -> genai_types.GenerateContentConfig:
+    def _get_generation_config(self, system_prompt: Optional[str] = None) -> genai_types.GenerateContentConfig:
         """Builds the generation config, handling special cases like 'thinking'."""
         config = self.generation_config_dict.copy()
 
-        thinking_config = None
-        if "flash" in self.model.value:
-            thinking_config = genai_types.ThinkingConfig(thinking_budget=0)
+        # Map thinking_level to token budget
+        # Values based on Gemini 3 API recommendations
+        THINKING_LEVEL_BUDGETS = {
+            "minimal": 0,
+            "low": 1024,
+            "medium": 4096,
+            "high": 16384,
+        }
         
-        # System instruction is now part of the config
-        system_instruction = self.system_message if self.system_message else None
+        # Read thinking_level from extra_params (set by user config)
+        # Default to "minimal" (0 tokens) for backward compatibility
+        thinking_level = self.config.extra_params.get("thinking_level", "minimal")
+        thinking_budget = THINKING_LEVEL_BUDGETS.get(thinking_level, 0)
         
+        include_thoughts = self.config.extra_params.get("include_thoughts", False)
+        if not isinstance(include_thoughts, bool):
+            include_thoughts = False
+        thinking_config = genai_types.ThinkingConfig(
+            thinking_budget=thinking_budget,
+            include_thoughts=include_thoughts
+        )
+        
+        # System instruction is now part of the config
+        system_instruction = system_prompt if system_prompt is not None else self.system_message
+
         return genai_types.GenerateContentConfig(
             **config,
             thinking_config=thinking_config,
             system_instruction=system_instruction
         )
 
-    async def _send_user_message_to_llm(self, user_message: LLMUserMessage, **kwargs) -> CompleteResponse:
-        self.add_user_message(user_message)
-        
+    async def _send_messages_to_llm(self, messages: List[Message], **kwargs) -> CompleteResponse:
         try:
-            history = _format_gemini_history(self.messages)
-            generation_config = self._get_generation_config()
+            system_prompt, non_system = _split_system_message(messages)
+            history = await self._renderer.render(non_system)
+            generation_config = self._get_generation_config(system_prompt=system_prompt)
 
+            # FIX: Removed 'models/' prefix to support Vertex AI
+            runtime_adjusted_model = resolve_model_for_runtime(
+                self.model.value,
+                modality="llm",
+                runtime=getattr(self, "runtime_info", None) and self.runtime_info.runtime,
+            )
             response = await self.async_client.models.generate_content(
-                model=f"models/{self.model.value}",
+                model=runtime_adjusted_model,
                 contents=history,
                 config=generation_config,
             )
             
-            assistant_message = response.text
-            self.add_assistant_message(assistant_message)
-            
+            assistant_message = response.text or ""
+            reasoning_summary = None
+            if response.candidates and response.candidates[0].content and response.candidates[0].content.parts:
+                parsed_text, parsed_thoughts = _split_gemini_parts(response.candidates[0].content.parts)
+                if parsed_text:
+                    assistant_message = parsed_text
+                if parsed_thoughts:
+                    reasoning_summary = parsed_thoughts
+
             token_usage = TokenUsage(
                 prompt_tokens=0,
                 completion_tokens=0,
@@ -93,35 +130,102 @@ class GeminiLLM(BaseLLM):
             
             return CompleteResponse(
                 content=assistant_message,
+                reasoning=reasoning_summary,
                 usage=token_usage
             )
         except Exception as e:
             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)
+    async def _stream_messages_to_llm(self, messages: List[Message], **kwargs) -> AsyncGenerator[ChunkResponse, None]:
         complete_response = ""
+        complete_reasoning = ""
+        
+        # Extract tools if provided
+        tools = kwargs.get("tools")
         
         try:
-            history = _format_gemini_history(self.messages)
-            generation_config = self._get_generation_config()
+            system_prompt, non_system = _split_system_message(messages)
+            history = await self._renderer.render(non_system)
+            generation_config = self._get_generation_config(system_prompt=system_prompt)
+            
+            # Add tools to config if present
+            # Note: In google.genai, tools can be passed in config
+            if tools:
+                # Auto-wrap tools if they appear to be raw function declarations
+                if isinstance(tools, list) and len(tools) > 0:
+                    first_tool = tools[0]
+                    # Check if it's a raw declaration (dict with name/description) but NOT a wrapper (dict with function_declarations)
+                    if isinstance(first_tool, dict):
+                        is_declaration = "name" in first_tool and "description" in first_tool
+                        is_wrapper = "function_declarations" in first_tool
+                        
+                        if is_declaration and not is_wrapper:
+                             # Wrap the list of declarations into a single Tool structure
+                             tools = [{"function_declarations": tools}]
 
-            response_stream = await self.async_client.models.generate_content_stream(
-                model=f"models/{self.model.value}",
-                contents=history,
-                config=generation_config,
+                try:
+                    generation_config.tools = tools
+                except Exception:
+                    # Fallback or strict strict typing issues
+                    pass
+
+            # FIX: Removed 'models/' prefix to support Vertex AI
+            runtime_adjusted_model = resolve_model_for_runtime(
+                self.model.value,
+                modality="llm",
+                runtime=getattr(self, "runtime_info", None) and self.runtime_info.runtime,
             )
+            
+            # Prepare call args
+            call_kwargs = {
+                "model": runtime_adjusted_model,
+                "contents": history,
+                "config": generation_config,
+            }
+            # If explicit tools argument is needed and not supported in config for this SDK version:
+            # call_kwargs['tools'] = tools 
+            # But usually config holds it.
+
+            response_stream = await self.async_client.models.generate_content_stream(**call_kwargs)
 
             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)
+                handled_parts = False
+                if chunk.candidates and chunk.candidates[0].content and chunk.candidates[0].content.parts:
+                    handled_parts = True
+                    for part in chunk.candidates[0].content.parts:
+                        part_text = getattr(part, "text", None)
+                        if part_text:
+                            if getattr(part, "thought", False):
+                                complete_reasoning += part_text
+                                yield ChunkResponse(
+                                    content="",
+                                    reasoning=part_text,
+                                    is_complete=False
+                                )
+                            else:
+                                complete_response += part_text
+                                yield ChunkResponse(
+                                    content=part_text,
+                                    is_complete=False
+                                )
+
+                        tool_calls = convert_gemini_tool_calls(part)
+                        if tool_calls:
+                            yield ChunkResponse(
+                                content="",
+                                tool_calls=tool_calls,
+                                is_complete=False
+                            )
+
+                if not handled_parts:
+                    chunk_text = chunk.text
+                    if chunk_text:
+                        complete_response += chunk_text
+                        yield ChunkResponse(
+                            content=chunk_text,
+                            is_complete=False
+                        )
 
             token_usage = TokenUsage(
                 prompt_tokens=0,

autobyteus/llm/api/grok_llm.py

@@ -10,7 +10,7 @@ class GrokLLM(OpenAICompatibleLLM):
     def __init__(self, model: LLMModel = None, llm_config: LLMConfig = None):
         # Provide defaults if not specified
         if model is None:
-            model = LLMModel['grok-2-1212']
+            model = LLMModel['grok-4-1-fast-reasoning']
         if llm_config is None:
             llm_config = LLMConfig()