ailoy-py 0.0.1__cp312-cp312-macosx_14_0_arm64.whl

ailoy/runtime.py

@@ -0,0 +1,273 @@
+from asyncio import Event, to_thread
+from collections import defaultdict
+from typing import Any, AsyncGenerator, Generator, Literal, Optional, TypedDict
+
+from .ailoy_py import BrokerClient, generate_uuid, start_threads, stop_threads
+
+__all__ = ["Runtime", "AsyncRuntime"]
+
+
+class Packet(TypedDict):
+    packet_type: Literal["respond", "respond_execute"]
+    instruction_type: Optional[Literal["call_function", "define_component", "delete_component", "call_method"]]
+    headers: list[bool | int | str]
+    body: dict[str, Any]
+
+
+class RuntimeBase:
+    def __init__(self, address: str = "inproc://"):
+        self.address: str = address
+        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)
+        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]
+
+    def __del__(self):
+        self.stop()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, traceback):
+        self.stop()
+
+    def stop(self):
+        if self._client:
+            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)
+
+    def _send_type1(self, ptype: Literal["connect", "disconnect"]) -> Optional[str]:
+        txid = generate_uuid()
+        if self._client.send_type1(txid, ptype):
+            return txid
+        raise RuntimeError("Failed to send packet")
+
+    def _send_type2(
+        self,
+        ptype: Literal["subscribe", "unsubscribe", "execute"],
+        itype: Literal["call_function", "define_component", "delete_component", "call_method"],
+        *args,
+    ):
+        txid = generate_uuid()
+        if self._client.send_type2(txid, ptype, itype, *args):
+            return txid
+        raise RuntimeError("Failed to send packet")
+
+    def _send_type3(
+        self,
+        ptype: Literal["respond", "respond_execute"],
+        status: bool,
+        *args,
+    ):
+        txid = generate_uuid()
+        if self._client.send_type2(txid, ptype, status, *args):
+            return txid
+        raise RuntimeError("Failed to send packet")
+
+    def _sync_listen(self) -> None:
+        packet = self._client.listen()
+        if packet is not None:
+            txid = packet["headers"][0]
+            if packet["packet_type"] == "respond_execute":
+                idx = packet["headers"][1]
+                self._exec_responses[txid][idx] = packet
+            else:
+                self._responses[txid] = packet
+
+    async def _listen(self) -> None:
+        # If listen lock exists -> wait
+        if self._listen_lock:
+            await self._listen_lock.wait()
+        else:
+            # Create a new event
+            self._listen_lock = Event()
+            # Listen packet
+            packet = await to_thread(self._client.listen)
+            if packet is not None:
+                txid = packet["headers"][0]
+                if packet["packet_type"] == "respond_execute":
+                    idx = packet["headers"][1]
+                    self._exec_responses[txid][idx] = packet
+                else:
+                    self._responses[txid] = packet
+            # Emit event
+            self._listen_lock.set()
+            self._listen_lock = None
+
+
+class Runtime(RuntimeBase):
+    def __init__(self, address: str = "inproc://"):
+        super().__init__(address)
+
+    def call(self, func_name: str, input: Any) -> Any:
+        rv = [v for v in self.call_iter(func_name, input)]
+        if len(rv) == 0:
+            return None
+        elif len(rv) == 1:
+            return rv[0]
+        else:
+            return rv
+
+    def call_iter(self, func_name: str, input: Any) -> Generator[Any, None, None]:
+        txid = self._send_type2("execute", "call_function", func_name, input)
+
+        def generator():
+            idx = 0
+            finished = False
+            while not finished:
+                while idx not in self._exec_responses[txid]:
+                    self._sync_listen()
+                packet = self._exec_responses[txid].pop(idx)
+                if not packet["body"]["status"]:
+                    raise RuntimeError(packet["body"]["reason"])
+                if packet["headers"][2]:
+                    finished = True
+                yield packet["body"]["out"]
+                idx += 1
+            del self._exec_responses[txid]
+
+        return generator()
+
+    def define(self, comp_type: str, comp_name: str, input: Any) -> None:
+        txid = self._send_type2("execute", "define_component", comp_type, comp_name, input)
+        while 0 not in self._exec_responses[txid]:
+            self._sync_listen()
+        packet = self._exec_responses[txid][0]
+        if not packet["body"]["status"]:
+            raise RuntimeError(packet["body"]["reason"])
+        del self._exec_responses[txid]
+
+    def delete(self, comp_name: str) -> None:
+        txid = self._send_type2("execute", "delete_component", comp_name)
+        while 0 not in self._exec_responses[txid]:
+            self._sync_listen()
+        packet = self._exec_responses[txid][0]
+        if not packet["body"]["status"]:
+            raise RuntimeError(packet["body"]["reason"])
+        del self._exec_responses[txid]
+
+    def call_method(self, comp_name: str, func_name: str, input: Any) -> Any:
+        rv = [v for v in self.call_iter_method(comp_name, func_name, input)]
+        if len(rv) == 0:
+            return None
+        elif len(rv) == 1:
+            return rv[0]
+        else:
+            return rv
+
+    def call_iter_method(self, comp_name: str, func_name: str, input: Any) -> Generator[Any, None, None]:
+        txid = self._send_type2("execute", "call_method", comp_name, func_name, input)
+
+        def generator():
+            idx = 0
+            finished = False
+            while not finished:
+                while idx not in self._exec_responses[txid]:
+                    self._sync_listen()
+                packet = self._exec_responses[txid].pop(idx)
+                if not packet["body"]["status"]:
+                    raise RuntimeError(packet["body"]["reason"])
+                if packet["headers"][2]:
+                    finished = True
+                yield packet["body"]["out"]
+                idx += 1
+            del self._exec_responses[txid]
+
+        return generator()
+
+
+class AsyncRuntime(RuntimeBase):
+    def __init__(self, address: str = "inproc://"):
+        super().__init__(address)
+
+    async def call(self, func_name: str, input: Any) -> Any:
+        rv = [v async for v in self.call_iter(func_name, input)]
+        if len(rv) == 0:
+            return None
+        elif len(rv) == 1:
+            return rv[0]
+        else:
+            return rv
+
+    def call_iter(self, func_name: str, input: Any) -> AsyncGenerator[Any, None]:
+        txid = self._send_type2("execute", "call_function", func_name, input)
+
+        async def generator():
+            idx = 0
+            finished = False
+            while not finished:
+                while idx not in self._exec_responses[txid]:
+                    await self._listen()
+                packet = self._exec_responses[txid].pop(idx)
+                if not packet["body"]["status"]:
+                    raise RuntimeError(packet["body"]["reason"])
+                if packet["headers"][2]:
+                    finished = True
+                yield packet["body"]["out"]
+                idx += 1
+            del self._exec_responses[txid]
+
+        return generator()
+
+    async def define(self, comp_type: str, comp_name: str, input: Any) -> None:
+        txid = self._send_type2("execute", "define_component", comp_type, comp_name, input)
+        while 0 not in self._exec_responses[txid]:
+            await self._listen()
+        packet = self._exec_responses[txid][0]
+        if not packet["body"]["status"]:
+            raise RuntimeError(packet["body"]["reason"])
+        del self._exec_responses[txid]
+
+    async def delete(self, comp_name: str) -> None:
+        txid = self._send_type2("execute", "delete_component", comp_name)
+        while 0 not in self._exec_responses[txid]:
+            await self._listen()
+        packet = self._exec_responses[txid][0]
+        if not packet["body"]["status"]:
+            raise RuntimeError(packet["body"]["reason"])
+        del self._exec_responses[txid]
+
+    async def call_method(self, comp_name: str, func_name: str, input: Any) -> Any:
+        rv = [v async for v in self.call_iter_method(comp_name, func_name, input)]
+        if len(rv) == 0:
+            return None
+        elif len(rv) == 1:
+            return rv[0]
+        else:
+            return rv
+
+    def call_iter_method(self, comp_name: str, func_name: str, input: Any) -> AsyncGenerator[Any, None]:
+        txid = self._send_type2("execute", "call_method", comp_name, func_name, input)
+
+        async def generator():
+            idx = 0
+            finished = False
+            while not finished:
+                while idx not in self._exec_responses[txid]:
+                    await self._listen()
+                packet = self._exec_responses[txid].pop(idx)
+                if not packet["body"]["status"]:
+                    raise RuntimeError(packet["body"]["reason"])
+                if packet["headers"][2]:
+                    finished = True
+                yield packet["body"]["out"]
+                idx += 1
+            del self._exec_responses[txid]
+
+        return generator()

ailoy/vector_store.py

@@ -0,0 +1,207 @@
+from typing import Any, Dict, List, Literal, Optional
+
+from pydantic import BaseModel, TypeAdapter
+
+from ailoy.runtime import Runtime, generate_uuid
+
+__all__ = ["VectorStore"]
+
+
+class VectorStoreInsertItem(BaseModel):
+    document: str
+    metadata: Optional[Dict[str, Any]] = None
+
+
+class VectorStoreRetrieveItem(VectorStoreInsertItem):
+    id: str
+    similarity: float
+
+
+class ComponentState(BaseModel):
+    embedding_model_name: str
+    vector_store_name: str
+    valid: bool
+
+
+class VectorStore:
+    """
+    The `VectorStore` class provides a high-level abstraction for storing and retrieving documents.
+    It mainly consists of two modules - embedding model and vector store.
+
+    It supports embedding text using AI and interfacing with pluggable vector store backends such as FAISS or ChromaDB.
+    This class handles initialization, insertion, similarity-based retrieval, and cleanup.
+
+    Typical usage involves:
+    1. Initializing the store via `initialize()`
+    2. Inserting documents via `insert()`
+    3. Querying similar documents via `retrieve()`
+
+    The embedding model and vector store are defined dynamically within the provided runtime.
+    """
+
+    def __init__(
+        self,
+        runtime: Runtime,
+        embedding_model_name: Literal["BAAI/bge-m3"],
+        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(),
+    ):
+        """
+        Creates an instance.
+
+        :param runtime: The runtime environment used to manage components.
+        :param embedding_model_name: The name of the embedding model to use. Currently it only supports `BAAI/bge-m3`.
+        :param url: (ChromaDB only) URL of the database.
+        :param collection: (ChromaDB only) The collection name of the database.
+        :param vector_store_name: The name of the vector store provider (One of faiss or chromaDB).
+        :param embedding_model_attrs: Additional initialization parameters (for `define_component` runtime call).
+        :param vector_store_attrs: Additional initialization parameters (for `define_component` runtime call).
+        """
+        self._runtime = runtime
+        self._component_state = ComponentState(
+            embedding_model_name=generate_uuid(),
+            vector_store_name=generate_uuid(),
+            valid=False,
+        )
+        self.define(
+            embedding_model_name,
+            vector_store_name,
+            url,
+            collection,
+            embedding_model_attrs,
+            vector_store_attrs,
+        )
+
+    def __del__(self):
+        self.delete()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, traceback):
+        self.delete()
+
+    def define(
+        self,
+        embedding_model_name: Literal["BAAI/bge-m3"],
+        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(),
+    ):
+        """
+        Defines the embedding model and vector store components to the runtime.
+        This must be called before using any other method in the class. If already defined, this is a no-op.
+        """
+        if self._component_state.valid:
+            return
+
+        # Dimension size may be different based on embedding model
+        dimension = 1024
+
+        # Initialize embedding model
+        if embedding_model_name == "BAAI/bge-m3":
+            dimension = 1024
+            self._runtime.define(
+                "tvm_embedding_model",
+                self._component_state.embedding_model_name,
+                {
+                    "model": "BAAI/bge-m3",
+                    **embedding_model_attrs,
+                },
+            )
+        else:
+            raise NotImplementedError(f"Unsupprted embedding model: {embedding_model_name}")
+
+        # Initialize vector store
+        if vector_store_name == "faiss":
+            if "dimension" not in vector_store_attrs:
+                vector_store_attrs["dimension"] = dimension
+            self._runtime.define("faiss_vector_store", self._component_state.vector_store_name, vector_store_attrs)
+        elif vector_store_name == "chromadb":
+            if "url" not in vector_store_attrs:
+                vector_store_attrs["url"] = url
+            if "collection" not in vector_store_attrs:
+                vector_store_attrs["collection"] = collection
+            self._runtime.define("chromadb_vector_store", self._component_state.vector_store_name, vector_store_attrs)
+        else:
+            raise NotImplementedError(f"Unsupprted vector store: {vector_store_name}")
+
+        self._component_state.valid = True
+
+    def delete(self):
+        """
+        Delete resources from the runtime.
+        This should be called when the VectorStore is no longer needed. If already deleted, this is a no-op.
+        """
+        if not self._component_state.valid:
+            return
+
+        self._runtime.delete(self._component_state.embedding_model_name)
+        self._runtime.delete(self._component_state.vector_store_name)
+        self._component_state.valid = False
+
+    # TODO: add NDArray typing
+    def embedding(self, text: str) -> Any:
+        """
+        Generates an embedding vector for the given input text using the embedding model.
+
+        :param text: Input text to embed.
+        :returns: The resulting embedding vector.
+        """
+        resp = self._runtime.call_method(
+            self._component_state.embedding_model_name,
+            "infer",
+            {
+                "prompt": text,
+            },
+        )
+        return resp["embedding"]
+
+    def insert(self, document: str, metadata: Optional[Dict[str, Any]] = None):
+        """
+        Inserts a new document into the vector store.
+
+        :param document: The raw text document to insert.
+        :param metadata: Metadata records additional information about the document.
+        """
+        embedding = self.embedding(document)
+        self._runtime.call_method(
+            self._component_state.vector_store_name,
+            "insert",
+            {
+                "embedding": embedding,
+                "document": document,
+                "metadata": metadata,
+            },
+        )
+
+    def retrieve(self, query: str, top_k: int = 5) -> List[VectorStoreRetrieveItem]:
+        """
+        Retrieves the top-K most similar documents to the given query.
+
+        :param query: The input query string to search for similar content.
+        :param top_k: Number of top similar documents to retrieve.
+        :returns: A list of retrieved items.
+        """
+        embedding = self.embedding(query)
+        resp = self._runtime.call_method(
+            self._component_state.vector_store_name,
+            "retrieve",
+            {
+                "query_embedding": embedding,
+                "top_k": top_k,
+            },
+        )
+        results = TypeAdapter(List[VectorStoreRetrieveItem]).validate_python(resp["results"])
+        return results
+
+    def clean(self):
+        """
+        Removes all entries from the vector store.
+        """
+        self._runtime.call_method(self._component_state.vector_store_name, "clean", {})

ailoy_py-0.0.1.dist-info/METADATA

@@ -0,0 +1,88 @@
+Metadata-Version: 2.4
+Name: ailoy-py
+Version: 0.0.1
+Summary: Python binding for Ailoy runtime APIs
+Author-Email: "Brekkylab Inc." <contact@brekkylab.com>
+License-Expression: Apache-2.0
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Python: >=3.10
+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: pydantic>=2.11.4
+Requires-Dist: rich>=14.0.0
+Requires-Dist: typer>=0.15.4
+Description-Content-Type: text/markdown
+
+# ailoy-py
+
+Python binding for Ailoy runtime APIs.
+
+See our [documentation](https://brekkylab.github.io/ailoy) for more details.
+
+## Install
+
+```bash
+pip install ailoy-py
+```
+
+## Quickstart
+
+```python
+from ailoy import Runtime, Agent
+
+# 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:
+    # This is where the actual LLM call happens
+    for resp in agent.query("Please give me a short poem about AI"):
+        agent.print(resp)
+
+# Stop the runtime
+rt.stop()
+```
+
+## Building from source
+
+### Prerequisites
+
+- Python 3.10 or higher
+- C/C++ compiler
+  (recommended versions are below)
+  - GCC >= 13
+  - LLVM Clang >= 17
+  - Apple Clang >= 15
+  - MSVC >= 19.29
+- CMake >= 3.24.0
+- Git
+- OpenSSL
+- Rust & Cargo >= 1.82.0
+- OpenMP
+- BLAS
+- LAPACK
+- Vulkan SDK (if you are using vulkan)
+
+
+### Setup development environment
+
+```bash
+pip install -e .
+```
+
+### Generate wheel
+
+```bash
+python -m build -w
+```

ailoy_py-0.0.1.dist-info/RECORD

@@ -0,0 +1,18 @@
+ailoy/ailoy_py.cpython-312-darwin.so,sha256=O174nGPlqj7xn4WdBaPa87C7TWtZN9NA_PrTXz13omo,17025088
+ailoy/ailoy_py.pyi,sha256=Yf90FEXkslpCpr1r2eqQ3-_1jLo65zmG94bBXDRqinU,991
+ailoy/vector_store.py,sha256=Ojhr4bcSfKKuMaldfsjz_G41AGgz4vyvpkWn5WMFE2c,7365
+ailoy/__init__.py,sha256=ArHu4OLbxU_re9cX-7zhshRqkn46gBcl-0SxDa5Wk00,148
+ailoy/runtime.py,sha256=tVdUaqqx9NB-h4grRkW_R2XYW5ihFn95aUhnTZkl_Zg,9997
+ailoy/agent.py,sha256=orrfgeH6Xo9nPkHqJJ1YqS_8-5-0_tkpKlF8YcdRIqU,24998
+ailoy/cli/model.py,sha256=cerCHE-VY9TOwqRcLBtmqnV-5vphpvyhtrfPFZiTKCM,2979
+ailoy/cli/__main__.py,sha256=HnBVb2em1F2NLPeNX5r3xRndRrnGaXVCduo8WBULAI0,179
+ailoy/.dylibs/libomp.dylib,sha256=GX9Rol1LBJAoha7wZJMQxuCUi_o41Yo5QLBTMN85rRI,735616
+ailoy/.dylibs/libtvm_runtime.dylib,sha256=XRFCL65dNaE9JNMvuLvfBoriquAs90OfMl_AmNd_BG8,3324976
+ailoy/presets/tools/tmdb.json,sha256=UGLN5uAJ2b-Hu3nLcW95WXDLB3mfC3rBYfQANp_e8Ps,7046
+ailoy/presets/tools/calculator.json,sha256=ePnZsjZChnvS08s9eVdIp4Bys_PlJBXPHCCjv6oMvzA,1040
+ailoy/presets/tools/nytimes.json,sha256=wrfe9bnAlSPzHladoGEX2oCAeE0wed3BvgXQ_Z2PdXg,918
+ailoy/presets/tools/frankfurter.json,sha256=bZ5vhszf_aR-B_QN4L2xrI5nR-f4AMZk41UUDq1dTXg,1152
+ailoy_py-0.0.1.dist-info/RECORD,,
+ailoy_py-0.0.1.dist-info/WHEEL,sha256=NdzsWaD5A_LKBhqfKDFIgjB3iRnxft5CRtv8RBeg8bg,141
+ailoy_py-0.0.1.dist-info/entry_points.txt,sha256=gVG45uDE6kef0wm6SEMYSgZgRNNRhSAeP2n2lPR00dI,50
+ailoy_py-0.0.1.dist-info/METADATA,sha256=FAyzKTteIMgZse5H99IVQ5q2ScjU9hlTuOggBNrgTF0,2010

ailoy_py-0.0.1.dist-info/WHEEL

@@ -0,0 +1,6 @@
+Wheel-Version: 1.0
+Generator: scikit-build-core 0.11.3
+Root-Is-Purelib: false
+Tag: cp312-cp312-macosx_14_0_arm64
+Generator: delocate 0.13.0
+

ailoy_py-0.0.1.dist-info/entry_points.txt

@@ -0,0 +1,3 @@
+[console_scripts]
+ailoy = ailoy.cli.__main__:app
+