stores 0.1.7.dev6__py3-none-any.whl → 0.1.8.dev2__py3-none-any.whl

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/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/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/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.dist-info → stores-0.1.8.dev2.dist-info}/METADATA

@@ -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.dist-info → stores-0.1.8.dev2.dist-info}/RECORD

@@ -1,15 +1,15 @@
 stores/__init__.py,sha256=KYpKkNrMLx6ssVUbxHnn9wFBq5F5KnaFchcimIfDf9g,186
 stores/constants.py,sha256=7WqFmoGCtmUKHA5WHxOJvvK7g-yYu_KGoqnuVFADNao,57
-stores/format.py,sha256=LduYBVDiUDB1J1HDyu9jHrRG1V97pw6C5g76OirpJJY,7792
+stores/format.py,sha256=L6DaRklE0CpUzMoVTckliaP-M5UrF6gW8nEG_m4y21M,8089
 stores/parse.py,sha256=HYPNPzQod2vpu1Cln7yQ8aVkZT1Mw2IN0sZ2A1DIaqE,4967
 stores/utils.py,sha256=GPWT6lCoGobwP3PlEOHyJfKyd0dobamjyErcR7lgm7M,242
 stores/indexes/__init__.py,sha256=s-RNqml8uGREQhxwSdDoxcbcxeD8soB9BcL5dBKsQfI,215
-stores/indexes/base_index.py,sha256=YrEwETZ5eXj3rXK5qxOllRXqFifQoteYdzPAasbvEyg,15536
+stores/indexes/base_index.py,sha256=oaLN_Tk6PdoTSlRdGsd2wTvBVQ4S3ylV8rbIVSSNkEQ,15608
 stores/indexes/index.py,sha256=Cub5mtnYGipHfPR8BexJYRSKfuJmcGPp0B3ou2bGNqs,2901
 stores/indexes/local_index.py,sha256=Gg9LkEo1L0_NZZYPItsF_-1y6nFP369C3QlPvPyvPx8,3708
-stores/indexes/remote_index.py,sha256=-GV4l2c7GL6_bcVOQnSK5rzYru237bwC5L6eiO-QFzM,3438
-stores/indexes/venv_utils.py,sha256=k664ptVhPfs4AHCpB6CiE-pWHGcqi982YjsZDDt64Ys,18577
-stores-0.1.7.dev6.dist-info/METADATA,sha256=Gg4CaLZDS0isd_R0Vhpn7IGCIQQ51Og5rBuBANnV7tk,3081
-stores-0.1.7.dev6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-stores-0.1.7.dev6.dist-info/licenses/LICENSE,sha256=VTidYE7_Dam0Dwyq095EhhDIqi47g03oVpLAHQgKws0,1066
-stores-0.1.7.dev6.dist-info/RECORD,,
+stores/indexes/remote_index.py,sha256=kbyfjqeOgH5IgC0gYfoq79loFM1AeOCCYcsKimudZAg,3666
+stores/indexes/venv_utils.py,sha256=oN0LMQxeZQvVlkqsu3vrshaYaSmSDE-VOFSmMKc8QeI,18822
+stores-0.1.8.dev2.dist-info/METADATA,sha256=xsNHI0bP0mF7StNBmRniznrzUFFsiPMptKqSmO8pWko,3081
+stores-0.1.8.dev2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+stores-0.1.8.dev2.dist-info/licenses/LICENSE,sha256=VTidYE7_Dam0Dwyq095EhhDIqi47g03oVpLAHQgKws0,1066
+stores-0.1.8.dev2.dist-info/RECORD,,