stores 0.0.0__py3-none-any.whl → 0.1.0__py3-none-any.whl

stores/__init__.py

@@ -0,0 +1,9 @@
+from stores.format import ProviderFormat
+from stores.indexes import Index
+from stores.parse import llm_parse_json
+
+__all__ = [
+    "Index",
+    "ProviderFormat",
+    "llm_parse_json",
+]

stores/constants.py

@@ -0,0 +1,2 @@
+VENV_NAME = ".venv"
+TOOLS_CONFIG_FILENAME = "tools.toml"

stores/format.py

@@ -0,0 +1,214 @@
+import inspect
+import logging
+import types as T
+from enum import Enum
+from itertools import chain
+from typing import (
+    Callable,
+    Dict,
+    GenericAlias,
+    List,
+    Literal,
+    Tuple,
+    Type,
+    Union,
+    get_args,
+    get_origin,
+    get_type_hints,
+)
+
+from stores.utils import check_duplicates
+
+logging.basicConfig()
+logger = logging.getLogger("stores.format")
+logger.setLevel(logging.INFO)
+
+
+class ProviderFormat(str, Enum):
+    ANTHROPIC = "anthropic"
+    GOOGLE_GEMINI = "google-gemini"
+    OPENAI_CHAT = "openai-chat-completions"
+    OPENAI_RESPONSES = "openai-responses"
+
+
+def get_type_repr(typ: Type | GenericAlias) -> list[str]:
+    origin = get_origin(typ)
+    args = get_args(typ)
+
+    if origin is Literal:
+        return list(dict.fromkeys(chain(*[get_type_repr(type(arg)) for arg in args])))
+    if inspect.isclass(typ) and issubclass(typ, Enum):
+        return list(dict.fromkeys(chain(*[get_type_repr(type(v.value)) for v in typ])))
+    if isinstance(typ, type) and typ.__class__.__name__ == "_TypedDictMeta":
+        return ["object"]
+    if origin in (list, List) or typ is list:
+        return ["array"]
+    if origin in (dict, Dict) or typ is dict:
+        return ["object"]
+    if origin in (tuple, Tuple) or typ is tuple:
+        return ["array"]
+    if origin is Union or origin is T.UnionType:
+        return list(dict.fromkeys(chain(*[get_type_repr(arg) for arg in args])))
+
+    type_mappings = {
+        "str": "string",
+        "int": "integer",
+        "bool": "boolean",
+        "float": "number",
+        "NoneType": "null",
+    }
+    if typ.__name__ in type_mappings:
+        return [type_mappings[typ.__name__]]
+
+
+def get_type_schema(typ: Type | GenericAlias):
+    origin = get_origin(typ)
+    args = get_args(typ)
+
+    schema = {
+        "type": get_type_repr(typ),
+        # TODO: Retrieve description from Annotation if available
+        "description": "",
+    }
+
+    if origin is Literal:
+        schema["enum"] = list(args)
+    elif inspect.isclass(typ) and issubclass(typ, Enum):
+        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["required"] = list(hints.keys())
+    elif origin in (list, List) or typ is dict:
+        if args:
+            schema["items"] = get_type_schema(args[0])
+        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])
+        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)
+            del subschema["type"]
+            schema = {
+                **schema,
+                **subschema,
+            }
+
+    # 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]
+
+    return schema
+
+
+def get_param_schema(param: inspect.Parameter, provider: ProviderFormat):
+    param_schema = get_type_schema(param.annotation)
+
+    if param_schema["type"] is None:
+        raise TypeError(f"Unsupported type: {param.annotation.__name__}")
+
+    if (
+        param.default is not inspect.Parameter.empty
+        and "null" not in param_schema["type"]
+    ):
+        if type(param_schema["type"]) is list:
+            param_schema["type"].append("null")
+        else:
+            param_schema["type"] = [param_schema["type"], "null"]
+
+    if provider == ProviderFormat.GOOGLE_GEMINI:
+        # Filter out "null" type
+        if type(param_schema["type"]) is list:
+            param_schema["type"] = [t for t in param_schema["type"] if t != "null"]
+            if len(param_schema["type"]) == 1:
+                param_schema["type"] = param_schema["type"][0]
+        # Check if there are still multiple types are provided for a single argument
+        if type(param_schema["type"]) is list:
+            logger.warning(
+                f"Gemini does not support a function argument with multiple types e.g. Union[str, int]; defaulting to first found non-null type: {param_schema['type'][0]}"
+            )
+            param_schema["type"] = param_schema["type"][0]
+        # Add nullable property for Gemini
+        param_schema["nullable"] = param.default is not inspect.Parameter.empty
+        if param_schema["type"] == "object":
+            logger.warning(
+                f'Type of argument {param.name} is {param.annotation}, which is being formatted as an "object" type. However, Gemini does not seem to officially support an "object" parameter type yet and success rate might be spotty. Proceed with caution, or refactor {param.name} into one of the basic supported types: [string, integer, boolean, array].'
+            )
+    return param_schema
+
+
+def format_tools(
+    tools: list[Callable],
+    provider: ProviderFormat,
+):
+    """Format tools based on the provider's requirements."""
+
+    # Check for duplicate tool names
+    check_duplicates([t.__name__ for t in tools])
+
+    formatted_tools = []
+    for tool in tools:
+        # Extract parameters and their types from the tool's function signature
+        signature = inspect.signature(tool)
+        parameters = {}
+        required_params = []
+        for param_name, param in signature.parameters.items():
+            parameters[param_name] = get_param_schema(param, provider)
+            required_params.append(param_name)
+
+        # Create formatted tool structure based on provider
+        description = inspect.getdoc(tool) or "No description available."
+        input_schema = {
+            "type": "object",
+            "properties": parameters,
+            "required": required_params,
+        }
+
+        # Format tool based on provider
+        if provider == ProviderFormat.OPENAI_CHAT:
+            formatted_tool = {
+                "type": "function",
+                "function": {
+                    # OpenAI only supports ^[a-zA-Z0-9_-]{1,64}$
+                    "name": tool.__name__.replace(".", "-"),
+                    "description": description,
+                    "parameters": {**input_schema, "additionalProperties": False},
+                    "strict": True,
+                },
+            }
+        elif provider == ProviderFormat.OPENAI_RESPONSES:
+            formatted_tool = {
+                "type": "function",
+                # OpenAI only supports ^[a-zA-Z0-9_-]{1,64}$
+                "name": tool.__name__.replace(".", "-"),
+                "description": description,
+                "parameters": {**input_schema, "additionalProperties": False},
+            }
+        elif provider == ProviderFormat.ANTHROPIC:
+            formatted_tool = {
+                # Claude only supports ^[a-zA-Z0-9_-]{1,64}$
+                "name": tool.__name__.replace(".", "-"),
+                "description": description,
+                "input_schema": input_schema,
+            }
+        elif provider == ProviderFormat.GOOGLE_GEMINI:
+            formatted_tool = {
+                "name": tool.__name__,
+                "parameters": {
+                    "type": "object",
+                    "description": description,
+                    "properties": parameters,
+                    "required": required_params,
+                },
+            }
+
+        formatted_tools.append(formatted_tool)
+    return formatted_tools

stores/indexes/__init__.py

@@ -0,0 +1,11 @@
+from .base_index import BaseIndex
+from .index import Index
+from .local_index import LocalIndex
+from .remote_index import RemoteIndex
+
+__all__ = [
+    "BaseIndex",
+    "Index",
+    "LocalIndex",
+    "RemoteIndex",
+]

stores/indexes/base_index.py

@@ -0,0 +1,269 @@
+import asyncio
+import functools
+import inspect
+import logging
+import re
+from inspect import Parameter
+from types import NoneType, UnionType
+from typing import (
+    Any,
+    Callable,
+    List,
+    Literal,
+    Optional,
+    Tuple,
+    Union,
+    get_args,
+    get_origin,
+    get_type_hints,
+)
+
+from stores.format import ProviderFormat, format_tools
+from stores.parse import llm_parse_json
+from stores.utils import check_duplicates
+
+logging.basicConfig()
+logger = logging.getLogger("stores.indexes.base_index")
+logger.setLevel(logging.INFO)
+
+
+def _cast_arg(value: Any, typ: type | tuple[type]):
+    try:
+        if isinstance(typ, tuple) and len(typ) == 1:
+            typ = typ[0]
+        typ_origin = get_origin(typ)
+        if typ in [float, int, str]:
+            return typ(value)
+        if typ is bool:
+            if isinstance(value, str) and value.lower() == "false":
+                return False
+            else:
+                return typ(value)
+        if typ_origin in (list, List) and isinstance(value, (list, tuple)):
+            return [_cast_arg(v, get_args(typ)) for v in value]
+        if typ_origin in (tuple, Tuple) and isinstance(value, (list, tuple)):
+            return tuple(_cast_arg(v, get_args(typ)) for v in value)
+        if isinstance(typ, type) and typ.__class__.__name__ == "_TypedDictMeta":
+            hints = get_type_hints(typ)
+            for k, v in value.items():
+                value[k] = _cast_arg(v, hints[k])
+            return value
+        if typ_origin in [Union, UnionType]:
+            if NoneType in get_args(typ) and value is None:
+                return value
+            valid_types = [a for a in get_args(typ) if a is not NoneType]
+            if len(valid_types) == 1:
+                return _cast_arg(value, valid_types[0])
+    except Exception:
+        pass
+    # If not in one of the cases above, we return value unchanged
+    return value
+
+
+def _cast_bound_args(bound_args: inspect.BoundArguments):
+    """
+    In some packages, passed argument types are incorrect
+    e.g. LangChain returns float even when argtype is int
+    This only casts basic argtypes
+    """
+    for arg, argparam in bound_args.signature.parameters.items():
+        argtype = argparam.annotation
+        value = bound_args.arguments[arg]
+        new_value = _cast_arg(value, argtype)
+        if new_value != value:
+            # Warn that we are modifying value since this might not be expected
+            logger.warning(
+                f'Argument "{arg}" is type {argtype} but passed value is {value} of type {type(value)} - modifying value to {value} instead.'
+            )
+        bound_args.arguments[arg] = new_value
+
+    return bound_args
+
+
+# TODO: Support more nested types
+def _handle_non_string_literal(annotation: type):
+    origin = get_origin(annotation)
+    if origin is Literal:
+        if any([not isinstance(a, str) for a in get_args(annotation)]):
+            # TODO: Handle duplicates
+            literal_map = {str(a): a for a in get_args(annotation)}
+            new_annotation = Literal.__getitem__(tuple(literal_map.keys()))
+            return new_annotation, literal_map
+        else:
+            return annotation, {}
+    if origin in (list, List):
+        args = get_args(annotation)
+        new_annotation, literal_map = _handle_non_string_literal(args[0])
+        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]
+        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:]:
+            new_annotation, literal_map = _handle_non_string_literal(child_argtype)
+            new_union = new_union | new_annotation
+            union_literal_maps[new_annotation.__name__] = literal_map
+        return new_union, union_literal_maps
+    return annotation, {}
+
+
+# TODO: Support more nested types
+def _undo_non_string_literal(annotation: type, value: Any, literal_map: dict):
+    origin = get_origin(annotation)
+    if origin is Literal:
+        return literal_map.get(value, value)
+    if origin in (list, List) and isinstance(value, (list, tuple)):
+        args = get_args(annotation)
+        return [
+            _undo_non_string_literal(args[0], v, literal_map["item"]) for v in value
+        ]
+    if origin is Union or origin is UnionType:
+        for arg in get_args(annotation):
+            try:
+                return _undo_non_string_literal(arg, value, literal_map[arg.__name__])
+            except Exception:
+                pass
+    return value
+
+
+def wrap_tool(tool: Callable):
+    """
+    Wrap tool to make it compatible with LLM libraries
+    - Gemini does not accept non-None default values
+        If there are any default args, we set default value to None
+        and inject the correct default value at runtime.
+    - Gemini does not accept non-string Literals
+        We convert non-string Literals to strings and reset this at runtime
+    """
+    if hasattr(tool, "_wrapped") and tool._wrapped:
+        return tool
+
+    # Retrieve default arguments
+    original_signature = inspect.signature(tool)
+    new_args = []
+    literal_maps = {}
+    for arg in original_signature.parameters.values():
+        new_arg = arg
+
+        # Handle non-string Literals
+        argtype = new_arg.annotation
+        new_annotation, literal_map = _handle_non_string_literal(argtype)
+        literal_maps[arg.name] = literal_map
+        new_arg = new_arg.replace(
+            kind=Parameter.POSITIONAL_OR_KEYWORD,
+            annotation=new_annotation,
+        )
+
+        # Handle defaults
+        argtype = new_arg.annotation
+        if new_arg.default is Parameter.empty:
+            # If it's annotated with Optional or Union[None, X]
+            # remove the Optional tag since no default value is supplied
+            origin = get_origin(argtype)
+            if (origin in [Union, UnionType]) and NoneType in get_args(argtype):
+                argtype_args = [a for a in get_args(argtype) if a != NoneType]
+                new_annotation = argtype_args[0]
+                for child_argtype in argtype_args[1:]:
+                    new_annotation = new_annotation | child_argtype
+                new_arg = new_arg.replace(
+                    kind=Parameter.POSITIONAL_OR_KEYWORD,
+                    annotation=new_annotation,
+                )
+        else:
+            # 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)]
+            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_arg = new_arg.replace(
+                default=None,
+                kind=Parameter.POSITIONAL_OR_KEYWORD,
+                annotation=new_annotation,
+            )
+        new_args.append(new_arg)
+    new_sig = original_signature.replace(parameters=new_args)
+
+    if inspect.iscoroutinefunction(tool):
+
+        async def wrapper(*args, **kwargs):
+            # Inject default values within wrapper
+            bound_args = original_signature.bind(*args, **kwargs)
+            bound_args.apply_defaults()
+            _cast_bound_args(bound_args)
+            # Inject correct Literals
+            for k, v in bound_args.arguments.items():
+                if k in literal_maps:
+                    param = original_signature.parameters[k]
+                    bound_args.arguments[k] = _undo_non_string_literal(
+                        param.annotation, v, literal_maps[k]
+                    )
+            return await tool(*bound_args.args, **bound_args.kwargs)
+    else:
+
+        def wrapper(*args, **kwargs):
+            # Inject default values within wrapper
+            bound_args = original_signature.bind(*args, **kwargs)
+            bound_args.apply_defaults()
+            _cast_bound_args(bound_args)
+            # Inject correct Literals
+            for k, v in bound_args.arguments.items():
+                if k in literal_maps:
+                    param = original_signature.parameters[k]
+                    bound_args.arguments[k] = _undo_non_string_literal(
+                        param.annotation, v, literal_maps[k]
+                    )
+            return tool(*bound_args.args, **bound_args.kwargs)
+
+    functools.update_wrapper(wrapper, tool)
+    wrapper.__signature__ = new_sig
+    wrapper._wrapped = True
+
+    return wrapper
+
+
+class BaseIndex:
+    def __init__(self, tools: list[Callable]):
+        check_duplicates([t.__name__ for t in tools])
+        self.tools = [wrap_tool(t) for t in tools]
+
+    @property
+    def tools_dict(self):
+        return {tool.__name__: tool for tool in self.tools}
+
+    def execute(self, toolname: str, kwargs: dict | None = None):
+        kwargs = kwargs or {}
+
+        # Use regex since we need to match cases where we perform
+        # substitutions such as replace(".", "-")
+        pattern = re.compile(":?" + re.sub("-|\\.", "(-|\\.)", toolname) + "$")
+
+        matching_tools = []
+        for key in self.tools_dict.keys():
+            if pattern.match(key):
+                matching_tools.append(key)
+        if len(matching_tools) == 0:
+            raise ValueError(f"No tool matching '{toolname}'")
+        elif len(matching_tools) > 1:
+            raise ValueError(f"'{toolname}' matches multiple tools - {matching_tools}")
+        else:
+            toolname = matching_tools[0]
+
+        tool = self.tools_dict[toolname]
+        if inspect.iscoroutinefunction(tool):
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            return loop.run_until_complete(tool(**kwargs))
+        else:
+            return tool(**kwargs)
+
+    def parse_and_execute(self, msg: str):
+        toolcall = llm_parse_json(msg, keys=["toolname", "kwargs"])
+        return self.execute(toolcall.get("toolname"), toolcall.get("kwargs"))
+
+    def format_tools(self, provider: ProviderFormat):
+        return format_tools(self.tools, provider)

stores/indexes/index.py

@@ -0,0 +1,56 @@
+import logging
+import os
+from pathlib import Path
+from typing import Callable
+
+from stores.indexes.base_index import BaseIndex
+from stores.indexes.local_index import LocalIndex
+from stores.indexes.remote_index import RemoteIndex
+
+logging.basicConfig()
+logger = logging.getLogger("stores.index")
+logger.setLevel(logging.INFO)
+
+
+class Index(BaseIndex):
+    def __init__(
+        self,
+        tools: list[Callable, os.PathLike] | None = None,
+        env_var: dict[str, dict] | None = None,
+    ):
+        self.env_var = env_var or {}
+        tools = tools or []
+
+        _tools = []
+        for tool in tools:
+            if isinstance(tool, (str, Path)):
+                index_name = tool
+                loaded_index = None
+                if Path(index_name).exists():
+                    # Load LocalIndex
+                    try:
+                        loaded_index = LocalIndex(index_name)
+                    except Exception:
+                        logger.warning(
+                            f'Unable to load index "{index_name}"', exc_info=True
+                        )
+                if loaded_index is None and isinstance(index_name, str):
+                    # Load RemoteIndex
+                    try:
+                        loaded_index = RemoteIndex(
+                            index_name, env_var=self.env_var.get(index_name)
+                        )
+                    except Exception:
+                        logger.warning(
+                            f'Unable to load index "{index_name}"\nIf this is a local index, make sure it can be found as a directory and contains a tools.toml file.',
+                            exc_info=True,
+                        )
+                if loaded_index is None:
+                    raise ValueError(
+                        f'Unable to load index "{index_name}"\nIf this is a local index, make sure it can be found as a directory and contains a tools.toml file.'
+                    )
+                _tools += loaded_index.tools
+            elif isinstance(tool, Callable):
+                _tools.append(tool)
+
+        super().__init__(_tools)

stores/indexes/local_index.py

@@ -0,0 +1,84 @@
+import importlib
+import logging
+import os
+import sys
+import venv
+from pathlib import Path
+
+from stores.constants import TOOLS_CONFIG_FILENAME, VENV_NAME
+from stores.indexes.base_index import BaseIndex
+from stores.indexes.venv_utils import init_venv_tools, install_venv_deps
+
+if sys.version_info >= (3, 11):
+    import tomllib
+else:
+    import tomli as tomllib
+
+logging.basicConfig()
+logger = logging.getLogger("stores.indexes.local_index")
+logger.setLevel(logging.INFO)
+
+
+class LocalIndex(BaseIndex):
+    def __init__(
+        self,
+        index_folder: os.PathLike,
+        create_venv: bool = False,
+        env_var: dict | None = None,
+    ):
+        self.index_folder = Path(index_folder)
+        self.env_var = env_var or {}
+
+        if not self.index_folder.exists():
+            raise ValueError(
+                f"Unable to load index - {self.index_folder} does not exist"
+            )
+
+        if create_venv:
+            # Create venv and install deps
+            self.venv = self.index_folder / VENV_NAME
+            if not self.venv.exists():
+                venv.create(self.venv, symlinks=True, with_pip=True, upgrade_deps=True)
+            install_venv_deps(self.index_folder)
+            # Initialize tools
+            tools = init_venv_tools(self.index_folder, self.env_var)
+        else:
+            if self.env_var:
+                raise ValueError(
+                    "Environment variables will only be restricted if create_venv=True when initializing LocalIndex"
+                )
+            tools = self._init_tools()
+        super().__init__(tools)
+
+    def _init_tools(self):
+        """
+        Load local tools.toml file and import tool functions
+
+        NOTE: Can we just add index_folder to sys.path and import the functions?
+        """
+        index_manifest = self.index_folder / TOOLS_CONFIG_FILENAME
+        if not index_manifest.exists():
+            raise ValueError(f"Unable to load index - {index_manifest} does not exist")
+
+        with open(index_manifest, "rb") as file:
+            manifest = tomllib.load(file)["index"]
+
+        tools = []
+        for tool_id in manifest.get("tools", []):
+            module_name = ".".join(tool_id.split(".")[:-1])
+            tool_name = tool_id.split(".")[-1]
+
+            module_file = self.index_folder / module_name.replace(".", "/")
+            if (module_file / "__init__.py").exists():
+                module_file = module_file / "__init__.py"
+            else:
+                module_file = Path(str(module_file) + ".py")
+
+            spec = importlib.util.spec_from_file_location(module_name, module_file)
+            module = importlib.util.module_from_spec(spec)
+            sys.modules[spec.name] = module
+            spec.loader.exec_module(module)
+            tool = getattr(module, tool_name)
+            tool.__name__ = tool_id
+            tools.append(tool)
+        return tools