stores 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

stores/indexes/index.py

@@ -1,7 +1,7 @@
 import logging
 import os
 from pathlib import Path
-from typing import Callable
+from typing import Callable, Optional
 
 from stores.indexes.base_index import BaseIndex
 from stores.indexes.local_index import LocalIndex
@@ -17,6 +17,8 @@ class Index(BaseIndex):
         self,
         tools: list[Callable, os.PathLike] | None = None,
         env_var: dict[str, dict] | None = None,
+        cache_dir: Optional[os.PathLike] = None,
+        reset_cache=False,
     ):
         self.env_var = env_var or {}
         tools = tools or []
@@ -38,7 +40,10 @@ class Index(BaseIndex):
                     # Load RemoteIndex
                     try:
                         loaded_index = RemoteIndex(
-                            index_name, env_var=self.env_var.get(index_name)
+                            index_name,
+                            env_var=self.env_var.get(index_name),
+                            cache_dir=cache_dir,
+                            reset_cache=reset_cache,
                         )
                     except Exception:
                         logger.warning(

stores/indexes/remote_index.py

@@ -1,7 +1,10 @@
 import json
 import logging
+import shutil
 import venv
+from os import PathLike
 from pathlib import Path
+from typing import Optional
 
 import requests
 from git import Repo
@@ -21,6 +24,10 @@ INDEX_LOOKUP_URL = (
 )
 
 
+def clear_default_cache():
+    shutil.rmtree(CACHE_DIR)
+
+
 def lookup_index(index_id: str, index_version: str | None = None):
     response = requests.post(
         INDEX_LOOKUP_URL,
@@ -39,11 +46,24 @@ def lookup_index(index_id: str, index_version: str | None = None):
 
 
 class RemoteIndex(BaseIndex):
-    def __init__(self, index_id: str, env_var: dict | None = None):
+    def __init__(
+        self,
+        index_id: str,
+        env_var: dict | None = None,
+        cache_dir: Optional[PathLike] = None,
+        reset_cache=False,
+    ):
         self.index_id = index_id
-        self.index_folder = CACHE_DIR / self.index_id
+        if cache_dir is None:
+            cache_dir = CACHE_DIR
+        else:
+            cache_dir = Path(cache_dir)
+        if reset_cache:
+            shutil.rmtree(cache_dir)
+        self.index_folder = cache_dir / self.index_id
         self.env_var = env_var or {}
         if not self.index_folder.exists():
+            logger.info(f"Installing {index_id}...")
             commit_like = None
             if ":" in index_id:
                 index_id, commit_like = index_id.split(":")

stores/indexes/venv_utils.py

@@ -93,6 +93,7 @@ def init_venv_tools(index_folder: os.PathLike, env_var: dict | None = None):
             tool_id=tool_id,
             index_folder=index_folder,
             venv=VENV_NAME,
+            env_var=env_var,
         )
         tool = parse_tool_signature(
             signature_dict=tool_sig,
@@ -105,9 +106,15 @@ def init_venv_tools(index_folder: os.PathLike, env_var: dict | None = None):
 
 
 # TODO: Sanitize tool_id, args, and kwargs
-def get_tool_signature(tool_id: str, index_folder: os.PathLike, venv: str = VENV_NAME):
+def get_tool_signature(
+    tool_id: str,
+    index_folder: os.PathLike,
+    venv: str = VENV_NAME,
+    env_var: dict | None = None,
+):
     module_name = ".".join(tool_id.split(".")[:-1])
     tool_name = tool_id.split(".")[-1]
+    env_var = env_var or {}
 
     runner = f"""
 import pickle, sys, traceback, inspect, enum
@@ -115,44 +122,49 @@ from typing import Any, Dict, List, Literal, Tuple, Union, get_args, get_origin,
 import types as T
 
 
-def extract_type_info(typ):
+def extract_type_info(typ, custom_types: list[str] | None = None):
+    custom_types = custom_types or []
+    if hasattr(typ, "__name__") and typ.__name__ in custom_types:
+        return typ.__name__
     origin = get_origin(typ)
     args = list(get_args(typ))
     if origin is Literal:
         return {{"type": "Literal", "values": args}}
     elif inspect.isclass(typ) and issubclass(typ, enum.Enum):
+        custom_types.append(typ.__name__)
         return {{
             "type": "Enum",
             "type_name": typ.__name__,
             "values": {{v.name: v.value for v in typ}},
         }}
     elif isinstance(typ, type) and typ.__class__.__name__ == "_TypedDictMeta":
+        custom_types.append(typ.__name__)
         hints = get_type_hints(typ)
         return {{
             "type": "TypedDict",
             "type_name": typ.__name__,
-            "fields": {{k: extract_type_info(v) for k, v in hints.items()}}
+            "fields": {{k: extract_type_info(v, custom_types) for k, v in hints.items()}}
         }}
     elif origin in (list, List) or typ is list:
         return {{
             "type": "List",
-            "item_type": extract_type_info(args[0]) if args else {{"type": Any}}
+            "item_type": extract_type_info(args[0], custom_types) if args else {{"type": Any}}
         }}
     elif origin in (dict, Dict) or typ is dict:
         return {{
             "type": "Dict",
-            "key_type": extract_type_info(args[0]) if args else {{"type": Any}},
-            "value_type": extract_type_info(args[1]) if len(args) > 1 else {{"type": Any}}
+            "key_type": extract_type_info(args[0], custom_types) if args else {{"type": Any}},
+            "value_type": extract_type_info(args[1], custom_types) if len(args) > 1 else {{"type": Any}}
         }}
     elif origin in (tuple, Tuple) or typ is tuple:
         return {{
             "type": "Tuple",
-            "item_types": [extract_type_info(arg) for arg in args] if args else [{{"type": Any}}]
+            "item_types": [extract_type_info(arg, custom_types) for arg in args] if args else [{{"type": Any}}]
         }}
     elif origin is Union or origin is T.UnionType:
         return {{
             "type": "Union",
-            "options": [extract_type_info(arg) for arg in args]
+            "options": [extract_type_info(arg, custom_types) for arg in args]
         }}
     else:
         return {{"type": typ}}
@@ -192,6 +204,7 @@ except Exception as e:
         [f"{venv}/bin/python", "-c", runner],
         capture_output=True,
         cwd=index_folder,
+        env=env_var,
     )
     try:
         response = pickle.loads(result.stdout)
@@ -205,33 +218,39 @@ except Exception as e:
         raise RuntimeError(f"Error loading tool {tool_id}:\n{response['error']}")
 
 
-def parse_param_type(param_info: dict):
+def parse_param_type(param_info: dict, custom_types: list[str] | None = None):
+    custom_types = custom_types or []
+    # Support ForwardRef
     param_type = param_info["type"]
+    if param_type in custom_types:
+        return param_type
     if not isinstance(param_type, str):
         return param_type
     if param_type == "Literal":
         return Literal.__getitem__(tuple(param_info["values"]))
     elif param_type == "Enum":
+        custom_types.append(param_info["type_name"])
         return Enum(param_info["type_name"], param_info["values"])
     elif param_type == "TypedDict":
+        custom_types.append(param_info["type_name"])
         properties = {}
         for k, v in param_info["fields"].items():
-            properties[k] = parse_param_type(v)
+            properties[k] = parse_param_type(v, custom_types)
         return TypedDict(param_info["type_name"], properties)
     elif param_type == "List":
-        return list[parse_param_type(param_info["item_type"])]
+        return list[parse_param_type(param_info["item_type"], custom_types)]
     elif param_type == "Dict":
         return Dict[
-            parse_param_type(param_info["key_type"]),
-            parse_param_type(param_info["value_type"]),
+            parse_param_type(param_info["key_type"], custom_types),
+            parse_param_type(param_info["value_type"], custom_types),
         ]
     elif param_type == "Tuple":
         return Tuple.__getitem__(
-            tuple([parse_param_type(i) for i in param_info["item_types"]])
+            tuple([parse_param_type(i, custom_types) for i in param_info["item_types"]])
         )
     elif param_type == "Union":
         return Union.__getitem__(
-            tuple([parse_param_type(i) for i in param_info["options"]])
+            tuple([parse_param_type(i, custom_types) for i in param_info["options"]])
         )
     else:
         raise TypeError(f"Invalid param type {param_type} in param info {param_info}")

{stores-0.1.1.dist-info → stores-0.1.3.dist-info}/METADATA

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

{stores-0.1.1.dist-info → stores-0.1.3.dist-info}/RECORD

@@ -5,11 +5,11 @@ 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=K4--_OAFyPNESaisD0JPNRIEffCnjw9nziahDgjdZHU,10708
-stores/indexes/index.py,sha256=ckyisaoD_0Uw3dWXpMYYhNWxMXcfykwJIY8lUdpNntA,2113
+stores/indexes/index.py,sha256=X8q2Pe6PCmzQ9j0KVvJmqZid1A1RB0jUe-JugcOrl7o,2330
 stores/indexes/local_index.py,sha256=NS4AWzeMVZ51YPo0eJ_udbF5aASuksG-88YXLDn41fQ,2880
-stores/indexes/remote_index.py,sha256=-iKO0q9MS1CiyzY1lu90FxHs0pVR6z82_7Lk7jX9wFk,2523
-stores/indexes/venv_utils.py,sha256=odeKBphn6H2Jwz6pyy4pvEZFB48Kd1rroAuMiJjcPLE,11696
-stores-0.1.1.dist-info/METADATA,sha256=FwbKEQMpPJNB_wPgKRU_kM2lhmUNocmL2VIwBP7UZOE,3076
-stores-0.1.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-stores-0.1.1.dist-info/licenses/LICENSE,sha256=VTidYE7_Dam0Dwyq095EhhDIqi47g03oVpLAHQgKws0,1066
-stores-0.1.1.dist-info/RECORD,,
+stores/indexes/remote_index.py,sha256=dp84keyEDyL-cQrmshwHd4QqDGMw0CTSJcQxqZQiULY,2983
+stores/indexes/venv_utils.py,sha256=U69UsHn9APTc615TchW2aZmrBORgbZhuZunv4zQv_hU,12514
+stores-0.1.3.dist-info/METADATA,sha256=4ta-OcEHDgCulY7lk_yxkBb_JDOxbY0keiLisgPS8_U,3076
+stores-0.1.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+stores-0.1.3.dist-info/licenses/LICENSE,sha256=VTidYE7_Dam0Dwyq095EhhDIqi47g03oVpLAHQgKws0,1066
+stores-0.1.3.dist-info/RECORD,,