ailoy-py 0.0.1__cp311-cp311-win_amd64.whl → 0.0.3__cp311-cp311-win_amd64.whl

ailoy/mcp.py

@@ -0,0 +1,171 @@
+import asyncio
+import json
+import multiprocessing
+import platform
+import tempfile
+from multiprocessing.connection import Connection
+from typing import Annotated, Any, Literal, Union
+
+import mcp.types as mcp_types
+from mcp import Tool as MCPTool
+from mcp.client.session import ClientSession
+from mcp.client.stdio import (
+    StdioServerParameters,
+    stdio_client,
+)
+from mcp.shared.exceptions import McpError
+from pydantic import BaseModel, Field, TypeAdapter
+
+__all__ = ["MCPServer"]
+
+
+class ListToolsRequest(BaseModel):
+    type: Literal["list_tools"] = "list_tools"
+
+
+class CallToolRequest(BaseModel):
+    type: Literal["call_tool"] = "call_tool"
+    tool: MCPTool
+    arguments: dict[str, Any]
+
+
+class ShutdownRequest(BaseModel):
+    type: Literal["shutdown"] = "shutdown"
+
+
+# Requests (main -> subprocess)
+RequestMessage = Annotated[Union[ListToolsRequest, CallToolRequest, ShutdownRequest], Field(discriminator="type")]
+
+
+class ResultMessage(BaseModel):
+    type: Literal["result"] = "result"
+    result: Any
+
+
+class ErrorMessage(BaseModel):
+    type: Literal["error"] = "error"
+    error: str
+
+
+# Response (subprocess -> main)
+ResponseMessage = Annotated[Union[ResultMessage, ErrorMessage], Field(discriminator="type")]
+
+
+class MCPServer:
+    """
+    MCPServer manages a subprocess that acts as a bridge between an MCP stdio server and the main process.
+
+    - The subprocess communicates with the MCP stdio server using the official MCP Python SDK.
+    - Communication between the main process and the subprocess is handled through a multiprocessing Pipe.
+      Messages sent over this Pipe are serialized and deserialized using structured Pydantic models:
+        - `RequestMessage` for requests from the main process to the subprocess.
+        - `ResponseMessage` for responses from the subprocess to the main process.
+
+    This design ensures:
+    - Type-safe, structured inter-process communication.
+    - Synchronous interaction with an asynchronous MCP session (via message passing).
+    - Subprocess lifecycle control (including initialization and shutdown).
+    """
+
+    def __init__(self, name: str, params: StdioServerParameters):
+        self.name = name
+        self.params = params
+
+        self._parent_conn, self._child_conn = multiprocessing.Pipe()
+
+        ctx = multiprocessing.get_context("fork" if platform.system() != "Windows" else "spawn")
+        self._proc: multiprocessing.Process = ctx.Process(target=self._run_process, args=(self._child_conn,))
+        self._proc.start()
+
+        # Wait for subprocess to signal initialization complete
+        try:
+            self._recv_response()
+        except RuntimeError as e:
+            self.cleanup()
+            raise e
+
+    def __del__(self):
+        self.cleanup()
+
+    def _run_process(self, conn: Connection):
+        asyncio.run(self._process_main(conn))
+
+    async def _process_main(self, conn: Connection):
+        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:
+                        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())
+
+    def _recv_response(self) -> ResultMessage:
+        raw = self._parent_conn.recv()
+        msg = TypeAdapter(ResponseMessage).validate_python(raw)
+        if isinstance(msg, ErrorMessage):
+            raise RuntimeError(msg.error)
+        return msg
+
+    def list_tools(self) -> list[MCPTool]:
+        self._send_request(ListToolsRequest())
+        msg = self._recv_response()
+        return [MCPTool.model_validate(tool) for tool in msg.result]
+
+    def call_tool(self, tool: MCPTool, arguments: dict[str, Any]) -> list[str]:
+        self._send_request(CallToolRequest(tool=tool, arguments=arguments))
+        msg = self._recv_response()
+        return msg.result
+
+    def cleanup(self) -> None:
+        if self._proc.is_alive():
+            self._send_request(ShutdownRequest())
+            self._proc.join()

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,71 @@
+from typing import Literal, Optional, Self, 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",
+]
+
+APIModelProvider = Literal["openai", "gemini", "claude"]
+
+
+@dataclass
+class APIModel:
+    id: OpenAIModelId | GeminiModelId | ClaudeModelId | str
+    api_key: str
+    provider: Optional[APIModelProvider] = None
+
+    @model_validator(mode="after")
+    def validate_provider(self) -> 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"
+            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/runtime.py

@@ -1,3 +1,4 @@
+import time
 from asyncio import Event, to_thread
 from collections import defaultdict
 from typing import Any, AsyncGenerator, Generator, Literal, Optional, TypedDict
@@ -15,21 +16,25 @@ class Packet(TypedDict):
 
 
 class RuntimeBase:
-    def __init__(self, address: str = "inproc://"):
-        self.address: str = address
+    __client_count: dict[str, int] = {}
+
+    def __init__(self, url: str = "inproc://"):
+        self.url: str = url
         self._responses: dict[str, Packet] = {}
         self._exec_responses: defaultdict[str, dict[int, Packet]] = defaultdict(dict)
         self._listen_lock: Optional[Event] = None
 
-        start_threads(self.address)
-        self._client: BrokerClient = BrokerClient(address)
+        if RuntimeBase.__client_count.get(self.url, 0) == 0:
+            start_threads(self.url)
+            RuntimeBase.__client_count[self.url] = 0
+
+        self._client: BrokerClient = BrokerClient(self.url)
         txid = self._send_type1("connect")
-        if not txid:
-            raise RuntimeError("Connection failed")
         self._sync_listen()
         if not self._responses[txid]["body"]["status"]:
             raise RuntimeError("Connection failed")
         del self._responses[txid]
+        RuntimeBase.__client_count[self.url] += 1
 
     def __del__(self):
         self.stop()
@@ -41,22 +46,32 @@ class RuntimeBase:
         self.stop()
 
     def stop(self):
-        if self._client:
+        if self.is_alive():
             txid = self._send_type1("disconnect")
-            if not txid:
-                raise RuntimeError("Disconnection failed")
             while txid not in self._responses:
                 self._sync_listen()
             if not self._responses[txid]["body"]["status"]:
                 raise RuntimeError("Disconnection failed")
             self._client = None
-            stop_threads(self.address)
+            RuntimeBase.__client_count[self.url] -= 1
+        if RuntimeBase.__client_count.get(self.url, 0) <= 0:
+            stop_threads(self.url)
+            RuntimeBase.__client_count.pop(self.url, 0)
 
-    def _send_type1(self, ptype: Literal["connect", "disconnect"]) -> Optional[str]:
+    def is_alive(self):
+        return self._client is not None
+
+    def _send_type1(self, ptype: Literal["connect", "disconnect"]) -> str:
         txid = generate_uuid()
-        if self._client.send_type1(txid, ptype):
-            return txid
-        raise RuntimeError("Failed to send packet")
+        retry_count = 0
+        # Since the broker thread might start slightly later than the runtime client,
+        # we retry sending the packat a few times to ensure delivery.
+        while retry_count < 3:
+            if self._client.send_type1(txid, ptype):
+                return txid
+            time.sleep(0.001)
+            retry_count += 1
+        raise RuntimeError(f'Failed to send packet "{ptype}"')
 
     def _send_type2(
         self,
@@ -76,7 +91,7 @@ class RuntimeBase:
         *args,
     ):
         txid = generate_uuid()
-        if self._client.send_type2(txid, ptype, status, *args):
+        if self._client.send_type3(txid, ptype, status, *args):
             return txid
         raise RuntimeError("Failed to send packet")
 
@@ -112,8 +127,8 @@ class RuntimeBase:
 
 
 class Runtime(RuntimeBase):
-    def __init__(self, address: str = "inproc://"):
-        super().__init__(address)
+    def __init__(self, url: str = "inproc://"):
+        super().__init__(url)
 
     def call(self, func_name: str, input: Any) -> Any:
         rv = [v for v in self.call_iter(func_name, input)]
@@ -193,8 +208,8 @@ class Runtime(RuntimeBase):
 
 
 class AsyncRuntime(RuntimeBase):
-    def __init__(self, address: str = "inproc://"):
-        super().__init__(address)
+    def __init__(self, url: str = "inproc://"):
+        super().__init__(url)
 
     async def call(self, func_name: str, input: Any) -> Any:
         rv = [v async for v in self.call_iter(func_name, input)]

ailoy/tools.py

@@ -0,0 +1,205 @@
+import inspect
+import json
+import re
+import types
+from typing import (
+    Any,
+    Callable,
+    Optional,
+    Union,
+    get_args,
+    get_origin,
+    get_type_hints,
+)
+
+description_re = re.compile(r"^(.*?)[\n\s]*(Args:|Returns:|Raises:|\Z)", re.DOTALL)
+# Extracts the Args: block from the docstring
+args_re = re.compile(r"\n\s*Args:\n\s*(.*?)[\n\s]*(Returns:|Raises:|\Z)", re.DOTALL)
+# Splits the Args: block into individual arguments
+args_split_re = re.compile(
+    r"""
+(?:^|\n)  # Match the start of the args block, or a newline
+\s*(\w+):\s*  # Capture the argument name and strip spacing
+(.*?)\s*  # Capture the argument description, which can span multiple lines, and strip trailing spacing
+(?=\n\s*\w+:|\Z)  # Stop when you hit the next argument or the end of the block
+""",
+    re.DOTALL | re.VERBOSE,
+)
+# Extracts the Returns: block from the docstring, if present. Note that most chat templates ignore the return type/doc!
+returns_re = re.compile(r"\n\s*Returns:\n\s*(.*?)[\n\s]*(Raises:|\Z)", re.DOTALL)
+
+
+class TypeHintParsingException(Exception):
+    """Exception raised for errors in parsing type hints to generate JSON schemas"""
+
+    pass
+
+
+class DocstringParsingException(Exception):
+    """Exception raised for errors in parsing docstrings to generate JSON schemas"""
+
+    pass
+
+
+def _get_json_schema_type(param_type: str) -> dict[str, str]:
+    type_mapping = {
+        int: {"type": "integer"},
+        float: {"type": "number"},
+        str: {"type": "string"},
+        bool: {"type": "boolean"},
+        type(None): {"type": "null"},
+        Any: {},
+    }
+    # if is_vision_available():
+    #     type_mapping[Image] = {"type": "image"}
+    # if is_torch_available():
+    #     type_mapping[Tensor] = {"type": "audio"}
+    return type_mapping.get(param_type, {"type": "object"})
+
+
+def _parse_type_hint(hint: str) -> dict:
+    origin = get_origin(hint)
+    args = get_args(hint)
+
+    if origin is None:
+        try:
+            return _get_json_schema_type(hint)
+        except KeyError:
+            raise TypeHintParsingException(
+                "Couldn't parse this type hint, likely due to a custom class or object: ", hint
+            )
+
+    elif origin is Union or (hasattr(types, "UnionType") and origin is types.UnionType):
+        # Recurse into each of the subtypes in the Union, except None, which is handled separately at the end
+        subtypes = [_parse_type_hint(t) for t in args if t is not type(None)]
+        if len(subtypes) == 1:
+            # A single non-null type can be expressed directly
+            return_dict = subtypes[0]
+        elif all(isinstance(subtype["type"], str) for subtype in subtypes):
+            # A union of basic types can be expressed as a list in the schema
+            return_dict = {"type": sorted([subtype["type"] for subtype in subtypes])}
+        else:
+            # A union of more complex types requires "anyOf"
+            return_dict = {"anyOf": subtypes}
+        if type(None) in args:
+            return_dict["nullable"] = True
+        return return_dict
+
+    elif origin is list:
+        if not args:
+            return {"type": "array"}
+        else:
+            # Lists can only have a single type argument, so recurse into it
+            return {"type": "array", "items": _parse_type_hint(args[0])}
+
+    elif origin is tuple:
+        if not args:
+            return {"type": "array"}
+        if len(args) == 1:
+            raise TypeHintParsingException(
+                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 "
+                "pass the element directly."
+            )
+        if ... in args:
+            raise TypeHintParsingException(
+                "Conversion of '...' is not supported in Tuple type hints. "
+                "Use List[] types for variable-length"
+                " inputs instead."
+            )
+        return {"type": "array", "prefixItems": [_parse_type_hint(t) for t in args]}
+
+    elif origin is dict:
+        # The JSON equivalent to a dict is 'object', which mandates that all keys are strings
+        # However, we can specify the type of the dict values with "additionalProperties"
+        out = {"type": "object"}
+        if len(args) == 2:
+            out["additionalProperties"] = _parse_type_hint(args[1])
+        return out
+
+    raise TypeHintParsingException("Couldn't parse this type hint, likely due to a custom class or object: ", hint)
+
+
+def _convert_type_hints_to_json_schema(func: Callable) -> dict:
+    type_hints = get_type_hints(func)
+    signature = inspect.signature(func)
+    required = []
+    for param_name, param in signature.parameters.items():
+        if param.annotation == inspect.Parameter.empty:
+            raise TypeHintParsingException(f"Argument {param.name} is missing a type hint in function {func.__name__}")
+        if param.default == inspect.Parameter.empty:
+            required.append(param_name)
+
+    properties = {}
+    for param_name, param_type in type_hints.items():
+        properties[param_name] = _parse_type_hint(param_type)
+
+    schema = {"type": "object", "properties": properties}
+    if required:
+        schema["required"] = required
+
+    return schema
+
+
+def parse_google_format_docstring(docstring: str) -> tuple[Optional[str], Optional[dict], Optional[str]]:
+    """
+    Parses a Google-style docstring to extract the function description,
+    argument descriptions, and return description.
+
+    Args:
+        docstring (str): The docstring to parse.
+
+    Returns:
+        The function description, arguments, and return description.
+    """
+
+    # Extract the sections
+    description_match = description_re.search(docstring)
+    args_match = args_re.search(docstring)
+    returns_match = returns_re.search(docstring)
+
+    # Clean and store the sections
+    description = description_match.group(1).strip() if description_match else None
+    docstring_args = args_match.group(1).strip() if args_match else None
+    returns = returns_match.group(1).strip() if returns_match else None
+
+    # Parsing the arguments into a dictionary
+    if docstring_args is not None:
+        docstring_args = "\n".join([line for line in docstring_args.split("\n") if line.strip()])  # Remove blank lines
+        matches = args_split_re.findall(docstring_args)
+        args_dict = {match[0]: re.sub(r"\s*\n+\s*", " ", match[1].strip()) for match in matches}
+    else:
+        args_dict = {}
+
+    return description, args_dict, returns
+
+
+def get_json_schema(func: Callable) -> dict:
+    doc = inspect.getdoc(func)
+    if not doc:
+        raise DocstringParsingException(f"Cannot generate JSON schema for {func.__name__} because it has no docstring!")
+    doc = doc.strip()
+    main_doc, param_descriptions, return_doc = parse_google_format_docstring(doc)
+
+    json_schema = _convert_type_hints_to_json_schema(func)
+    if (return_dict := json_schema["properties"].pop("return", None)) is not None:
+        if return_doc is not None:  # We allow a missing return docstring since most templates ignore it
+            return_dict["description"] = return_doc
+    for arg, schema in json_schema["properties"].items():
+        if arg not in param_descriptions:
+            raise DocstringParsingException(
+                f"Cannot generate JSON schema for {func.__name__} because the docstring has no description for the argument '{arg}'"
+            )
+        desc = param_descriptions[arg]
+        enum_choices = re.search(r"\(choices:\s*(.*?)\)\s*$", desc, flags=re.IGNORECASE)
+        if enum_choices:
+            schema["enum"] = [c.strip() for c in json.loads(enum_choices.group(1))]
+            desc = enum_choices.string[: enum_choices.start()].strip()
+        schema["description"] = desc
+
+    output = {"name": func.__name__, "description": main_doc, "parameters": json_schema}
+    if return_dict is not None:
+        output["return"] = return_dict
+    return {"type": "function", "function": output}

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/vector_store.py

@@ -46,8 +46,8 @@ class VectorStore:
         vector_store_name: Literal["faiss", "chromadb"],
         url: Optional[str] = None,
         collection: Optional[str] = None,
-        embedding_model_attrs: dict[str, Any] = dict(),
-        vector_store_attrs: dict[str, Any] = dict(),
+        embedding_model_attrs: Optional[dict[str, Any]] = None,
+        vector_store_attrs: Optional[dict[str, Any]] = None,
     ):
         """
         Creates an instance.
@@ -69,10 +69,10 @@ class VectorStore:
         self.define(
             embedding_model_name,
             vector_store_name,
-            url,
-            collection,
-            embedding_model_attrs,
-            vector_store_attrs,
+            url=url,
+            collection=collection,
+            embedding_model_attrs=embedding_model_attrs,
+            vector_store_attrs=vector_store_attrs,
         )
 
     def __del__(self):
@@ -90,8 +90,8 @@ class VectorStore:
         vector_store_name: Literal["faiss", "chromadb"],
         url: Optional[str] = None,
         collection: Optional[str] = None,
-        embedding_model_attrs: dict[str, Any] = dict(),
-        vector_store_attrs: dict[str, Any] = dict(),
+        embedding_model_attrs: Optional[dict[str, Any]] = None,
+        vector_store_attrs: Optional[dict[str, Any]] = None,
     ):
         """
         Defines the embedding model and vector store components to the runtime.
@@ -111,13 +111,14 @@ class VectorStore:
                 self._component_state.embedding_model_name,
                 {
                     "model": "BAAI/bge-m3",
-                    **embedding_model_attrs,
+                    **(embedding_model_attrs or {}),
                 },
             )
         else:
             raise NotImplementedError(f"Unsupprted embedding model: {embedding_model_name}")
 
         # Initialize vector store
+        vector_store_attrs = vector_store_attrs or {}
         if vector_store_name == "faiss":
             if "dimension" not in vector_store_attrs:
                 vector_store_attrs["dimension"] = dimension

ailoy_py-0.0.3.dist-info/DELVEWHEEL

@@ -0,0 +1,2 @@
+Version: 1.10.1
+Arguments: ['C:\\hostedtoolcache\\windows\\Python\\3.11.9\\x64\\Scripts\\delvewheel', 'repair', '-w', 'wheelhouse', 'D:\\a\\ailoy\\ailoy\\bindings\\python\\dist\\ailoy_py-0.0.3-cp311-cp311-win_amd64.whl', '--add-path', 'D:\\a\\ailoy\\ailoy\\bindings\\python\\build\\_deps\\tvm-build\\Release', '--exclude', 'vulkan-1.dll']