alita-sdk 0.3.528__py3-none-any.whl → 0.3.554__py3-none-any.whl

alita_sdk/runtime/tools/image_generation.py

@@ -1,7 +1,9 @@
 """
 Image generation tool for Alita SDK.
 """
+import json
 import logging
+import uuid
 from typing import Optional, Type, Any, List, Literal
 from langchain_core.tools import BaseTool, BaseToolkit
 from pydantic import BaseModel, Field, create_model, ConfigDict
@@ -76,7 +78,12 @@ class ImageGenerationTool(BaseTool):
     """Tool for generating images using the Alita client."""
     
     name: str = "generate_image"
-    description: str = "Generate images from text prompts using AI models"
+    description: str = (
+        "Generate images from text prompts using AI models. "
+        "Returns a JSON object with 'cached_image_id' field containing a reference to the generated image data. "
+        "The cached_image_id can be used to save or process the image. "
+        "The actual image data is stored temporarily and can be retrieved using the cached_image_id reference."
+    )
     args_schema: Type[BaseModel] = ImageGenerationInput
     alita_client: Any = None
     
@@ -85,10 +92,10 @@ class ImageGenerationTool(BaseTool):
         self.alita_client = client
     
     def _run(self, prompt: str, n: int = 1, size: str = "auto",
-             quality: str = "auto", style: Optional[str] = None) -> list:
+             quality: str = "auto", style: Optional[str] = None) -> str:
         """Generate an image based on the provided parameters."""
         try:
-            logger.info(f"Generating image with prompt: {prompt[:50]}...")
+            logger.debug(f"Generating image with prompt: {prompt[:50]}...")
             
             result = self.alita_client.generate_image(
                 prompt=prompt,
@@ -98,57 +105,56 @@ class ImageGenerationTool(BaseTool):
                 style=style
             )
             
-            # Return multimodal content format for LLM consumption
+            # Return simple JSON structure with reference ID instead of full base64
             if 'data' in result:
                 images = result['data']
-                content_chunks = []
                 
-                # Add a text description of what was generated
-                if len(images) == 1:
-                    content_chunks.append({
-                        "type": "text",
-                        "text": f"Generated image for prompt: '{prompt}'"
-                    })
-                else:
-                    content_chunks.append({
-                        "type": "text",
-                        "text": f"Generated {len(images)} images for "
-                                f"prompt: '{prompt}'"
+                # Process all images with unified structure
+                images_list = []
+                for idx, image_data in enumerate(images, 1):
+                    if not image_data.get('b64_json'):
+                        continue
+                    
+                    cached_image_id = f"img_{uuid.uuid4().hex[:12]}"
+                    
+                    # Store in cache
+                    if hasattr(self.alita_client, '_generated_images_cache'):
+                        self.alita_client._generated_images_cache[cached_image_id] = {
+                            'base64_data': image_data['b64_json']
+                        }
+                        logger.debug(f"Stored generated image in cache with ID: {cached_image_id}")
+                    
+                    images_list.append({
+                        "image_number": idx,
+                        "image_type": "png",
+                        "cached_image_id": cached_image_id
                     })
                 
-                # Add image content for each generated image
-                for image_data in images:
-                    if image_data.get('url'):
-                        content_chunks.append({
-                            "type": "image_url",
-                            "image_url": {
-                                "url": image_data['url']
-                            }
-                        })
-                    elif image_data.get('b64_json'):
-                        content_chunks.append({
-                            "type": "image_url",
-                            "image_url": {
-                                "url": f"data:image/png;base64,"
-                                       f"{image_data['b64_json']}"
-                            }
-                        })
+                if not images_list:
+                    return json.dumps({
+                        "status": "error",
+                        "message": "No base64 image data found"
+                    })
                 
-                return content_chunks
+                return json.dumps({
+                    "status": "success",
+                    "prompt": prompt,
+                    "total_images": len(images_list),
+                    "images": images_list
+                })
             
-            # Fallback to text response if no images in result
-            return [{
-                "type": "text",
-                "text": f"Image generation completed but no images "
-                        f"returned: {result}"
-            }]
+            # Fallback to error response if no images in result
+            return json.dumps({
+                "status": "error",
+                "message": f"Image generation completed but no images returned: {result}"
+            })
             
         except Exception as e:
             logger.error(f"Error generating image: {e}")
-            return [{
-                "type": "text",
-                "text": f"Error generating image: {str(e)}"
-            }]
+            return json.dumps({
+                "status": "error",
+                "message": f"Error generating image: {str(e)}"
+            })
     
     async def _arun(self, prompt: str, n: int = 1, size: str = "256x256",
                     quality: str = "auto",

alita_sdk/runtime/tools/llm.py

@@ -6,6 +6,7 @@ from typing import Any, Optional, List, Union, Literal
 from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
 from langchain_core.runnables import RunnableConfig
 from langchain_core.tools import BaseTool, ToolException
+from langchain_core.callbacks import dispatch_custom_event
 from pydantic import Field
 
 from ..langchain.constants import ELITEA_RS
@@ -42,6 +43,17 @@ logger = logging.getLogger(__name__)
     
 #     return supports_reasoning
 
+JSON_INSTRUCTION_TEMPLATE = (
+        "\n\n**IMPORTANT: You MUST respond with ONLY a valid JSON object.**\n\n"
+        "Required JSON fields:\n{field_descriptions}\n\n"
+        "Example format:\n"
+        "{{\n{example_fields}\n}}\n\n"
+        "Rules:\n"
+        "1. Output ONLY the JSON object - no markdown, no explanations, no extra text\n"
+        "2. Ensure all required fields are present\n"
+        "3. Use proper JSON syntax with double quotes for strings\n"
+        "4. Do not wrap the JSON in code blocks or backticks"
+    )
 
 class LLMNode(BaseTool):
     """Enhanced LLM node with chat history and tool binding support"""
@@ -65,6 +77,221 @@ class LLMNode(BaseTool):
     steps_limit: Optional[int] = Field(default=25, description='Maximum steps for tool execution')
     tool_execution_timeout: Optional[int] = Field(default=900, description='Timeout (seconds) for tool execution. Default is 15 minutes.')
 
+    def _prepare_structured_output_params(self) -> dict:
+        """
+        Prepare structured output parameters from structured_output_dict.
+
+        Returns:
+            Dictionary with parameter definitions for creating Pydantic model
+        """
+        struct_params = {
+            key: {
+                "type": 'list[str]' if 'list' in value else value,
+                "description": ""
+            }
+            for key, value in (self.structured_output_dict or {}).items()
+        }
+        # Add default output field for proper response to user
+        struct_params[ELITEA_RS] = {
+            'description': 'final output to user (summarized output from LLM)',
+            'type': 'str',
+            "default": None
+        }
+        return struct_params
+
+    def _invoke_with_structured_output(self, llm_client: Any, messages: List, struct_model: Any, config: RunnableConfig):
+        """
+        Invoke LLM with structured output, handling tool calls if present.
+
+        Args:
+            llm_client: LLM client instance
+            messages: List of conversation messages
+            struct_model: Pydantic model for structured output
+            config: Runnable configuration
+
+        Returns:
+            Tuple of (completion, initial_completion, final_messages)
+        """
+        initial_completion = llm_client.invoke(messages, config=config)
+
+        if hasattr(initial_completion, 'tool_calls') and initial_completion.tool_calls:
+            # Handle tool calls first, then apply structured output
+            new_messages, _ = self._run_async_in_sync_context(
+                self.__perform_tool_calling(initial_completion, messages, llm_client, config)
+            )
+            llm = self.__get_struct_output_model(llm_client, struct_model)
+            completion = llm.invoke(new_messages, config=config)
+            return completion, initial_completion, new_messages
+        else:
+            # Direct structured output without tool calls
+            llm = self.__get_struct_output_model(llm_client, struct_model)
+            completion = llm.invoke(messages, config=config)
+            return completion, initial_completion, messages
+
+    def _build_json_instruction(self, struct_model: Any) -> str:
+        """
+        Build JSON instruction message for fallback handling.
+
+        Args:
+            struct_model: Pydantic model with field definitions
+
+        Returns:
+            Formatted JSON instruction string
+        """
+        field_descriptions = []
+        for name, field in struct_model.model_fields.items():
+            field_type = field.annotation.__name__ if hasattr(field.annotation, '__name__') else str(field.annotation)
+            field_desc = field.description or field_type
+            field_descriptions.append(f"  - {name} ({field_type}): {field_desc}")
+
+        example_fields = ",\n".join([
+            f'  "{k}": <{field.annotation.__name__ if hasattr(field.annotation, "__name__") else "value"}>'
+            for k, field in struct_model.model_fields.items()
+        ])
+
+        return JSON_INSTRUCTION_TEMPLATE.format(
+            field_descriptions="\n".join(field_descriptions),
+            example_fields=example_fields
+        )
+
+    def _create_fallback_completion(self, content: str, struct_model: Any) -> Any:
+        """
+        Create a fallback completion object when JSON parsing fails.
+
+        Args:
+            content: Plain text content from LLM
+            struct_model: Pydantic model to construct
+
+        Returns:
+            Pydantic model instance with fallback values
+        """
+        result_dict = {}
+        for k, field in struct_model.model_fields.items():
+            if k == ELITEA_RS:
+                result_dict[k] = content
+            elif field.is_required():
+                # Set default values for required fields based on type
+                result_dict[k] = field.default if field.default is not None else None
+            else:
+                result_dict[k] = field.default
+        return struct_model.model_construct(**result_dict)
+
+    def _handle_structured_output_fallback(self, llm_client: Any, messages: List, struct_model: Any,
+                                          config: RunnableConfig, original_error: Exception) -> Any:
+        """
+        Handle structured output fallback through multiple strategies.
+
+        Tries fallback methods in order:
+        1. json_mode with explicit instructions
+        2. function_calling method
+        3. Plain text with JSON extraction
+
+        Args:
+            llm_client: LLM client instance
+            messages: Original conversation messages
+            struct_model: Pydantic model for structured output
+            config: Runnable configuration
+            original_error: The original ValueError that triggered fallback
+
+        Returns:
+            Completion with structured output (best effort)
+
+        Raises:
+            Propagates exceptions from LLM invocation
+        """
+        logger.error(f"Error invoking structured output model: {format_exc()}")
+        logger.info("Attempting to fall back to json mode")
+
+        # Build JSON instruction once
+        json_instruction = self._build_json_instruction(struct_model)
+
+        # Add instruction to messages
+        modified_messages = messages.copy()
+        if modified_messages and isinstance(modified_messages[-1], HumanMessage):
+            modified_messages[-1] = HumanMessage(
+                content=modified_messages[-1].content + json_instruction
+            )
+        else:
+            modified_messages.append(HumanMessage(content=json_instruction))
+
+        # Try json_mode with explicit instructions
+        try:
+            completion = self.__get_struct_output_model(
+                llm_client, struct_model, method="json_mode"
+            ).invoke(modified_messages, config=config)
+            return completion
+        except Exception as json_mode_error:
+            logger.warning(f"json_mode also failed: {json_mode_error}")
+            logger.info("Falling back to function_calling method")
+
+            # Try function_calling as a third fallback
+            try:
+                completion = self.__get_struct_output_model(
+                    llm_client, struct_model, method="function_calling"
+                ).invoke(modified_messages, config=config)
+                return completion
+            except Exception as function_calling_error:
+                logger.error(f"function_calling also failed: {function_calling_error}")
+                logger.info("Final fallback: using plain LLM response")
+
+                # Last resort: get plain text response and wrap in structure
+                plain_completion = llm_client.invoke(modified_messages, config=config)
+                content = plain_completion.content.strip() if hasattr(plain_completion, 'content') else str(plain_completion)
+
+                # Try to extract JSON from the response
+                import json
+                import re
+
+                json_match = re.search(r'\{.*\}', content, re.DOTALL)
+                if json_match:
+                    try:
+                        parsed_json = json.loads(json_match.group(0))
+                        # Validate it has expected fields and wrap in pydantic model
+                        completion = struct_model(**parsed_json)
+                        return completion
+                    except (json.JSONDecodeError, Exception) as parse_error:
+                        logger.warning(f"Could not parse extracted JSON: {parse_error}")
+                        return self._create_fallback_completion(content, struct_model)
+                else:
+                    # No JSON found, create response with content in elitea_response
+                    return self._create_fallback_completion(content, struct_model)
+
+    def _format_structured_output_result(self, result: dict, messages: List, initial_completion: Any) -> dict:
+        """
+        Format structured output result with properly formatted messages.
+
+        Args:
+            result: Result dictionary from model_dump()
+            messages: Original conversation messages
+            initial_completion: Initial completion before tool calls
+
+        Returns:
+            Formatted result dictionary with messages
+        """
+        # Ensure messages are properly formatted
+        if result.get('messages') and isinstance(result['messages'], list):
+            result['messages'] = [{'role': 'assistant', 'content': '\n'.join(result['messages'])}]
+        else:
+            # Extract content from initial_completion, handling thinking blocks
+            fallback_content = result.get(ELITEA_RS, '')
+            if not fallback_content and initial_completion:
+                content_parts = self._extract_content_from_completion(initial_completion)
+                fallback_content = content_parts.get('text') or ''
+                thinking = content_parts.get('thinking')
+
+                # Log thinking if present
+                if thinking:
+                    logger.debug(f"Thinking content present in structured output: {thinking[:100]}...")
+
+                if not fallback_content:
+                    # Final fallback to raw content
+                    content = initial_completion.content
+                    fallback_content = content if isinstance(content, str) else str(content)
+
+            result['messages'] = messages + [AIMessage(content=fallback_content)]
+
+        return result
+
     def get_filtered_tools(self) -> List[BaseTool]:
         """
         Filter available tools based on tool_names list.
@@ -162,8 +389,6 @@ class LLMNode(BaseTool):
             if func_args.get('system') is None or func_args.get('task') is None:
                 raise ToolException(f"LLMNode requires 'system' and 'task' parameters in input mapping. "
                                     f"Actual params: {func_args}")
-                raise ToolException(f"LLMNode requires 'system' and 'task' parameters in input mapping. "
-                                    f"Actual params: {func_args}")
             # cast to str in case user passes variable different from str
             messages = [SystemMessage(content=str(func_args.get('system'))), *func_args.get('chat_history', []), HumanMessage(content=str(func_args.get('task')))]
             # Remove pre-last item if last two messages are same type and content
@@ -195,44 +420,23 @@ class LLMNode(BaseTool):
         try:
             if self.structured_output and self.output_variables:
                 # Handle structured output
-                struct_params = {
-                    key: {
-                        "type": 'list[str]' if 'list' in value else value,
-                        "description": ""
-                    }
-                    for key, value in (self.structured_output_dict or {}).items()
-                }
-                # Add default output field for proper response to user
-                struct_params['elitea_response'] = {
-                    'description': 'final output to user (summarized output from LLM)', 'type': 'str',
-                    "default": None}
+                struct_params = self._prepare_structured_output_params()
                 struct_model = create_pydantic_model(f"LLMOutput", struct_params)
-                initial_completion = llm_client.invoke(messages, config=config)
-                if hasattr(initial_completion, 'tool_calls') and initial_completion.tool_calls:
-                    new_messages, _ = self._run_async_in_sync_context(
-                        self.__perform_tool_calling(initial_completion, messages, llm_client, config)
+
+                try:
+                    completion, initial_completion, final_messages = self._invoke_with_structured_output(
+                        llm_client, messages, struct_model, config
                     )
-                    llm = self.__get_struct_output_model(llm_client, struct_model)
-                    completion = llm.invoke(new_messages, config=config)
-                    result = completion.model_dump()
-                else:
-                    try:
-                        llm = self.__get_struct_output_model(llm_client, struct_model)
-                        completion = llm.invoke(messages, config=config)
-                    except ValueError as e:
-                        logger.error(f"Error invoking structured output model: {format_exc()}")
-                        logger.info("Attemping to fall back to json mode")
-                        # Fallback to regular LLM with JSON extraction
-                        completion = self.__get_struct_output_model(llm_client, struct_model,
-                                                                    method="json_mode").invoke(messages, config=config)
-                    result = completion.model_dump()
-
-                # Ensure messages are properly formatted
-                if result.get('messages') and isinstance(result['messages'], list):
-                    result['messages'] = [{'role': 'assistant', 'content': '\n'.join(result['messages'])}]
-                else:
-                    result['messages'] = messages + [
-                        AIMessage(content=result.get(ELITEA_RS, '') or initial_completion.content)]
+                except ValueError as e:
+                    # Handle fallback for structured output failures
+                    completion = self._handle_structured_output_fallback(
+                        llm_client, messages, struct_model, config, e
+                    )
+                    initial_completion = None
+                    final_messages = messages
+
+                result = completion.model_dump()
+                result = self._format_structured_output_result(result, final_messages, initial_completion or completion)
 
                 return result
             else:
@@ -250,24 +454,89 @@ class LLMNode(BaseTool):
                     if self.output_variables:
                         if self.output_variables[0] == 'messages':
                             return output_msgs
-                        output_msgs[self.output_variables[0]] = current_completion.content if current_completion else None
+                        # Extract content properly from thinking-enabled responses
+                        if current_completion:
+                            content_parts = self._extract_content_from_completion(current_completion)
+                            text_content = content_parts.get('text')
+                            thinking = content_parts.get('thinking')
+                            
+                            # Dispatch thinking event if present
+                            if thinking:
+                                try:
+                                    model_name = getattr(llm_client, 'model_name', None) or getattr(llm_client, 'model', 'LLM')
+                                    dispatch_custom_event(
+                                        name="thinking_step",
+                                        data={
+                                            "message": thinking,
+                                            "tool_name": f"LLM ({model_name})",
+                                            "toolkit": "reasoning",
+                                        },
+                                        config=config,
+                                    )
+                                except Exception as e:
+                                    logger.warning(f"Failed to dispatch thinking event: {e}")
+                            
+                            if text_content:
+                                output_msgs[self.output_variables[0]] = text_content
+                            else:
+                                # Fallback to raw content
+                                content = current_completion.content
+                                output_msgs[self.output_variables[0]] = content if isinstance(content, str) else str(content)
+                        else:
+                            output_msgs[self.output_variables[0]] = None
 
                     return output_msgs
                 else:
-                    # Regular text response
-                    content = completion.content.strip() if hasattr(completion, 'content') else str(completion)
+                    # Regular text response - handle both simple strings and thinking-enabled responses
+                    content_parts = self._extract_content_from_completion(completion)
+                    thinking = content_parts.get('thinking')
+                    text_content = content_parts.get('text') or ''
+                    
+                    # Fallback to string representation if no content extracted
+                    if not text_content:
+                        if hasattr(completion, 'content'):
+                            content = completion.content
+                            text_content = content.strip() if isinstance(content, str) else str(content)
+                        else:
+                            text_content = str(completion)
+                    
+                    # Dispatch thinking step event to chat if present
+                    if thinking:
+                        logger.info(f"Model thinking: {thinking[:200]}..." if len(thinking) > 200 else f"Model thinking: {thinking}")
+                        
+                        # Dispatch custom event for thinking step to be displayed in chat
+                        try:
+                            model_name = getattr(llm_client, 'model_name', None) or getattr(llm_client, 'model', 'LLM')
+                            dispatch_custom_event(
+                                name="thinking_step",
+                                data={
+                                    "message": thinking,
+                                    "tool_name": f"LLM ({model_name})",
+                                    "toolkit": "reasoning",
+                                },
+                                config=config,
+                            )
+                        except Exception as e:
+                            logger.warning(f"Failed to dispatch thinking event: {e}")
+                    
+                    # Build the AI message with both thinking and text
+                    # Store thinking in additional_kwargs for potential future use
+                    ai_message_kwargs = {'content': text_content}
+                    if thinking:
+                        ai_message_kwargs['additional_kwargs'] = {'thinking': thinking}
+                    ai_message = AIMessage(**ai_message_kwargs)
 
                     # Try to extract JSON if output variables are specified (but exclude 'messages' which is handled separately)
                     json_output_vars = [var for var in (self.output_variables or []) if var != 'messages']
                     if json_output_vars:
                         # set response to be the first output variable for non-structured output
-                        response_data = {json_output_vars[0]: content}
-                        new_messages = messages + [AIMessage(content=content)]
+                        response_data = {json_output_vars[0]: text_content}
+                        new_messages = messages + [ai_message]
                         response_data['messages'] = new_messages
                         return response_data
 
                     # Simple text response (either no output variables or JSON parsing failed)
-                    new_messages = messages + [AIMessage(content=content)]
+                    new_messages = messages + [ai_message]
                     return {"messages": new_messages}
 
         except Exception as e:
@@ -285,6 +554,56 @@ class LLMNode(BaseTool):
         # Legacy support for old interface
         return self.invoke(kwargs, **kwargs)
     
+    @staticmethod
+    def _extract_content_from_completion(completion) -> dict:
+        """Extract thinking and text content from LLM completion.
+        
+        Handles Anthropic's extended thinking format where content is a list
+        of blocks with types: 'thinking' and 'text'.
+        
+        Args:
+            completion: LLM completion object with content attribute
+            
+        Returns:
+            dict with 'thinking' and 'text' keys
+        """
+        result = {'thinking': None, 'text': None}
+        
+        if not hasattr(completion, 'content'):
+            return result
+            
+        content = completion.content
+        
+        # Handle list of content blocks (Anthropic extended thinking format)
+        if isinstance(content, list):
+            thinking_blocks = []
+            text_blocks = []
+            
+            for block in content:
+                if isinstance(block, dict):
+                    block_type = block.get('type', '')
+                    if block_type == 'thinking':
+                        thinking_blocks.append(block.get('thinking', ''))
+                    elif block_type == 'text':
+                        text_blocks.append(block.get('text', ''))
+                elif hasattr(block, 'type'):
+                    # Handle object format
+                    if block.type == 'thinking':
+                        thinking_blocks.append(getattr(block, 'thinking', ''))
+                    elif block.type == 'text':
+                        text_blocks.append(getattr(block, 'text', ''))
+            
+            if thinking_blocks:
+                result['thinking'] = '\n\n'.join(thinking_blocks)
+            if text_blocks:
+                result['text'] = '\n\n'.join(text_blocks)
+        
+        # Handle simple string content
+        elif isinstance(content, str):
+            result['text'] = content
+        
+        return result
+    
     def _run_async_in_sync_context(self, coro):
         """Run async coroutine from sync context.
         

alita_sdk/runtime/tools/loop.py

@@ -102,7 +102,9 @@ Input Data:
         logger.debug(f"LoopNode input: {predict_input}")
         completion = self.client.invoke(predict_input, config=config)
         logger.debug(f"LoopNode pure output: {completion}")
-        loop_data = _old_extract_json(completion.content.strip())
+        from ..langchain.utils import extract_text_from_completion
+        content_text = extract_text_from_completion(completion)
+        loop_data = _old_extract_json(content_text.strip())
         logger.debug(f"LoopNode output: {loop_data}")
         if self.return_type == "str":
             accumulated_response = ''

alita_sdk/runtime/tools/loop_output.py

@@ -93,7 +93,9 @@ Answer must be JSON only extractable by JSON.LOADS."""
         else:
             input_[-1].content += self.unstructured_output
             completion = self.client.invoke(input_, config=config)
-            result = _extract_json(completion.content.strip())
+            from ..langchain.utils import extract_text_from_completion
+            content_text = extract_text_from_completion(completion)
+            result = _extract_json(content_text.strip())
         try:
             tool_result: dict | List[dict] = self.tool.invoke(result, config=config, kwargs=kwargs)
             dispatch_custom_event(