alita-sdk 0.3.382__py3-none-any.whl → 0.3.383__py3-none-any.whl

alita_sdk/runtime/langchain/assistant.py

@@ -184,17 +184,6 @@ class Assistant:
         """
         # Exclude compiled graph runnables from simple tool agents
         simple_tools = [t for t in self.tools if isinstance(t, (BaseTool, CompiledStateGraph))]
-        # Add image generation tool if model is configured
-        if self.alita_client.model_image_generation is not None:
-            try:
-                from ..tools.image_generation import (
-                    create_image_generation_tool
-                )
-                image_tool = create_image_generation_tool(self.alita_client)
-                simple_tools.append(image_tool)
-                logger.info("Added ImageGenerationTool to react agent")
-            except Exception as e:
-                logger.error(f"Error adding ImageGenerationTool: {e}")
         
         # Set up memory/checkpointer if available
         checkpointer = None

alita_sdk/runtime/toolkits/tools.py

@@ -13,6 +13,7 @@ from .subgraph import SubgraphToolkit
 from .vectorstore import VectorStoreToolkit
 from ..tools.mcp_server_tool import McpServerTool
 from ..tools.sandbox import SandboxToolkit
+from ..tools.image_generation import ImageGenerationToolkit
 # Import community tools
 from ...community import get_toolkits as community_toolkits, get_tools as community_tools
 from ...tools.memory import MemoryToolkit
@@ -26,7 +27,8 @@ def get_toolkits():
         ArtifactToolkit.toolkit_config_schema(),
         MemoryToolkit.toolkit_config_schema(),
         VectorStoreToolkit.toolkit_config_schema(),
-        SandboxToolkit.toolkit_config_schema()
+        SandboxToolkit.toolkit_config_schema(),
+        ImageGenerationToolkit.toolkit_config_schema()
     ]
 
     return core_toolkits + community_toolkits() + alita_toolkits()
@@ -74,6 +76,14 @@ def get_tools(tools_list: list, alita_client, llm, memory_store: BaseStore = Non
                     stateful=False,
                     allow_net=True,
                 ).get_tools()
+            elif tool['name'] == 'image_generation':
+                if alita_client and alita_client.model_image_generation:
+                    tools += ImageGenerationToolkit.get_toolkit(
+                        client=alita_client,
+                    ).get_tools()
+                else:
+                    logger.warning("Image generation internal tool requested "
+                                   "but no image generation model configured")
         elif tool['type'] == 'artifact':
             tools.extend(ArtifactToolkit.get_toolkit(
                 client=alita_client,

alita_sdk/runtime/tools/__init__.py

@@ -5,7 +5,11 @@ This module provides various tools that can be used within LangGraph agents.
 
 from .sandbox import PyodideSandboxTool, StatefulPyodideSandboxTool, create_sandbox_tool
 from .echo import EchoTool
-from .image_generation import ImageGenerationTool, create_image_generation_tool
+from .image_generation import (
+    ImageGenerationTool,
+    create_image_generation_tool,
+    ImageGenerationToolkit
+)
 
 __all__ = [
     "PyodideSandboxTool",
@@ -13,5 +17,6 @@ __all__ = [
     "create_sandbox_tool",
     "EchoTool",
     "ImageGenerationTool",
+    "ImageGenerationToolkit",
     "create_image_generation_tool"
-]
+]

alita_sdk/runtime/tools/image_generation.py

@@ -2,16 +2,59 @@
 Image generation tool for Alita SDK.
 """
 import logging
-from typing import Optional, Type, Any
-from langchain_core.tools import BaseTool
-from pydantic import BaseModel, Field
+from typing import Optional, Type, Any, List, Literal
+from langchain_core.tools import BaseTool, BaseToolkit
+from pydantic import BaseModel, Field, create_model, ConfigDict
 
 logger = logging.getLogger(__name__)
 
+name = "image_generation"
+
+
+def get_tools(tools_list: list, alita_client=None, llm=None,
+              memory_store=None):
+    """
+    Get image generation tools for the provided tool configurations.
+
+    Args:
+        tools_list: List of tool configurations
+        alita_client: Alita client instance (required for image generation)
+        llm: LLM client instance (unused for image generation)
+        memory_store: Optional memory store instance (unused)
+
+    Returns:
+        List of image generation tools
+    """
+    all_tools = []
+
+    for tool in tools_list:
+        if (tool.get('type') == 'image_generation' or
+                tool.get('toolkit_name') == 'image_generation'):
+            try:
+                if not alita_client:
+                    logger.error("Alita client is required for image "
+                                 "generation tools")
+                    continue
+
+                toolkit_instance = ImageGenerationToolkit.get_toolkit(
+                    client=alita_client,
+                    toolkit_name=tool.get('toolkit_name', '')
+                )
+                all_tools.extend(toolkit_instance.get_tools())
+            except Exception as e:
+                logger.error(f"Error in image generation toolkit "
+                             f"get_tools: {e}")
+                logger.error(f"Tool config: {tool}")
+                raise
+
+    return all_tools
+
 
 class ImageGenerationInput(BaseModel):
     """Input schema for image generation tool."""
-    prompt: str = Field(description="Text prompt describing the image to generate")
+    prompt: str = Field(
+        description="Text prompt describing the image to generate"
+    )
     n: int = Field(
         default=1, description="Number of images to generate (1-10)",
         ge=1, le=10
@@ -22,7 +65,7 @@ class ImageGenerationInput(BaseModel):
     )
     quality: str = Field(
         default="auto",
-        description="Quality of the generated image ('low', 'medium', 'high', 'auto')"
+        description="Quality of the generated image ('low', 'medium', 'high')"
     )
     style: Optional[str] = Field(
         default=None, description="Style of the generated image (optional)"
@@ -69,7 +112,8 @@ class ImageGenerationTool(BaseTool):
                 else:
                     content_chunks.append({
                         "type": "text",
-                        "text": f"Generated {len(images)} images for prompt: '{prompt}'"
+                        "text": f"Generated {len(images)} images for "
+                                f"prompt: '{prompt}'"
                     })
                 
                 # Add image content for each generated image
@@ -85,7 +129,8 @@ class ImageGenerationTool(BaseTool):
                         content_chunks.append({
                             "type": "image_url",
                             "image_url": {
-                                "url": f"data:image/png;base64,{image_data['b64_json']}"
+                                "url": f"data:image/png;base64,"
+                                       f"{image_data['b64_json']}"
                             }
                         })
                 
@@ -94,7 +139,8 @@ class ImageGenerationTool(BaseTool):
             # Fallback to text response if no images in result
             return [{
                 "type": "text",
-                "text": f"Image generation completed but no images returned: {result}"
+                "text": f"Image generation completed but no images "
+                        f"returned: {result}"
             }]
             
         except Exception as e:
@@ -114,3 +160,53 @@ class ImageGenerationTool(BaseTool):
 def create_image_generation_tool(client):
     """Create an image generation tool with the provided Alita client."""
     return ImageGenerationTool(client=client)
+
+
+class ImageGenerationToolkit(BaseToolkit):
+    """Toolkit for image generation tools."""
+    tools: List[BaseTool] = []
+
+    @staticmethod
+    def toolkit_config_schema() -> BaseModel:
+        """Get the configuration schema for the image generation toolkit."""
+        # Create sample tool to get schema
+        sample_tool = ImageGenerationTool(client=None)
+        selected_tools = {sample_tool.name: sample_tool.args_schema.schema()}
+
+        return create_model(
+            'image_generation',
+            selected_tools=(
+                List[Literal[tuple(selected_tools)]],
+                Field(
+                    default=[],
+                    json_schema_extra={'args_schemas': selected_tools}
+                )
+            ),
+            __config__=ConfigDict(json_schema_extra={
+                'metadata': {
+                    "label": "Image Generation",
+                    "icon_url": "image_generation.svg",
+                    "hidden": True,
+                    "categories": ["internal_tool"],
+                    "extra_categories": ["image generation"],
+                }
+            })
+        )
+
+    @classmethod
+    def get_toolkit(cls, client=None, **kwargs):
+        """
+        Get toolkit with image generation tools.
+
+        Args:
+            client: Alita client instance (required)
+            **kwargs: Additional arguments
+        """
+        if not client:
+            raise ValueError("Alita client is required for image generation")
+            
+        tools = [ImageGenerationTool(client=client)]
+        return cls(tools=tools)
+
+    def get_tools(self):
+        return self.tools

{alita_sdk-0.3.382.dist-info → alita_sdk-0.3.383.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: alita_sdk
-Version: 0.3.382
+Version: 0.3.383
 Summary: SDK for building langchain agents using resources from Alita
 Author-email: Artem Rozumenko <artyom.rozumenko@gmail.com>, Mikalai Biazruchka <mikalai_biazruchka@epam.com>, Roman Mitusov <roman_mitusov@epam.com>, Ivan Krakhmaliuk <lifedj27@gmail.com>, Artem Dubrovskiy <ad13box@gmail.com>
 License-Expression: Apache-2.0

{alita_sdk-0.3.382.dist-info → alita_sdk-0.3.383.dist-info}/RECORD

@@ -41,7 +41,7 @@ alita_sdk/runtime/clients/datasource.py,sha256=HAZovoQN9jBg0_-lIlGBQzb4FJdczPhkH
 alita_sdk/runtime/clients/prompt.py,sha256=li1RG9eBwgNK_Qf0qUaZ8QNTmsncFrAL2pv3kbxZRZg,1447
 alita_sdk/runtime/clients/sandbox_client.py,sha256=OhEasE0MxBBDw4o76xkxVCpNpr3xJ8spQsrsVxMrjUA,16192
 alita_sdk/runtime/langchain/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-alita_sdk/runtime/langchain/assistant.py,sha256=HqFQ8nWrSdbXRx-huf0srA_zYCxcMjEndBCJVVriRMQ,15607
+alita_sdk/runtime/langchain/assistant.py,sha256=w6OUXNOWxWkju5n99Tfs_bbXlRsu397QRIlhig9SGmg,15041
 alita_sdk/runtime/langchain/chat_message_template.py,sha256=kPz8W2BG6IMyITFDA5oeb5BxVRkHEVZhuiGl4MBZKdc,2176
 alita_sdk/runtime/langchain/constants.py,sha256=eHVJ_beJNTf1WJo4yq7KMK64fxsRvs3lKc34QCXSbpk,3319
 alita_sdk/runtime/langchain/indexer.py,sha256=0ENHy5EOhThnAiYFc7QAsaTNp9rr8hDV_hTK8ahbatk,37592
@@ -102,9 +102,9 @@ alita_sdk/runtime/toolkits/configurations.py,sha256=kIDAlnryPQfbZyFxV-9SzN2-Vefz
 alita_sdk/runtime/toolkits/datasource.py,sha256=qk78OdPoReYPCWwahfkKLbKc4pfsu-061oXRryFLP6I,2498
 alita_sdk/runtime/toolkits/prompt.py,sha256=WIpTkkVYWqIqOWR_LlSWz3ug8uO9tm5jJ7aZYdiGRn0,1192
 alita_sdk/runtime/toolkits/subgraph.py,sha256=wwUK8JjPXkGzyVZ3tAukmvST6eGbqx_U11rpnmbrvtg,2105
-alita_sdk/runtime/toolkits/tools.py,sha256=Lo90zr6G_Mi378hpMNTRFmk05A7dDJDMbtVHEeFd5cM,8708
+alita_sdk/runtime/toolkits/tools.py,sha256=MV9sX7YTWqXbqhzhc7Xj3PG1Y_NZJ02r8XV8gg__WSI,9274
 alita_sdk/runtime/toolkits/vectorstore.py,sha256=BGppQADa1ZiLO17fC0uCACTTEvPHlodEDYEzUcBRbAA,2901
-alita_sdk/runtime/tools/__init__.py,sha256=TbHPnDtCdQvNzK1YQnk_ufkuI7FgHfvY1-JWUgycZhQ,497
+alita_sdk/runtime/tools/__init__.py,sha256=Fx7iHqkzA90-KfjdcUUzMUI_7kDarjuTsSpSzOW2pN0,568
 alita_sdk/runtime/tools/agent.py,sha256=m98QxOHwnCRTT9j18Olbb5UPS8-ZGeQaGiUyZJSyFck,3162
 alita_sdk/runtime/tools/application.py,sha256=z3vLZODs-_xEEnZFmGF0fKz1j3VtNJxqsAmg5ovExpQ,3129
 alita_sdk/runtime/tools/artifact.py,sha256=u3szFwZqguHrPZ3tZJ7S_TiZl7cxlT3oHYd6zbdpRDE,13842
@@ -112,7 +112,7 @@ alita_sdk/runtime/tools/datasource.py,sha256=pvbaSfI-ThQQnjHG-QhYNSTYRnZB0rYtZFp
 alita_sdk/runtime/tools/echo.py,sha256=spw9eCweXzixJqHnZofHE1yWiSUa04L4VKycf3KCEaM,486
 alita_sdk/runtime/tools/function.py,sha256=4r-VbGtm8gN_RTU2I-1iUNVs_MWNgRqn-aQrr__lCTc,7001
 alita_sdk/runtime/tools/graph.py,sha256=7jImBBSEdP5Mjnn2keOiyUwdGDFhEXLUrgUiugO3mgA,3503
-alita_sdk/runtime/tools/image_generation.py,sha256=8ZH4SoRrbS4EzmtF6cpNMRvuFephCYD2S8uqNC9KGE4,4274
+alita_sdk/runtime/tools/image_generation.py,sha256=Kls9D_ke_SK7xmVr7I9SlQcAEBJc86gf66haN0qIj9k,7469
 alita_sdk/runtime/tools/indexer_tool.py,sha256=whSLPevB4WD6dhh2JDXEivDmTvbjiMV1MrPl9cz5eLA,4375
 alita_sdk/runtime/tools/llm.py,sha256=AOxC3OEQ8FQCiI6f1we6s_X3a7DGyd2Z6yi1P_1jqoY,15376
 alita_sdk/runtime/tools/loop.py,sha256=uds0WhZvwMxDVFI6MZHrcmMle637cQfBNg682iLxoJA,8335
@@ -353,8 +353,8 @@ alita_sdk/tools/zephyr_scale/api_wrapper.py,sha256=kT0TbmMvuKhDUZc0i7KO18O38JM9S
 alita_sdk/tools/zephyr_squad/__init__.py,sha256=0ne8XLJEQSLOWfzd2HdnqOYmQlUliKHbBED5kW_Vias,2895
 alita_sdk/tools/zephyr_squad/api_wrapper.py,sha256=kmw_xol8YIYFplBLWTqP_VKPRhL_1ItDD0_vXTe_UuI,14906
 alita_sdk/tools/zephyr_squad/zephyr_squad_cloud_client.py,sha256=R371waHsms4sllHCbijKYs90C-9Yu0sSR3N4SUfQOgU,5066
-alita_sdk-0.3.382.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-alita_sdk-0.3.382.dist-info/METADATA,sha256=-Iut0sm4fsmy3rvpKk8q-fPFyzuIvdbBwy6UYkYztAU,19071
-alita_sdk-0.3.382.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-alita_sdk-0.3.382.dist-info/top_level.txt,sha256=0vJYy5p_jK6AwVb1aqXr7Kgqgk3WDtQ6t5C-XI9zkmg,10
-alita_sdk-0.3.382.dist-info/RECORD,,
+alita_sdk-0.3.383.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+alita_sdk-0.3.383.dist-info/METADATA,sha256=mtWM9Y6ppwLMV2IEeNtQlNYLyXsCdI0zRIP922y7jdY,19071
+alita_sdk-0.3.383.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+alita_sdk-0.3.383.dist-info/top_level.txt,sha256=0vJYy5p_jK6AwVb1aqXr7Kgqgk3WDtQ6t5C-XI9zkmg,10
+alita_sdk-0.3.383.dist-info/RECORD,,