stores 0.1.7.dev6__tar.gz → 0.1.8.dev2__tar.gz

{stores-0.1.7.dev6 → stores-0.1.8.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stores
-Version: 0.1.7.dev6
+Version: 0.1.8.dev2
 Summary: Repository of Python functions and tools for LLMs
 License-File: LICENSE
 Requires-Python: >=3.10

{stores-0.1.7.dev6 → stores-0.1.8.dev2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "stores"
-version = "0.1.7dev6"
+version = "0.1.8dev2"
 description = "Repository of Python functions and tools for LLMs"
 readme = "README.md"
 requires-python = ">=3.10"

{stores-0.1.7.dev6 → stores-0.1.8.dev2}/run_browser_use.py

@@ -1,16 +1,16 @@
 import stores  # noqa
 
+
 index = stores.indexes.LocalIndex(
     "/drive3/Silanthro/tools/basic-browser-use",
     create_venv=True,
-    include=["basic_browser_use.stream_browser_agent"],
+    include=["basic_browser_use.stream_browser_agent_gui"],
 )
-# print(index.tools)
 
 
 async def collect():
     async for value in index.astream_execute(
-        "basic_browser_use.stream_browser_agent",
+        "basic_browser_use.stream_browser_agent_gui",
         {"task": "What are the latest OpenAI models?"},
     ):
         print(value)
@@ -28,7 +28,7 @@ async def collect():
 # asyncio.run(collect())
 
 for value in index.stream_execute(
-    "basic_browser_use.stream_browser_agent",
+    "basic_browser_use.stream_browser_agent_gui",
     {"task": "What are the latest OpenAI models?"},
 ):
     print(value)

{stores-0.1.7.dev6 → stores-0.1.8.dev2}/stores/format.py

@@ -61,7 +61,7 @@ def get_type_repr(typ: Type | GenericAlias) -> list[str]:
         return [type_mappings[typ.__name__]]
 
 
-def get_type_schema(typ: Type | GenericAlias):
+def get_type_schema(typ: Type | GenericAlias, provider: ProviderFormat):
     origin = get_origin(typ)
     args = get_args(typ)
 
@@ -77,24 +77,27 @@ def get_type_schema(typ: Type | GenericAlias):
         schema["enum"] = [v.value for v in typ]
     elif isinstance(typ, type) and typ.__class__.__name__ == "_TypedDictMeta":
         hints = get_type_hints(typ)
-        schema["properties"] = {k: get_type_schema(v) for k, v in hints.items()}
-        schema["additionalProperties"] = False
+        schema["properties"] = {
+            k: get_type_schema(v, provider) for k, v in hints.items()
+        }
+        if provider != ProviderFormat.GOOGLE_GEMINI:
+            schema["additionalProperties"] = False
         schema["required"] = list(hints.keys())
     elif origin in (list, List) or typ is dict:
         if args:
-            schema["items"] = get_type_schema(args[0])
+            schema["items"] = get_type_schema(args[0], provider)
         else:
             raise TypeError("Insufficient argument type information")
     elif origin in (dict, Dict) or typ is dict:
         raise TypeError("Insufficient argument type information")
     elif origin in (tuple, Tuple) or typ is tuple:
         if args:
-            schema["items"] = get_type_schema(args[0])
+            schema["items"] = get_type_schema(args[0], provider)
         else:
             raise TypeError("Insufficient argument type information")
     elif origin is Union or origin is T.UnionType:
         for arg in args:
-            subschema = get_type_schema(arg)
+            subschema = get_type_schema(arg, provider)
             del subschema["type"]
             schema = {
                 **schema,
@@ -103,14 +106,17 @@ def get_type_schema(typ: Type | GenericAlias):
 
     # Un-nest single member type lists since Gemini does not accept list of types
     # Optional for OpenAI or Anthropic
-    if schema["type"] and len(schema["type"]) == 1:
-        schema["type"] = schema["type"][0]
+    if schema["type"]:
+        if len(schema["type"]) == 1:
+            schema["type"] = schema["type"][0]
+        elif len(schema["type"]) > 1 and provider == ProviderFormat.GOOGLE_GEMINI:
+            schema["type"] = schema["type"][0]
 
     return schema
 
 
 def get_param_schema(param: inspect.Parameter, provider: ProviderFormat):
-    param_schema = get_type_schema(param.annotation)
+    param_schema = get_type_schema(param.annotation, provider)
 
     if param_schema["type"] is None:
         raise TypeError(f"Unsupported type: {param.annotation.__name__}")

{stores-0.1.7.dev6 → stores-0.1.8.dev2}/stores/indexes/base_index.py

@@ -9,7 +9,6 @@ from typing import (
     Callable,
     List,
     Literal,
-    Optional,
     Tuple,
     Union,
     get_args,
@@ -98,7 +97,7 @@ def _handle_non_string_literal(annotation: type):
         return list[new_annotation], {"item": literal_map}
     if origin is Union or origin is UnionType:
         union_literal_maps = {}
-        argtype_args = [a for a in get_args(annotation) if a != NoneType]
+        argtype_args = [a for a in get_args(annotation)]
         new_union, literal_map = _handle_non_string_literal(argtype_args[0])
         union_literal_maps[new_union.__name__] = literal_map
         for child_argtype in argtype_args[1:]:
@@ -198,12 +197,12 @@ def wrap_tool(tool: Callable):
             # Process args with default values: make sure type includes None
             new_annotation = argtype
             if new_annotation is Parameter.empty:
-                new_annotation = Optional[type(new_arg.default)]
+                new_annotation = type(new_arg.default) | None
             origin = get_origin(new_annotation)
             if origin not in [Union, UnionType] or NoneType not in get_args(
                 new_annotation
             ):
-                new_annotation = Optional[new_annotation]
+                new_annotation = new_annotation | None
             new_arg = new_arg.replace(
                 default=None,
                 kind=Parameter.POSITIONAL_OR_KEYWORD,
@@ -402,13 +401,17 @@ class BaseIndex:
             # Handle sync
             yield tool_fn(**kwargs)
 
-    def parse_and_execute(self, msg: str):
+    def parse_and_execute(self, msg: str, collect_results=False):
         toolcall = llm_parse_json(msg, keys=["toolname", "kwargs"])
-        return self.execute(toolcall.get("toolname"), toolcall.get("kwargs"))
+        return self.execute(
+            toolcall.get("toolname"), toolcall.get("kwargs"), collect_results
+        )
 
-    async def async_parse_and_execute(self, msg: str):
+    async def aparse_and_execute(self, msg: str, collect_results=False):
         toolcall = llm_parse_json(msg, keys=["toolname", "kwargs"])
-        return await self.aexecute(toolcall.get("toolname"), toolcall.get("kwargs"))
+        return await self.aexecute(
+            toolcall.get("toolname"), toolcall.get("kwargs"), collect_results
+        )
 
     def stream_parse_and_execute(self, msg: str):
         toolcall = llm_parse_json(msg, keys=["toolname", "kwargs"])
@@ -416,7 +419,7 @@ class BaseIndex:
 
     async def astream_parse_and_execute(self, msg: str):
         toolcall = llm_parse_json(msg, keys=["toolname", "kwargs"])
-        async for value in self.astream_parse_and_execute(
+        async for value in self.astream_execute(
             toolcall.get("toolname"), toolcall.get("kwargs")
         ):
             yield value

{stores-0.1.7.dev6 → stores-0.1.8.dev2}/stores/indexes/remote_index.py

@@ -8,7 +8,7 @@ from pathlib import Path
 from typing import Optional
 
 import requests
-from git import Repo
+from git import GitCommandError, Repo
 
 from stores.constants import VENV_NAME
 from stores.indexes.base_index import BaseIndex
@@ -44,6 +44,8 @@ def lookup_index(index_id: str, index_version: str | None = None):
     )
     if response.ok:
         return response.json()
+    else:
+        raise ValueError(f"Index {index_id} not found in database")
 
 
 class RemoteIndex(BaseIndex):
@@ -88,7 +90,10 @@ class RemoteIndex(BaseIndex):
             if not repo_url:
                 # Otherwise, assume index references a GitHub repo
                 repo_url = f"https://github.com/{index_id}.git"
-            repo = Repo.clone_from(repo_url, self.index_folder)
+            try:
+                repo = Repo.clone_from(repo_url, self.index_folder)
+            except GitCommandError as e:
+                raise ValueError(f"Index {index_id} not found") from e
             if commit_like:
                 repo.git.checkout(commit_like)
 

{stores-0.1.7.dev6 → stores-0.1.8.dev2}/stores/indexes/venv_utils.py

@@ -327,8 +327,7 @@ def parse_tool_signature(
     if signature_dict.get("isasyncgenfunction"):
 
         async def func_handler(*args, **kwargs):
-            # TODO: Make this truly async
-            async for value in run_remote_tool(
+            async for value in run_remote_tool_async(
                 tool_id=signature_dict["tool_id"],
                 index_folder=index_folder,
                 args=args,
@@ -338,6 +337,7 @@ def parse_tool_signature(
                 stream=True,
             ):
                 yield value
+
     elif signature_dict.get("isgeneratorfunction"):
 
         def func_handler(*args, **kwargs):
@@ -347,7 +347,7 @@ def parse_tool_signature(
             def run():
                 async def runner():
                     try:
-                        async for item in run_remote_tool(
+                        async for item in run_remote_tool_async(
                             tool_id=signature_dict["tool_id"],
                             index_folder=index_folder,
                             args=args,
@@ -381,26 +381,63 @@ def parse_tool_signature(
     elif signature_dict.get("iscoroutinefunction"):
 
         async def func_handler(*args, **kwargs):
-            # TODO: Make this truly async
-            return run_remote_tool(
+            result = []
+            async for item in run_remote_tool_async(
                 tool_id=signature_dict["tool_id"],
                 index_folder=index_folder,
                 args=args,
                 kwargs=kwargs,
                 venv=venv,
                 env_var=env_var,
-            )
+                stream=True,
+            ):
+                result.append(item)
+            return result[-1] if result else None
     else:
 
-        def func_handler(*args, **kwargs):
-            return run_remote_tool(
+        async def func_handler_async_fallback(*args, **kwargs):
+            result = []
+            async for item in run_remote_tool_async(
                 tool_id=signature_dict["tool_id"],
                 index_folder=index_folder,
                 args=args,
                 kwargs=kwargs,
                 venv=venv,
                 env_var=env_var,
-            )
+                stream=True,
+            ):
+                result.append(item)
+            return result[-1] if result else None
+
+        def func_handler(*args, **kwargs):
+            coro = func_handler_async_fallback(*args, **kwargs)
+            try:
+                # Check if we're in an async context
+                asyncio.get_running_loop()
+                in_async = True
+            except RuntimeError:
+                in_async = False
+
+            if not in_async:
+                # Safe to run directly
+                return asyncio.run(coro)
+
+            q = queue.Queue()
+
+            def runner():
+                try:
+                    result = asyncio.run(coro)
+                    q.put(result)
+                except Exception as e:
+                    q.put(e)
+
+            t = threading.Thread(target=runner)
+            t.start()
+            result = q.get()
+            t.join()
+            if isinstance(result, Exception):
+                raise result
+            return result
 
     # Reconstruct signature from list of args
     params = []
@@ -426,92 +463,7 @@ def parse_tool_signature(
     return func
 
 
-# TODO: Sanitize tool_id, args, and kwargs
-def run_remote_tool(
-    tool_id: str,
-    index_folder: os.PathLike,
-    args: list | None = None,
-    kwargs: dict | None = None,
-    venv: str = VENV_NAME,
-    env_var: dict | None = None,
-    stream: bool = False,
-):
-    args = args or []
-    kwargs = kwargs or {}
-    env_var = env_var or {}
-
-    module_name = ".".join(tool_id.split(".")[:-1])
-    tool_name = tool_id.split(".")[-1]
-    payload = json.dumps(
-        {
-            "args": args,
-            "kwargs": kwargs,
-        }
-    ).encode("utf-8")
-
-    # We use sockets to pass function output
-    listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-    listener.bind(("localhost", 0))
-    listener.listen(1)
-    _, port = listener.getsockname()
-
-    result_data = {}
-
-    def handle_connection_sync():
-        conn, _ = listener.accept()
-        with conn:
-            buffer = ""
-            while True:
-                chunk = conn.recv(4096).decode("utf-8")
-                if not chunk:
-                    break
-                buffer += chunk
-                while "\n" in buffer:
-                    line, buffer = buffer.split("\n", 1)
-                    if not line.strip():
-                        continue
-                    msg = json.loads(line)
-                    if msg.get("ok") and "stream" in msg:
-                        result_data.setdefault("stream", []).append(msg["stream"])
-                    elif msg.get("ok") and "result" in msg:
-                        result_data["result"] = msg["result"]
-                    elif "error" in msg:
-                        result_data["error"] = msg["error"]
-                    elif msg.get("done"):
-                        return
-
-    async def handle_connection_async():
-        loop = asyncio.get_running_loop()
-        conn, _ = await loop.sock_accept(listener)
-        conn.setblocking(False)
-        buffer = ""
-        try:
-            while True:
-                chunk = await loop.sock_recv(conn, 4096)
-                if not chunk:
-                    break
-                buffer += chunk.decode("utf-8")
-                while "\n" in buffer:
-                    line, buffer = buffer.split("\n", 1)
-                    if not line.strip():
-                        continue
-                    msg = json.loads(line)
-                    if msg.get("ok") and "stream" in msg:
-                        yield msg["stream"]
-                    elif msg.get("ok") and "result" in msg:
-                        yield msg["result"]
-                    elif "error" in msg:
-                        raise RuntimeError(f"Subprocess error:\n{msg['error']}")
-                    elif msg.get("done"):
-                        return
-        finally:
-            conn.close()
-
-    if not stream:
-        thread = threading.Thread(target=lambda: handle_connection_sync())
-        thread.start()
-
-    runner = f"""
+tool_runner = """
 import asyncio, inspect, json, socket, sys, traceback
 sys.path.insert(0, "{index_folder}")
 
@@ -560,26 +512,98 @@ finally:
         pass
 """
 
-    proc = subprocess.Popen(
-        [get_python_command(Path(index_folder) / venv), "-c", runner],
-        stdin=subprocess.PIPE,
-        stdout=subprocess.DEVNULL,
-        stderr=subprocess.PIPE,
+
+# TODO: Sanitize tool_id, args, and kwargs
+async def run_remote_tool_async(
+    tool_id: str,
+    index_folder: os.PathLike,
+    args: list | None = None,
+    kwargs: dict | None = None,
+    venv: str = VENV_NAME,
+    env_var: dict | None = None,
+    stream: bool = True,
+):
+    args = args or []
+    kwargs = kwargs or {}
+    env_var = env_var or {}
+
+    module_name = ".".join(tool_id.split(".")[:-1])
+    tool_name = tool_id.split(".")[-1]
+    payload = json.dumps({"args": args, "kwargs": kwargs}).encode("utf-8")
+
+    listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    listener.bind(("localhost", 0))
+    listener.listen(1)
+    listener.setblocking(False)
+    _, port = listener.getsockname()
+
+    loop = asyncio.get_running_loop()
+    conn_task = loop.create_task(loop.sock_accept(listener))
+
+    runner = tool_runner.format(
+        index_folder=index_folder,
+        port=port,
+        module_name=module_name,
+        tool_name=tool_name,
+    )
+
+    proc = await asyncio.create_subprocess_exec(
+        get_python_command(Path(index_folder) / venv),
+        "-c",
+        runner,
+        stdin=asyncio.subprocess.PIPE,
+        stdout=asyncio.subprocess.DEVNULL,
+        stderr=asyncio.subprocess.PIPE,
         env=env_var or None,
     )
-    proc.stdin.write(payload)
-    proc.stdin.close()
-
-    if not stream:
-        thread.join()
-
-        if "error" in result_data:
-            raise RuntimeError(f"Subprocess failed with error:\n{result_data['error']}")
-        elif "result" in result_data:
-            return result_data["result"]
-        elif "stream" in result_data:
-            return result_data["stream"]
-        else:
-            raise RuntimeError("Subprocess completed without returning data.")
-    else:
-        return handle_connection_async()
+
+    try:
+        proc.stdin.write(payload)
+        await proc.stdin.drain()
+        proc.stdin.close()
+
+        conn, _ = await conn_task
+        conn.setblocking(False)
+
+        buffer = ""
+        result = None
+        while True:
+            chunk = await loop.sock_recv(conn, 4096)
+            if not chunk:
+                break
+            buffer += chunk.decode("utf-8")
+            while "\n" in buffer:
+                line, buffer = buffer.split("\n", 1)
+                if not line.strip():
+                    continue
+                msg = json.loads(line)
+
+                if msg.get("ok") and "stream" in msg:
+                    if stream:
+                        yield msg["stream"]
+                    else:
+                        result = msg["stream"]
+                elif msg.get("ok") and "result" in msg:
+                    result = msg["result"]
+                elif "error" in msg:
+                    raise RuntimeError(f"Subprocess error:\n{msg['error']}")
+                elif "done" in msg and result is not None:
+                    yield result
+                    return
+
+    except asyncio.CancelledError:
+        proc.kill()
+        await proc.wait()
+        raise
+    finally:
+        try:
+            conn.close()
+        except Exception:
+            pass
+        try:
+            listener.close()
+        except Exception:
+            pass
+        if proc.returncode is None:
+            proc.kill()
+            await proc.wait()

{stores-0.1.7.dev6 → stores-0.1.8.dev2}/tests/mock_index/tools.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import TypedDict
+from typing import Literal, TypedDict
 
 
 def foo(bar: str):
@@ -46,3 +46,25 @@ class Animal(TypedDict):
 
 def typed_dict_input(bar: Animal):
     return bar
+
+
+def literal_input(bar: Literal["red", "green", "blue"]):
+    return bar
+
+
+def literal_nonstring_input(bar: Literal[1, 2, 3]):
+    return bar
+
+
+def default_input(bar: str = "foo"):
+    return bar
+
+
+def stream_input(bar: str):
+    for _ in range(3):
+        yield bar
+
+
+async def astream_input(bar: str):
+    for _ in range(3):
+        yield bar

{stores-0.1.7.dev6 → stores-0.1.8.dev2}/tests/mock_index/tools.toml

@@ -6,5 +6,10 @@ tools = [
     "tools.async_foo",
     "tools.enum_input",
     "tools.typed_dict_input",
+    "tools.literal_input",
+    "tools.literal_nonstring_input",
+    "tools.default_input",
+    "tools.stream_input",
+    "tools.astream_input",
     "hello.world",
 ]

{stores-0.1.7.dev6 → stores-0.1.8.dev2}/tests/mock_index_w_deps/mock_index/__init__.py

@@ -55,3 +55,13 @@ def tuple_input(bar: tuple[Animal]) -> tuple[Animal]:
 
 def union_input(bar: Union[Color, Animal]) -> Union[Color, Animal]:
     return bar
+
+
+def stream_input(bar: str):
+    for _ in range(3):
+        yield bar
+
+
+async def astream_input(bar: str):
+    for _ in range(3):
+        yield bar

{stores-0.1.7.dev6 → stores-0.1.8.dev2}/tests/mock_index_w_deps/tools.toml

@@ -11,4 +11,7 @@ tools = [
   "mock_index.dict_input",
   "mock_index.tuple_input",
   "mock_index.union_input",
+  "mock_index.stream_input",
+  "mock_index.astream_input",
+  "mock_index.not_a_function",
 ]

{stores-0.1.7.dev6 → stores-0.1.8.dev2}/tests/test_format/conftest.py

@@ -128,6 +128,8 @@ def a_tool_with_complex_args(provider):
     elif provider == ProviderFormat.GOOGLE_GEMINI:
         for k in parameters.keys():
             parameters[k]["nullable"] = False
+            if parameters[k]["type"] == "object":
+                del parameters[k]["additionalProperties"]
         schema = {
             "name": tool.__name__,
             "parameters": {

{stores-0.1.7.dev6 → stores-0.1.8.dev2}/tests/test_indexes/conftest.py

@@ -196,15 +196,15 @@ def foo_w_optional_and_default(bar: Optional[str] = "test"):
     params=[
         {
             "function": foo_w_default,
-            "signature": "(bar: Optional[str] = None)",
+            "signature": "(bar: str | None = None)",
         },
         {
             "function": foo_w_default_notype,
-            "signature": "(bar: Optional[str] = None)",
+            "signature": "(bar: str | None = None)",
         },
         {
             "function": foo_w_optional_and_default,
-            "signature": "(bar: Optional[str] = None)",
+            "signature": "(bar: str | None = None)",
         },
     ],
 )
@@ -466,3 +466,28 @@ def cast_foo_unhandled(bar: list[str, int]):
 )
 def cast_tool(request):
     yield request.param
+
+
+def foo(bar: str):
+    return bar
+
+
+async def afoo(bar: str):
+    return bar
+
+
+def stream_foo(bar: str):
+    for _ in range(3):
+        yield bar
+
+
+async def astream_foo(bar: str):
+    for _ in range(3):
+        yield bar
+
+
+@pytest.fixture(
+    params=[foo, afoo, stream_foo, astream_foo],
+)
+def various_runtype_tool(request):
+    yield request.param

{stores-0.1.7.dev6 → stores-0.1.8.dev2}/tests/test_indexes/test_base_index.py

@@ -1,4 +1,6 @@
+import asyncio
 import inspect
+import json
 import logging
 
 import pytest
@@ -197,3 +199,162 @@ def test_base_index_format_tools(sample_tool, provider):
                 },
             }
         ]
+
+
+def handle_various_run_type(tool_fn, kwargs: dict, collect_results=False):
+    if inspect.isasyncgenfunction(tool_fn):
+        # Handle async generator
+
+        async def collect():
+            results = []
+            async for value in tool_fn(**kwargs):
+                results.append(value)
+            if collect_results:
+                return results
+            else:
+                return results[-1]
+
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        return loop.run_until_complete(collect())
+    elif inspect.isgeneratorfunction(tool_fn):
+        # Handle sync generator
+        results = []
+        for value in tool_fn(**kwargs):
+            results.append(value)
+        if collect_results:
+            return results
+        else:
+            return results[-1]
+    elif inspect.iscoroutinefunction(tool_fn):
+        # Handle async
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        return loop.run_until_complete(tool_fn(**kwargs))
+    else:
+        # Handle sync
+        return tool_fn(**kwargs)
+
+
+async def ahandle_various_run_type(tool_fn, kwargs: dict, collect_results=False):
+    kwargs = kwargs or {}
+    if inspect.isasyncgenfunction(tool_fn):
+        # Handle async generator
+        results = []
+        async for value in tool_fn(**kwargs):
+            results.append(value)
+        if collect_results:
+            return results
+        else:
+            return results[-1]
+    elif inspect.isgeneratorfunction(tool_fn):
+        # Handle sync generator
+        results = []
+        for value in tool_fn(**kwargs):
+            results.append(value)
+        if collect_results:
+            return results
+        else:
+            return results[-1]
+    elif inspect.iscoroutinefunction(tool_fn):
+        # Handle async
+        return await tool_fn(**kwargs)
+    else:
+        # Handle sync
+        return tool_fn(**kwargs)
+
+
+def test_base_index_execute(various_runtype_tool):
+    index = BaseIndex([various_runtype_tool])
+
+    assert index.parse_and_execute(
+        json.dumps(
+            {
+                "toolname": various_runtype_tool.__name__,
+                "kwargs": {"bar": "hello"},
+            }
+        )
+    ) == handle_various_run_type(various_runtype_tool, {"bar": "hello"})
+
+
+async def test_base_index_aexecute(various_runtype_tool):
+    index = BaseIndex([various_runtype_tool])
+
+    assert await index.aparse_and_execute(
+        json.dumps(
+            {
+                "toolname": various_runtype_tool.__name__,
+                "kwargs": {"bar": "hello"},
+            }
+        )
+    ) == await ahandle_various_run_type(various_runtype_tool, {"bar": "hello"})
+
+
+def test_base_index_execute_collect(various_runtype_tool):
+    index = BaseIndex([various_runtype_tool])
+
+    assert index.execute(
+        various_runtype_tool.__name__, {"bar": "hello"}, collect_results=True
+    ) == handle_various_run_type(
+        various_runtype_tool, {"bar": "hello"}, collect_results=True
+    )
+
+
+async def test_base_index_aexecute_collect(various_runtype_tool):
+    index = BaseIndex([various_runtype_tool])
+
+    assert await index.aexecute(
+        various_runtype_tool.__name__, {"bar": "hello"}, collect_results=True
+    ) == await ahandle_various_run_type(
+        various_runtype_tool, {"bar": "hello"}, collect_results=True
+    )
+
+
+def test_base_index_stream_execute(various_runtype_tool):
+    index = BaseIndex([various_runtype_tool])
+
+    results = []
+    for result in index.stream_parse_and_execute(
+        json.dumps(
+            {
+                "toolname": various_runtype_tool.__name__,
+                "kwargs": {"bar": "hello"},
+            }
+        )
+    ):
+        results.append(result)
+    if not inspect.isgeneratorfunction(
+        various_runtype_tool
+    ) and not inspect.isasyncgenfunction(various_runtype_tool):
+        results = results[0]
+
+    answer = handle_various_run_type(
+        various_runtype_tool, {"bar": "hello"}, collect_results=True
+    )
+
+    assert results == answer
+
+
+async def test_base_index_astream_execute(various_runtype_tool):
+    index = BaseIndex([various_runtype_tool])
+
+    results = []
+    async for result in index.astream_parse_and_execute(
+        json.dumps(
+            {
+                "toolname": various_runtype_tool.__name__,
+                "kwargs": {"bar": "hello"},
+            }
+        )
+    ):
+        results.append(result)
+    if not inspect.isgeneratorfunction(
+        various_runtype_tool
+    ) and not inspect.isasyncgenfunction(various_runtype_tool):
+        results = results[0]
+
+    answer = await ahandle_various_run_type(
+        various_runtype_tool, {"bar": "hello"}, collect_results=True
+    )
+
+    assert results == answer

{stores-0.1.7.dev6 → stores-0.1.8.dev2}/tests/test_indexes/test_index.py

@@ -31,6 +31,11 @@ async def test_index(local_index_folder):
         "tools.async_foo",
         "tools.enum_input",
         "tools.typed_dict_input",
+        "tools.literal_input",
+        "tools.literal_nonstring_input",
+        "tools.default_input",
+        "tools.stream_input",
+        "tools.astream_input",
         "hello.world",
         "foo",
     ]

{stores-0.1.7.dev6 → stores-0.1.8.dev2}/tests/test_indexes/test_local_index.py

@@ -13,6 +13,11 @@ def test_local_index_basic(local_index_folder):
         "tools.async_foo",
         "tools.enum_input",
         "tools.typed_dict_input",
+        "tools.literal_input",
+        "tools.literal_nonstring_input",
+        "tools.default_input",
+        "tools.stream_input",
+        "tools.astream_input",
         "hello.world",
     ]
     # Test tool execution
@@ -22,6 +27,11 @@ def test_local_index_basic(local_index_folder):
         "tools.async_foo": "hello world",
         "tools.enum_input": "red",
         "tools.typed_dict_input": {"name": "Tiger", "num_legs": 4},
+        "tools.literal_input": "red",
+        "tools.literal_nonstring_input": 1,
+        "tools.default_input": "foo",
+        "tools.stream_input": "hello world",
+        "tools.astream_input": "hello world",
         "hello.world": "hello world",
     }
     for tool in index.tools[:-1]:
@@ -53,9 +63,11 @@ def test_local_index_invalid_folder(buggy_index_folder):
 def test_local_index_with_deps(remote_index_folder):
     # Regular LocalIndex should not load index folder with deps since deps are not installed
     with pytest.raises(ModuleNotFoundError):
-        LocalIndex(remote_index_folder)
+        LocalIndex(remote_index_folder, exclude=["mock_index.not_a_function"])
     # LocalIndex can load index folder with deps if create_venv=True
-    LocalIndex(remote_index_folder, create_venv=True)
+    LocalIndex(
+        remote_index_folder, exclude=["mock_index.not_a_function"], create_venv=True
+    )
 
 
 def test_local_index_with_env_var():

{stores-0.1.7.dev6 → stores-0.1.8.dev2}/tests/test_indexes/test_remote_index.py

@@ -21,6 +21,12 @@ def test_lookup_index():
     }
 
 
+async def test_fake_remote_index():
+    fake_index = "no_such_index"
+    with pytest.raises(ValueError, match=f"Index {fake_index} not found"):
+        stores.indexes.RemoteIndex(fake_index)
+
+
 async def test_remote_index():
     # Check that env_vars are set correctly
     index = stores.indexes.RemoteIndex(

{stores-0.1.7.dev6 → stores-0.1.8.dev2}/tests/test_indexes/test_venv_utils.py

@@ -1,6 +1,7 @@
 import inspect
 import logging
 import venv
+from typing import get_args, get_origin, get_type_hints
 
 import pytest
 
@@ -48,7 +49,11 @@ async def test_install_venv_deps(remote_index_folder):
     assert reinstall_result == "Already installed"
 
     # Tools should load and be runnable
-    tools = venv_utils.init_venv_tools(remote_index_folder)
+    tools = venv_utils.init_venv_tools(
+        remote_index_folder,
+        # Test exclude
+        exclude=["mock_index.not_a_function"],
+    )
 
     # Tools should run successfully
     sample_string = "red"
@@ -73,15 +78,55 @@ async def test_install_venv_deps(remote_index_folder):
             kwargs = {"bar": [sample_animal, sample_animal]}
         elif tool.__name__ == "mock_index.union_input":
             kwargs = {"bar": sample_animal}
+        elif tool.__name__ in ["mock_index.stream_input", "mock_index.astream_input"]:
+            kwargs = {"bar": sample_string}
         else:
             kwargs = {}
         output = kwargs.get("bar", "pip_install_test")
-        if inspect.iscoroutinefunction(tool):
+        if inspect.isasyncgenfunction(tool):
+            async for value in tool(**kwargs):
+                assert value == output
+        elif inspect.isgeneratorfunction(tool):
+            for value in tool(**kwargs):
+                assert value == output
+        elif inspect.iscoroutinefunction(tool):
             assert await tool(**kwargs) == output
         else:
             assert tool(**kwargs) == output
 
 
+async def test_install_venv_deps_with_include(remote_index_folder):
+    # Create venv
+    venv_folder = remote_index_folder / VENV_NAME
+    venv.create(venv_folder, symlinks=True, with_pip=True)
+    # Test installation
+    result = venv_utils.install_venv_deps(remote_index_folder)
+    assert (remote_index_folder / venv_utils.HASH_FILE).exists()
+
+    for config_file in venv_utils.SUPPORTED_CONFIGS:
+        if (remote_index_folder / config_file).exists():
+            assert result.endswith(
+                f'"{" ".join(venv_utils.get_pip_command(venv_folder, config_file))}"'
+            )
+            assert venv_utils.has_installed(remote_index_folder / config_file)
+            break
+
+    # Running install again should show "Already installed"
+    reinstall_result = venv_utils.install_venv_deps(remote_index_folder)
+    assert reinstall_result == "Already installed"
+
+    # Tools should load and be runnable
+    tools = venv_utils.init_venv_tools(
+        remote_index_folder,
+        # Test include
+        include=["mock_index.typed_function"],
+    )
+
+    assert len(tools) == 1
+    tool = tools[0]
+    assert tool.__name__ == "mock_index.typed_function"
+
+
 def test_index_with_invalid_tool(index_folder_custom_class):
     with pytest.raises(RuntimeError, match="Error loading tool"):
         venv_utils.init_venv_tools(index_folder_custom_class)
@@ -89,10 +134,41 @@ def test_index_with_invalid_tool(index_folder_custom_class):
 
 def test_index_with_tool_error(index_folder_function_error):
     tools = venv_utils.init_venv_tools(index_folder_function_error)
-    with pytest.raises(RuntimeError, match="Subprocess failed with error"):
+    with pytest.raises(RuntimeError, match="ZeroDivisionError"):
         tools[0]()
 
 
+def test_parse_param_type_with_forward_ref():
+    param_info = {
+        "type": "TypedDict",
+        "type_name": "Person",
+        "fields": {
+            "name": {"type": str},
+            "friends": {
+                "type": "List",
+                "item_type": {
+                    "type": "Person",
+                },
+            },
+        },
+    }
+    typ = venv_utils.parse_param_type(param_info)
+    assert typ.__class__.__name__ == "_TypedDictMeta"
+    assert typ.__name__ == "Person"
+    hints = get_type_hints(typ)
+    for k, v in hints.items():
+        if k == "name":
+            assert v is str
+        elif k == "friends":
+            origin = get_origin(v)
+            assert origin is list
+            args = list(get_args(v))
+            assert len(args) == 1
+            assert args[0] == "Person"
+        else:
+            raise AssertionError("Invalid attribute")
+
+
 def test_parse_param_type_with_invalid_type():
     with pytest.raises(TypeError, match="Invalid param type"):
         venv_utils.parse_param_type({"type": "not a type"})

{stores-0.1.7.dev6 → stores-0.1.8.dev2}/uv.lock

@@ -2963,7 +2963,7 @@ wheels = [
 
 [[package]]
 name = "stores"
-version = "0.1.7.dev5"
+version = "0.1.8.dev1"
 source = { editable = "." }
 dependencies = [
     { name = "dirtyjson" },