ailoy-py 0.0.2__cp310-cp310-manylinux_2_28_x86_64.whl → 0.0.5__cp310-cp310-manylinux_2_28_x86_64.whl

ailoy/__init__.py

@@ -16,6 +16,7 @@ if __doc__ is None:
     else:  # fallback docstring
         __doc__ = "# ailoy-py\n\nPython binding for Ailoy runtime APIs"
 
-from .agent import Agent  # noqa: F401
+from .agent import Agent, AudioContent, BearerAuthenticator, ImageContent, TextContent, ToolAuthenticator  # noqa: F401
+from .models import APIModel, LocalModel  # noqa: F401
 from .runtime import AsyncRuntime, Runtime  # noqa: F401
 from .vector_store import VectorStore  # noqa: F401

ailoy/agent.py

@@ -1,69 +1,102 @@
+import base64
 import json
 import warnings
 from abc import ABC, abstractmethod
-from collections.abc import Awaitable, Callable, Generator
+from collections.abc import Callable, Generator
 from functools import partial
 from pathlib import Path
 from typing import (
+    Annotated,
     Any,
     Literal,
     Optional,
-    TypeVar,
     Union,
 )
 from urllib.parse import urlencode, urlparse, urlunparse
 
 import jmespath
-from pydantic import BaseModel, ConfigDict, Field
+from PIL.Image import Image
+from pydantic import BaseModel, ConfigDict, Field, TypeAdapter
 from rich.console import Console
 from rich.panel import Panel
 
 from ailoy.ailoy_py import generate_uuid
 from ailoy.mcp import MCPServer, MCPTool, StdioServerParameters
+from ailoy.models import APIModel, LocalModel
 from ailoy.runtime import Runtime
 from ailoy.tools import DocstringParsingException, TypeHintParsingException, get_json_schema
-
-__all__ = ["Agent"]
+from ailoy.utils.image import pillow_image_to_base64
 
 ## Types for internal data structures
 
 
-class TextData(BaseModel):
-    type: Literal["text"]
+class TextContent(BaseModel):
+    type: Literal["text"] = "text"
     text: str
 
 
+class ImageContent(BaseModel):
+    class UrlData(BaseModel):
+        url: str
+
+    type: Literal["image_url"] = "image_url"
+    image_url: UrlData
+
+    @staticmethod
+    def from_url(url: str):
+        return ImageContent(image_url={"url": url})
+
+    @staticmethod
+    def from_pillow(image: Image):
+        return ImageContent(image_url={"url": pillow_image_to_base64(image)})
+
+
+class AudioContent(BaseModel):
+    class AudioData(BaseModel):
+        data: str
+        format: Literal["mp3", "wav"]
+
+    type: Literal["input_audio"] = "input_audio"
+    input_audio: AudioData
+
+    @staticmethod
+    def from_bytes(data: bytes, format: Literal["mp3", "wav"]):
+        return AudioContent(input_audio={"data": base64.b64encode(data).decode("utf-8"), "format": format})
+
+
 class FunctionData(BaseModel):
     class FunctionBody(BaseModel):
         name: str
         arguments: Any
 
-    type: Literal["function"]
+    type: Literal["function"] = "function"
     id: Optional[str] = None
     function: FunctionBody
 
 
 class SystemMessage(BaseModel):
-    role: Literal["system"]
-    content: list[TextData]
+    role: Literal["system"] = "system"
+    content: str | list[TextContent]
 
 
 class UserMessage(BaseModel):
-    role: Literal["user"]
-    content: list[TextData]
+    role: Literal["user"] = "user"
+    content: str | list[TextContent | ImageContent | AudioContent]
 
 
 class AssistantMessage(BaseModel):
-    role: Literal["assistant"]
-    reasoning: Optional[list[TextData]] = None
-    content: Optional[list[TextData]] = None
+    role: Literal["assistant"] = "assistant"
+    content: Optional[str | list[TextContent]] = None
+    name: Optional[str] = None
     tool_calls: Optional[list[FunctionData]] = None
 
+    # Non-OpenAI fields
+    reasoning: Optional[list[TextContent]] = None
+
 
 class ToolMessage(BaseModel):
-    role: Literal["tool"]
-    name: str
-    content: list[TextData]
+    role: Literal["tool"] = "tool"
+    content: str | list[TextContent]
     tool_call_id: Optional[str] = None
 
 
@@ -76,72 +109,10 @@ Message = Union[
 
 
 class MessageOutput(BaseModel):
-    class AssistantMessageDelta(BaseModel):
-        content: Optional[list[TextData]] = None
-        reasoning: Optional[list[TextData]] = None
-        tool_calls: Optional[list[FunctionData]] = None
-
-    message: AssistantMessageDelta
+    message: AssistantMessage
     finish_reason: Optional[Literal["stop", "tool_calls", "invalid_tool_call", "length", "error"]] = None
 
 
-## Types for LLM Model Definitions
-
-TVMModelName = Literal["Qwen/Qwen3-0.6B", "Qwen/Qwen3-1.7B", "Qwen/Qwen3-4B", "Qwen/Qwen3-8B"]
-OpenAIModelName = Literal["gpt-4o"]
-ModelName = Union[TVMModelName, OpenAIModelName]
-
-
-class TVMModel(BaseModel):
-    name: TVMModelName
-    quantization: Optional[Literal["q4f16_1"]] = None
-    mode: Optional[Literal["interactive"]] = None
-
-
-class OpenAIModel(BaseModel):
-    name: OpenAIModelName
-    api_key: str
-
-
-class ModelDescription(BaseModel):
-    model_id: str
-    component_type: str
-    default_system_message: Optional[str] = None
-
-
-model_descriptions: dict[ModelName, ModelDescription] = {
-    "Qwen/Qwen3-0.6B": ModelDescription(
-        model_id="Qwen/Qwen3-0.6B",
-        component_type="tvm_language_model",
-        default_system_message="You are Qwen, created by Alibaba Cloud. You are a helpful assistant.",
-    ),
-    "Qwen/Qwen3-1.7B": ModelDescription(
-        model_id="Qwen/Qwen3-1.7B",
-        component_type="tvm_language_model",
-        default_system_message="You are Qwen, created by Alibaba Cloud. You are a helpful assistant.",
-    ),
-    "Qwen/Qwen3-4B": ModelDescription(
-        model_id="Qwen/Qwen3-4B",
-        component_type="tvm_language_model",
-        default_system_message="You are Qwen, created by Alibaba Cloud. You are a helpful assistant.",
-    ),
-    "Qwen/Qwen3-8B": ModelDescription(
-        model_id="Qwen/Qwen3-8B",
-        component_type="tvm_language_model",
-        default_system_message="You are Qwen, created by Alibaba Cloud. You are a helpful assistant.",
-    ),
-    "gpt-4o": ModelDescription(
-        model_id="gpt-4o",
-        component_type="openai",
-    ),
-}
-
-
-class ComponentState(BaseModel):
-    name: str
-    valid: bool
-
-
 ## Types for agent's responses
 
 _console = Console(highlight=False, force_jupyter=False, force_terminal=True)
@@ -149,7 +120,7 @@ _console = Console(highlight=False, force_jupyter=False, force_terminal=True)
 
 class AgentResponseOutputText(BaseModel):
     type: Literal["output_text", "reasoning"]
-    role: Literal["assistant"]
+    role: Literal["assistant"] = "assistant"
     is_type_switched: bool = False
     content: str
 
@@ -160,14 +131,14 @@ class AgentResponseOutputText(BaseModel):
 
 
 class AgentResponseToolCall(BaseModel):
-    type: Literal["tool_call"]
-    role: Literal["assistant"]
+    type: Literal["tool_call"] = "tool_call"
+    role: Literal["assistant"] = "assistant"
     is_type_switched: bool = False
     content: FunctionData
 
     def print(self):
         title = f"[magenta]Tool Call[/magenta]: [bold]{self.content.function.name}[/bold]"
-        if self.content.id is not None:
+        if self.content.id is not None and len(self.content.id) > 0:
             title += f" ({self.content.id})"
         panel = Panel(
             json.dumps(self.content.function.arguments, indent=2),
@@ -178,8 +149,8 @@ class AgentResponseToolCall(BaseModel):
 
 
 class AgentResponseToolResult(BaseModel):
-    type: Literal["tool_call_result"]
-    role: Literal["tool"]
+    type: Literal["tool_call_result"] = "tool_call_result"
+    role: Literal["tool"] = "tool"
     is_type_switched: bool = False
     content: ToolMessage
 
@@ -194,8 +165,8 @@ class AgentResponseToolResult(BaseModel):
         if len(content) > 500:
             content = content[:500] + "...(truncated)"
 
-        title = f"[green]Tool Result[/green]: [bold]{self.content.name}[/bold]"
-        if self.content.tool_call_id is not None:
+        title = "[green]Tool Result[/green]"
+        if self.content.tool_call_id is not None and len(self.content.tool_call_id) > 0:
             title += f" ({self.content.tool_call_id})"
         panel = Panel(
             content,
@@ -206,8 +177,8 @@ class AgentResponseToolResult(BaseModel):
 
 
 class AgentResponseError(BaseModel):
-    type: Literal["error"]
-    role: Literal["assistant"]
+    type: Literal["error"] = "error"
+    role: Literal["assistant"] = "assistant"
     is_type_switched: bool = False
     content: str
 
@@ -245,8 +216,11 @@ class ToolParameters(BaseModel):
     required: Optional[list[str]] = []
 
 
+JsonSchemaTypes = Literal["string", "integer", "number", "boolean", "object", "array", "null"]
+
+
 class ToolParametersProperty(BaseModel):
-    type: Literal["string", "number", "boolean", "object", "array", "null"]
+    type: JsonSchemaTypes | list[JsonSchemaTypes]
     description: Optional[str] = None
     model_config = ConfigDict(extra="allow")
 
@@ -308,22 +282,6 @@ class BearerAuthenticator(ToolAuthenticator):
         return {**request, "headers": headers}
 
 
-T_Retval = TypeVar("T_Retval")
-
-
-def run_async(coro: Callable[..., Awaitable[T_Retval]]) -> T_Retval:
-    try:
-        import anyio
-
-        # Running outside async loop
-        return anyio.run(lambda: coro)
-    except RuntimeError:
-        import anyio.from_thread
-
-        # Already in a running event loop: use anyio from_thread
-        return anyio.from_thread.run(coro)
-
-
 class Agent:
     """
     The `Agent` class provides a high-level interface for interacting with large language models (LLMs) in Ailoy.
@@ -337,28 +295,22 @@ class Agent:
     def __init__(
         self,
         runtime: Runtime,
-        model_name: ModelName,
+        model: APIModel | LocalModel,
         system_message: Optional[str] = None,
-        api_key: Optional[str] = None,
-        **attrs,
     ):
         """
         Create an instance.
 
         :param runtime: The runtime environment associated with the agent.
-        :param model_name: The name of the LLM model to use.
+        :param model: The model instance.
         :param system_message: Optional system message to set the initial assistant context.
-        :param api_key: (web agent only) The API key for AI API.
-        :param attrs: Additional initialization parameters (for `define_component` runtime call)
         :raises ValueError: If model name is not supported or validation fails.
         """
         self._runtime = runtime
 
         # Initialize component state
-        self._component_state = ComponentState(
-            name=generate_uuid(),
-            valid=False,
-        )
+        self._component_name = generate_uuid()
+        self._component_ready = False
 
         # Initialize messages
         self._messages: list[Message] = []
@@ -373,7 +325,7 @@ class Agent:
         self._mcp_servers: list[MCPServer] = []
 
         # Define the component
-        self.define(model_name, api_key=api_key, **attrs)
+        self.define(model)
 
     def __del__(self):
         self.delete()
@@ -384,70 +336,55 @@ class Agent:
     def __exit__(self, type, value, traceback):
         self.delete()
 
-    def define(self, model_name: ModelName, api_key: Optional[str] = None, **attrs) -> None:
+    def define(self, model: APIModel | LocalModel) -> None:
         """
         Initializes the agent by defining its model in the runtime.
         This must be called before running the agent. If already initialized, this is a no-op.
-        :param model_name: The name of the LLM model to use.
-        :param api_key: (web agent only) The API key for AI API.
-        :param attrs: Additional initialization parameters (for `define_component` runtime call)
+        :param model: The model instance.
         """
-        if self._component_state.valid:
+        if self._component_ready:
             return
 
         if not self._runtime.is_alive():
             raise ValueError("Runtime is currently stopped.")
 
-        if model_name not in model_descriptions:
-            raise ValueError(f"Model `{model_name}` not supported")
-
-        model_desc = model_descriptions[model_name]
-
-        # Add model name into attrs
-        if "model" not in attrs:
-            attrs["model"] = model_desc.model_id
-
         # Set default system message if not given; still can be None
         if self._system_message is None:
-            self._system_message = model_desc.default_system_message
+            self._system_message = getattr(model, "default_system_message", None)
 
         self.clear_messages()
 
-        # Add API key
-        if api_key:
-            attrs["api_key"] = api_key
-
         # Call runtime's define
         self._runtime.define(
-            model_descriptions[model_name].component_type,
-            self._component_state.name,
-            attrs,
+            model.component_type,
+            self._component_name,
+            model.to_attrs(),
         )
 
         # Mark as defined
-        self._component_state.valid = True
+        self._component_ready = True
 
     def delete(self) -> None:
         """
         Deinitializes the agent and releases resources in the runtime.
         This should be called when the agent is no longer needed. If already deinitialized, this is a no-op.
         """
-        if not self._component_state.valid:
+        if not self._component_ready:
             return
 
         if self._runtime.is_alive():
-            self._runtime.delete(self._component_state.name)
+            self._runtime.delete(self._component_name)
 
         self.clear_messages()
 
         for mcp_server in self._mcp_servers:
             mcp_server.cleanup()
 
-        self._component_state.valid = False
+        self._component_ready = False
 
     def query(
         self,
-        message: str,
+        message: str | list[str | Image | dict | TextContent | ImageContent | AudioContent],
         reasoning: bool = False,
     ) -> Generator[AgentResponse, None, None]:
         """
@@ -458,13 +395,36 @@ class Agent:
         :return: An iterator over the output, where each item represents either a generated token from the assistant or a tool call.
         :rtype: Iterator[:class:`AgentResponse`]
         """  # noqa: E501
-        if not self._component_state.valid:
+        if not self._component_ready:
             raise ValueError("Agent is not valid. Create one or define newly.")
 
         if not self._runtime.is_alive():
             raise ValueError("Runtime is currently stopped.")
 
-        self._messages.append(UserMessage(role="user", content=[{"type": "text", "text": message}]))
+        if isinstance(message, str):
+            self._messages.append(UserMessage(content=[TextContent(text=message)]))
+        elif isinstance(message, list):
+            if len(message) == 0:
+                raise ValueError("Message is empty")
+
+            contents = []
+            for content in message:
+                if isinstance(content, str):
+                    contents.append(TextContent(text=content))
+                elif isinstance(content, Image):
+                    contents.append(ImageContent.from_pillow(image=content))
+                elif isinstance(content, dict):
+                    ta: TypeAdapter[TextContent | ImageContent | AudioContent] = TypeAdapter(
+                        Annotated[TextContent | ImageContent | AudioContent, Field(discriminator="type")]
+                    )
+                    validated_content = ta.validate_python(content)
+                    contents.append(validated_content)
+                else:
+                    contents.append(content)
+
+            self._messages.append(UserMessage(content=contents))
+        else:
+            raise ValueError(f"Invalid message type: {type(message)}")
 
         prev_resp_type = None
 
@@ -480,7 +440,7 @@ class Agent:
             assistant_content = None
             assistant_tool_calls = None
             finish_reason = ""
-            for result in self._runtime.call_iter_method(self._component_state.name, "infer", infer_args):
+            for result in self._runtime.call_iter_method(self._component_name, "infer", infer_args):
                 msg = MessageOutput.model_validate(result)
 
                 if msg.message.reasoning:
@@ -491,13 +451,16 @@ class Agent:
                             assistant_reasoning[0].text += v.text
                         resp = AgentResponseOutputText(
                             type="reasoning",
-                            role="assistant",
                             is_type_switched=(prev_resp_type != "reasoning"),
                             content=v.text,
                         )
                         prev_resp_type = resp.type
                         yield resp
-                if msg.message.content:
+                if msg.message.content is not None:
+                    # Canonicalize message content to the array of TextContent
+                    if isinstance(msg.message.content, str):
+                        msg.message.content = [TextContent(text=msg.message.content)]
+
                     for v in msg.message.content:
                         if not assistant_content:
                             assistant_content = [v]
@@ -505,7 +468,6 @@ class Agent:
                             assistant_content[0].text += v.text
                         resp = AgentResponseOutputText(
                             type="output_text",
-                            role="assistant",
                             is_type_switched=(prev_resp_type != "output_text"),
                             content=v.text,
                         )
@@ -518,8 +480,6 @@ class Agent:
                         else:
                             assistant_tool_calls.append(v)
                         resp = AgentResponseToolCall(
-                            type="tool_call",
-                            role="assistant",
                             is_type_switched=True,
                             content=v,
                         )
@@ -532,7 +492,6 @@ class Agent:
             # Append output
             self._messages.append(
                 AssistantMessage(
-                    role="assistant",
                     reasoning=assistant_reasoning,
                     content=assistant_content,
                     tool_calls=assistant_tool_calls,
@@ -550,18 +509,16 @@ class Agent:
                         raise RuntimeError("Tool not found")
                     tool_result = tool_.call(**tool_call.function.arguments)
                     return ToolMessage(
-                        role="tool",
-                        name=tool_call.function.name,
-                        content=[TextData(type="text", text=json.dumps(tool_result))],
-                        tool_call_id=tool_call.id if tool_call.id else None,
+                        content=[
+                            TextContent(text=tool_result if isinstance(tool_result, str) else json.dumps(tool_result))
+                        ],
+                        tool_call_id=tool_call.id,
                     )
 
                 tool_call_results = [run_tool(tc) for tc in assistant_tool_calls]
                 for result_msg in tool_call_results:
                     self._messages.append(result_msg)
                     resp = AgentResponseToolResult(
-                        type="tool_call_result",
-                        role="tool",
                         is_type_switched=True,
                         content=result_msg,
                     )
@@ -571,6 +528,7 @@ class Agent:
                 continue
 
             # Finish this generator
+            yield AgentResponseOutputText(type="output_text", content="\n")
             break
 
     def get_messages(self) -> list[Message]:
@@ -589,9 +547,7 @@ class Agent:
         """
         self._messages.clear()
         if self._system_message is not None:
-            self._messages.append(
-                SystemMessage(role="system", content=[TextData(type="text", text=self._system_message)])
-            )
+            self._messages.append(SystemMessage(role="system", content=[TextContent(text=self._system_message)]))
 
     def print(self, resp: AgentResponse):
         resp.print()
@@ -779,7 +735,7 @@ class Agent:
                 continue
 
             desc = ToolDescription(
-                name=f"{name}/{tool.name}", description=tool.description, parameters=tool.inputSchema
+                name=f"{name}-{tool.name}", description=tool.description, parameters=tool.inputSchema
             )
 
             def call(tool: MCPTool, **inputs: dict[str, Any]) -> list[str]:
@@ -803,4 +759,16 @@ class Agent:
         mcp_server.cleanup()
 
         # Remove tools registered from the MCP server
-        self._tools = list(filter(lambda t: not t.desc.name.startswith(f"{mcp_server.name}/"), self._tools))
+        self._tools = list(filter(lambda t: not t.desc.name.startswith(f"{mcp_server.name}-"), self._tools))
+
+    def get_tools(self):
+        """
+        Get the list of registered tools.
+        """
+        return self._tools
+
+    def clear_tools(self):
+        """
+        Clear the registered tools.
+        """
+        self._tools.clear()

ailoy/mcp.py

@@ -2,7 +2,7 @@ import asyncio
 import json
 import multiprocessing
 import platform
-import subprocess
+import tempfile
 from multiprocessing.connection import Connection
 from typing import Annotated, Any, Literal, Union
 
@@ -13,6 +13,7 @@ from mcp.client.stdio import (
     StdioServerParameters,
     stdio_client,
 )
+from mcp.shared.exceptions import McpError
 from pydantic import BaseModel, Field, TypeAdapter
 
 __all__ = ["MCPServer"]
@@ -73,11 +74,15 @@ class MCPServer:
         self._parent_conn, self._child_conn = multiprocessing.Pipe()
 
         ctx = multiprocessing.get_context("fork" if platform.system() != "Windows" else "spawn")
-        self._proc = ctx.Process(target=self._run_process, args=(self._child_conn,))
+        self._proc: multiprocessing.Process = ctx.Process(target=self._run_process, args=(self._child_conn,))
         self._proc.start()
 
         # Wait for subprocess to signal initialization complete
-        self._recv_response()
+        try:
+            self._recv_response()
+        except RuntimeError as e:
+            self.cleanup()
+            raise e
 
     def __del__(self):
         self.cleanup()
@@ -86,52 +91,59 @@ class MCPServer:
         asyncio.run(self._process_main(conn))
 
     async def _process_main(self, conn: Connection):
-        async with stdio_client(self.params, errlog=subprocess.PIPE) as (read, write):
-            async with ClientSession(read, write) as session:
-                # Notify to main process that the initialization has been finished and ready to receive requests
-                try:
-                    await session.initialize()
-                    conn.send(ResultMessage(result=True).model_dump())
-                except Exception as e:
-                    conn.send(ErrorMessage(error=f"Failed to initialize MCP subprocess: {e}").model_dump())
-
-                while True:
-                    if not conn.poll(0.1):
-                        await asyncio.sleep(0.1)
-                        continue
-
+        with tempfile.TemporaryFile(mode="w+t") as _errlog:
+            async with stdio_client(self.params, errlog=_errlog) as (read, write):
+                async with ClientSession(read, write) as session:
+                    # Notify to main process that the initialization has been finished and ready to receive requests
                     try:
-                        raw = conn.recv()
-                        req = TypeAdapter(RequestMessage).validate_python(raw)
-
-                        if isinstance(req, ListToolsRequest):
-                            result = await session.list_tools()
-                            conn.send(ResultMessage(result=result.tools).model_dump())
-
-                        elif isinstance(req, CallToolRequest):
-                            result = await session.call_tool(req.tool.name, req.arguments)
-                            contents: list[str] = []
-                            for item in result.content:
-                                if isinstance(item, mcp_types.TextContent):
-                                    try:
-                                        content = json.loads(item.text)
-                                        contents.append(json.dumps(content))
-                                    except json.JSONDecodeError:
-                                        contents.append(item.text)
-                                elif isinstance(item, mcp_types.ImageContent):
-                                    contents.append(item.data)
-                                elif isinstance(item, mcp_types.EmbeddedResource):
-                                    if isinstance(item.resource, mcp_types.TextResourceContents):
-                                        contents.append(item.resource.text)
-                                    else:
-                                        contents.append(item.resource.blob)
-                            conn.send(ResultMessage(result=contents).model_dump())
-
-                        elif isinstance(req, ShutdownRequest):
-                            break
-
-                    except Exception as e:
-                        conn.send(ErrorMessage(error=str(e)).model_dump())
+                        await session.initialize()
+                        conn.send(ResultMessage(result=True).model_dump())
+                    except McpError:
+                        _errlog.seek(0)
+                        error = _errlog.read()
+                        conn.send(
+                            ErrorMessage(
+                                error=f"Failed to initialize MCP subprocess. Check the error output below.\n\n{error}"
+                            ).model_dump()
+                        )
+
+                    while True:
+                        if not conn.poll(0.1):
+                            await asyncio.sleep(0.1)
+                            continue
+
+                        try:
+                            raw = conn.recv()
+                            req = TypeAdapter(RequestMessage).validate_python(raw)
+
+                            if isinstance(req, ListToolsRequest):
+                                result = await session.list_tools()
+                                conn.send(ResultMessage(result=result.tools).model_dump())
+
+                            elif isinstance(req, CallToolRequest):
+                                result = await session.call_tool(req.tool.name, req.arguments)
+                                contents: list[str] = []
+                                for item in result.content:
+                                    if isinstance(item, mcp_types.TextContent):
+                                        try:
+                                            content = json.loads(item.text)
+                                            contents.append(json.dumps(content))
+                                        except json.JSONDecodeError:
+                                            contents.append(item.text)
+                                    elif isinstance(item, mcp_types.ImageContent):
+                                        contents.append(item.data)
+                                    elif isinstance(item, mcp_types.EmbeddedResource):
+                                        if isinstance(item.resource, mcp_types.TextResourceContents):
+                                            contents.append(item.resource.text)
+                                        else:
+                                            contents.append(item.resource.blob)
+                                conn.send(ResultMessage(result=contents).model_dump())
+
+                            elif isinstance(req, ShutdownRequest):
+                                break
+
+                        except Exception as e:
+                            conn.send(ErrorMessage(error=str(e)).model_dump())
 
     def _send_request(self, msg: RequestMessage):
         self._parent_conn.send(msg.model_dump())

ailoy/models/__init__.py

@@ -0,0 +1,7 @@
+from .api_model import APIModel
+from .local_model import LocalModel
+
+__all__ = [
+    "APIModel",
+    "LocalModel",
+]

ailoy/models/api_model.py

@@ -0,0 +1,86 @@
+from typing import Literal, Optional, get_args
+
+from pydantic import model_validator
+from pydantic.dataclasses import dataclass
+
+OpenAIModelId = Literal[
+    "o4-mini",
+    "o3",
+    "o3-pro",
+    "o3-mini",
+    "gpt-4o",
+    "gpt-4o-mini",
+    "gpt-4.1",
+    "gpt-4.1-mini",
+    "gpt-4.1-nano",
+]
+
+GeminiModelId = Literal[
+    "gemini-2.5-flash",
+    "gemini-2.5-pro",
+    "gemini-2.0-flash",
+    "gemini-1.5-flash",
+    "gemini-1.5-pro",
+]
+
+ClaudeModelId = Literal[
+    "claude-sonnet-4-20250514",
+    "claude-3-7-sonnet-20250219",
+    "claude-3-5-sonnet-20241022",
+    "claude-3-5-sonnet-20240620",
+    "claude-opus-4-20250514",
+    "claude-3-opus-20240229",
+    "claude-3-5-haiku-20241022",
+    "claude-3-haiku-20240307",
+]
+
+GrokModelId = Literal[
+    "grok-4",
+    "grok-4-0709",
+    "grok-3",
+    "grok-3-fast",
+    "grok-3-mini",
+    "grok-3-mini-fast",
+    "grok-2",
+    "grok-2-1212",
+    "grok-2-vision-1212",
+    "grok-2-image-1212",
+]
+
+APIModelProvider = Literal["openai", "gemini", "claude", "grok"]
+
+
+@dataclass
+class APIModel:
+    id: OpenAIModelId | GeminiModelId | ClaudeModelId | str
+    api_key: str
+    provider: Optional[APIModelProvider] = None
+
+    @model_validator(mode="after")
+    def validate_provider(self):
+        if self.provider is None:
+            if self.id in get_args(OpenAIModelId):
+                self.provider = "openai"
+            elif self.id in get_args(GeminiModelId):
+                self.provider = "gemini"
+            elif self.id in get_args(ClaudeModelId):
+                self.provider = "claude"
+            elif self.id in get_args(GrokModelId):
+                self.provider = "grok"
+            else:
+                raise ValueError(
+                    f'Failed to infer the model provider based on the model id "{self.id}". '
+                    "Please provide an explicit model provider."
+                )
+
+        return self
+
+    @property
+    def component_type(self) -> str:
+        return self.provider
+
+    def to_attrs(self):
+        return {
+            "model": self.id,
+            "api_key": self.api_key,
+        }

ailoy/models/local_model.py

@@ -0,0 +1,44 @@
+from typing import Literal, Optional
+
+from pydantic.dataclasses import dataclass
+
+LocalModelBackend = Literal["tvm"]
+LocalModelId = Literal[
+    "Qwen/Qwen3-0.6B",
+    "Qwen/Qwen3-1.7B",
+    "Qwen/Qwen3-4B",
+    "Qwen/Qwen3-8B",
+    "Qwen/Qwen3-14B",
+    "Qwen/Qwen3-32B",
+    "Qwen/Qwen3-30B-A3B",
+]
+Quantization = Literal["q4f16_1"]
+
+
+@dataclass
+class LocalModel:
+    id: LocalModelId
+    backend: LocalModelBackend = "tvm"
+    quantization: Quantization = "q4f16_1"
+    device: int = 0
+
+    @property
+    def default_system_message(self) -> Optional[str]:
+        if self.id.startswith("Qwen"):
+            return "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."
+        return None
+
+    @property
+    def component_type(self) -> str:
+        if self.backend == "tvm":
+            return "tvm_language_model"
+        raise ValueError(f"Unknown local model backend: {self.backend}")
+
+    def to_attrs(self) -> dict:
+        if self.backend == "tvm":
+            return {
+                "model": self.id,
+                "quantization": self.quantization,
+                "device": self.device,
+            }
+        raise ValueError(f"Unknown local model backend: {self.backend}")

ailoy/tools.py

@@ -5,6 +5,7 @@ import types
 from typing import (
     Any,
     Callable,
+    Literal,
     Optional,
     Union,
     get_args,
@@ -41,7 +42,7 @@ class DocstringParsingException(Exception):
     pass
 
 
-def _get_json_schema_type(param_type: str) -> dict[str, str]:
+def _get_json_schema_type(param_type: type) -> dict[str, str]:
     type_mapping = {
         int: {"type": "integer"},
         float: {"type": "number"},
@@ -85,6 +86,20 @@ def _parse_type_hint(hint: str) -> dict:
             return_dict["nullable"] = True
         return return_dict
 
+    elif origin is Literal and len(args) > 0:
+        LITERAL_TYPES = (int, float, str, bool, type(None))
+        args_types = []
+        for arg in args:
+            if type(arg) not in LITERAL_TYPES:
+                raise TypeHintParsingException("Only the valid python literals can be listed in typing.Literal.")
+            arg_type = _get_json_schema_type(type(arg)).get("type")
+            if arg_type is not None and arg_type not in args_types:
+                args_types.append(arg_type)
+        return {
+            "type": args_types.pop() if len(args_types) == 1 else list(args_types),
+            "enum": list(args),
+        }
+
     elif origin is list:
         if not args:
             return {"type": "array"}
@@ -100,13 +115,13 @@ def _parse_type_hint(hint: str) -> dict:
                 f"The type hint {str(hint).replace('typing.', '')} is a Tuple with a single element, which "
                 "we do not automatically convert to JSON schema as it is rarely necessary. If this input can contain "
                 "more than one element, we recommend "
-                "using a List[] type instead, or if it really is a single element, remove the Tuple[] wrapper and just "
+                "using a list[] type instead, or if it really is a single element, remove the tuple[] wrapper and just "
                 "pass the element directly."
             )
         if ... in args:
             raise TypeHintParsingException(
                 "Conversion of '...' is not supported in Tuple type hints. "
-                "Use List[] types for variable-length"
+                "Use list[] types for variable-length"
                 " inputs instead."
             )
         return {"type": "array", "prefixItems": [_parse_type_hint(t) for t in args]}

ailoy/utils/image.py

@@ -0,0 +1,11 @@
+import base64
+import io
+
+from PIL.Image import Image
+
+
+def pillow_image_to_base64(img: Image):
+    buffered = io.BytesIO()
+    img.save(buffered, format=img.format)
+    b64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
+    return f"data:image/{img.format.lower()};base64,{b64}"

{ailoy_py-0.0.2.dist-info → ailoy_py-0.0.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ailoy-py
-Version: 0.0.2
+Version: 0.0.5
 Summary: Python binding for Ailoy runtime APIs
 Author-Email: "Brekkylab Inc." <contact@brekkylab.com>
 License-Expression: Apache-2.0
@@ -18,6 +18,7 @@ Requires-Dist: anyio>=4.9.0
 Requires-Dist: jmespath>=1.0.1
 Requires-Dist: mcp>=1.8.0
 Requires-Dist: numpy>=2.0.2
+Requires-Dist: pillow>=11.2.1
 Requires-Dist: pydantic>=2.11.4
 Requires-Dist: rich>=14.0.0
 Requires-Dist: typer>=0.15.4
@@ -38,14 +39,14 @@ pip install ailoy-py
 ## Quickstart
 
 ```python
-from ailoy import Runtime, Agent
+from ailoy import Runtime, Agent, LocalModel
 
 # The runtime must be started to use Ailoy
 rt = Runtime()
 
 # Defines an agent
 # During this step, the model parameters are downloaded and the LLM is set up for execution
-with Agent(rt, model_name="Qwen/Qwen3-0.6B") as agent:
+with Agent(rt, LocalModel("Qwen/Qwen3-0.6B")) as agent:
     # This is where the actual LLM call happens
     for resp in agent.query("Please give me a short poem about AI"):
         agent.print(resp)

ailoy_py-0.0.5.dist-info/RECORD

@@ -0,0 +1,25 @@
+ailoy/__init__.py,sha256=INeFsnHtrqTwkMKUPSdD_9l8DhudnW574e9rsRtrMJc,783
+ailoy/agent.py,sha256=7v2E1ft1Aaufem9lxLEra3S-nX2fzeutJDgT-4LXwL8,26954
+ailoy/ailoy_py.cpython-310-x86_64-linux-gnu.so,sha256=Tkdn4uQLJYq-0SO0dyFx6Q5zVWCA9qg7j5VyKWeQrKs,22810457
+ailoy/ailoy_py.pyi,sha256=Yf90FEXkslpCpr1r2eqQ3-_1jLo65zmG94bBXDRqinU,991
+ailoy/mcp.py,sha256=wVzXfwUh4UcU60PYq17kCFG7ZClmEDBRt8LxetuSkns,6800
+ailoy/runtime.py,sha256=-75KawEMQSwxGvX5wtECVCWiTNdcHojsQ1e-OVB4IQ8,10545
+ailoy/tools.py,sha256=hLdKe3TN_yn2qSbNG0uX2extokPIdJ6inLbSbhXdYTo,8861
+ailoy/vector_store.py,sha256=ZfIuGYKv2dQmjOuDlSKDc-BBPlQ8no_70mZwnPzbBzo,7515
+ailoy/cli/__main__.py,sha256=HnBVb2em1F2NLPeNX5r3xRndRrnGaXVCduo8WBULAI0,179
+ailoy/cli/model.py,sha256=cerCHE-VY9TOwqRcLBtmqnV-5vphpvyhtrfPFZiTKCM,2979
+ailoy/models/__init__.py,sha256=1AtlJV9gYThw_3snu0jPEH_aQGI74ip7ZcVJLtN5nMU,117
+ailoy/models/api_model.py,sha256=kLDD1-R4hjgtTbIQkD-boT-ZK8UocmVyyHkI6tuxiFU,2090
+ailoy/models/local_model.py,sha256=Iyur0UEUSbLKzptx9croP_OAF-qh9S-ZDukDthHNz9w,1206
+ailoy/presets/tools/calculator.json,sha256=ePnZsjZChnvS08s9eVdIp4Bys_PlJBXPHCCjv6oMvzA,1040
+ailoy/presets/tools/frankfurter.json,sha256=bZ5vhszf_aR-B_QN4L2xrI5nR-f4AMZk41UUDq1dTXg,1152
+ailoy/presets/tools/nytimes.json,sha256=wrfe9bnAlSPzHladoGEX2oCAeE0wed3BvgXQ_Z2PdXg,918
+ailoy/presets/tools/tmdb.json,sha256=UGLN5uAJ2b-Hu3nLcW95WXDLB3mfC3rBYfQANp_e8Ps,7046
+ailoy/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ailoy/utils/image.py,sha256=zkufequcQVTmIkArreYUyB8r2nrcOL_8O6KOv5B-yis,288
+ailoy_py.libs/libgomp-870cb1d0.so.1.0.0,sha256=Ta6ZPLbakQH8LP74JzBt0DuJIBHS4nicjkSCjKnyWDw,253289
+ailoy_py.libs/libtvm_runtime-c9b1997b.so,sha256=OLiYdpcijjzZc_5wzEpaEkxSiavSA5JCfuhcmiWWG-4,5617177
+ailoy_py-0.0.5.dist-info/METADATA,sha256=WpOovIib4t4AWYNmIBRRixucdz9-wBetyxyDaLjzOk4,2053
+ailoy_py-0.0.5.dist-info/WHEEL,sha256=zoVaZapIUnOpzt0hU8TDEPq4sy8rbtXAzSVdE0SY4gI,118
+ailoy_py-0.0.5.dist-info/entry_points.txt,sha256=gVG45uDE6kef0wm6SEMYSgZgRNNRhSAeP2n2lPR00dI,50
+ailoy_py-0.0.5.dist-info/RECORD,,

{ailoy_py-0.0.2.dist-info → ailoy_py-0.0.5.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: scikit-build-core 0.11.4
+Generator: scikit-build-core 0.11.5
 Root-Is-Purelib: false
 Tag: cp310-cp310-manylinux_2_28_x86_64
 

ailoy_py-0.0.2.dist-info/RECORD

@@ -1,20 +0,0 @@
-ailoy/__init__.py,sha256=mzkLUc95OCc2okURWm9iA5xR8WZdxwvPgaanc9fwoH4,647
-ailoy/agent.py,sha256=uQ1o4CjQEO1vP7y0frGtVknN-A_iLNtkqt8h3vobgiM,27613
-ailoy/ailoy_py.cpython-310-x86_64-linux-gnu.so,sha256=vcWOM7sIPwrsURLwp2A8dp9ENkJzb3TuWN1XXuXSwC0,21241233
-ailoy/ailoy_py.pyi,sha256=Yf90FEXkslpCpr1r2eqQ3-_1jLo65zmG94bBXDRqinU,991
-ailoy/mcp.py,sha256=bC58tAWqhvMdZVCKHSdOVNUoAuYfZiou1hSH1oa_9Ag,6190
-ailoy/runtime.py,sha256=-75KawEMQSwxGvX5wtECVCWiTNdcHojsQ1e-OVB4IQ8,10545
-ailoy/tools.py,sha256=RnTfmWlqYY1q0V377CpAAyAK-yET7k45GgEhgM9G8eI,8207
-ailoy/vector_store.py,sha256=ZfIuGYKv2dQmjOuDlSKDc-BBPlQ8no_70mZwnPzbBzo,7515
-ailoy/cli/__main__.py,sha256=HnBVb2em1F2NLPeNX5r3xRndRrnGaXVCduo8WBULAI0,179
-ailoy/cli/model.py,sha256=cerCHE-VY9TOwqRcLBtmqnV-5vphpvyhtrfPFZiTKCM,2979
-ailoy/presets/tools/calculator.json,sha256=ePnZsjZChnvS08s9eVdIp4Bys_PlJBXPHCCjv6oMvzA,1040
-ailoy/presets/tools/frankfurter.json,sha256=bZ5vhszf_aR-B_QN4L2xrI5nR-f4AMZk41UUDq1dTXg,1152
-ailoy/presets/tools/nytimes.json,sha256=wrfe9bnAlSPzHladoGEX2oCAeE0wed3BvgXQ_Z2PdXg,918
-ailoy/presets/tools/tmdb.json,sha256=UGLN5uAJ2b-Hu3nLcW95WXDLB3mfC3rBYfQANp_e8Ps,7046
-ailoy_py.libs/libgomp-870cb1d0.so.1.0.0,sha256=Ta6ZPLbakQH8LP74JzBt0DuJIBHS4nicjkSCjKnyWDw,253289
-ailoy_py.libs/libtvm_runtime-2d14ca42.so,sha256=qPtn3HaKtxt-sL0wdu6Wqz7QsTmKY2ZWOPwO92TPfzU,5061889
-ailoy_py-0.0.2.dist-info/METADATA,sha256=B5RbxeITquJfdiw9bhA6w02q9OvLkuFH7jMRg6Lxc2A,2010
-ailoy_py-0.0.2.dist-info/WHEEL,sha256=VVjTMh1gjoiSjlh90KmkjL10Nq1LgP-3G8MTl99dyME,118
-ailoy_py-0.0.2.dist-info/entry_points.txt,sha256=gVG45uDE6kef0wm6SEMYSgZgRNNRhSAeP2n2lPR00dI,50
-ailoy_py-0.0.2.dist-info/RECORD,,