indexify 0.0.42__py3-none-any.whl → 0.2__py3-none-any.whl

indexify/functions_sdk/graph.py

@@ -0,0 +1,276 @@
+import inspect
+from collections import defaultdict
+from typing import (
+    Annotated,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Set,
+    Type,
+    Union,
+    get_args,
+    get_origin,
+)
+
+import cloudpickle
+import msgpack
+from pydantic import BaseModel
+from typing_extensions import get_args, get_origin
+
+from .data_objects import IndexifyData, RouterOutput
+from .graph_validation import validate_node, validate_route
+from .indexify_functions import (
+    IndexifyFunction,
+    IndexifyFunctionWrapper,
+    IndexifyRouter,
+)
+from .object_serializer import CloudPickleSerializer, get_serializer
+
+RouterFn = Annotated[
+    Callable[[IndexifyData], Optional[List[IndexifyFunction]]], "RouterFn"
+]
+GraphNode = Annotated[Union[IndexifyFunctionWrapper, RouterFn], "GraphNode"]
+
+
+def is_pydantic_model_from_annotation(type_annotation):
+    # If it's a string representation
+    if isinstance(type_annotation, str):
+        # Extract the class name from the string
+        class_name = type_annotation.split("'")[-2].split(".")[-1]
+        # This part is tricky and might require additional context or imports
+        # You might need to import the actual class or module where it's defined
+        # For example:
+        # from indexify.functions_sdk.data_objects import File
+        # return issubclass(eval(class_name), BaseModel)
+        return False  # Default to False if we can't evaluate
+
+    # If it's a Type object
+    origin = get_origin(type_annotation)
+    if origin is not None:
+        # Handle generic types like List[File], Optional[File], etc.
+        args = get_args(type_annotation)
+        if args:
+            return is_pydantic_model_from_annotation(args[0])
+
+    # If it's a direct class reference
+    if isinstance(type_annotation, type):
+        return issubclass(type_annotation, BaseModel)
+
+    return False
+
+
+class FunctionMetadata(BaseModel):
+    name: str
+    fn_name: str
+    description: str
+    reducer: bool = False
+    payload_encoder: str = "cloudpickle"
+
+
+class RouterMetadata(BaseModel):
+    name: str
+    description: str
+    source_fn: str
+    target_fns: List[str]
+    payload_encoder: str = "cloudpickle"
+
+
+class NodeMetadata(BaseModel):
+    dynamic_router: Optional[RouterMetadata] = None
+    compute_fn: Optional[FunctionMetadata] = None
+
+
+class ComputeGraphMetadata(BaseModel):
+    name: str
+    description: str
+    start_node: NodeMetadata
+    nodes: Dict[str, NodeMetadata]
+    edges: Dict[str, List[str]]
+    accumulator_zero_values: Dict[str, bytes] = {}
+
+    def get_input_payload_serializer(self):
+        return get_serializer(self.start_node.compute_fn.payload_encoder)
+
+
+class Graph:
+    def __init__(
+        self, name: str, start_node: IndexifyFunction, description: Optional[str] = None
+    ):
+        self.name = name
+        self.description = description
+        self.nodes: Dict[str, Union[IndexifyFunction, IndexifyRouter]] = {}
+        self.routers: Dict[str, List[str]] = defaultdict(list)
+        self.edges: Dict[str, List[str]] = defaultdict(list)
+        self.accumulator_zero_values: Dict[str, Any] = {}
+
+        self.add_node(start_node)
+        self._start_node: str = start_node.name
+
+    def get_function(self, name: str) -> IndexifyFunctionWrapper:
+        if name not in self.nodes:
+            raise ValueError(f"Function {name} not found in graph")
+        return IndexifyFunctionWrapper(self.nodes[name])
+
+    def get_accumulators(self) -> Dict[str, Any]:
+        return self.accumulator_zero_values
+
+    def deserialize_fn_output(self, name: str, output: IndexifyData) -> Any:
+        serializer = get_serializer(self.nodes[name].payload_encoder)
+        return serializer.deserialize(output.payload)
+
+    def add_node(
+        self, indexify_fn: Union[Type[IndexifyFunction], Type[IndexifyRouter]]
+    ) -> "Graph":
+        validate_node(indexify_fn=indexify_fn)
+
+        if indexify_fn.name in self.nodes:
+            return self
+
+        if issubclass(indexify_fn, IndexifyFunction) and indexify_fn.accumulate:
+            self.accumulator_zero_values[
+                indexify_fn.name
+            ] = indexify_fn.accumulate().model_dump()
+
+        self.nodes[indexify_fn.name] = indexify_fn
+        return self
+
+    def route(
+        self, from_node: Type[IndexifyRouter], to_nodes: List[Type[IndexifyFunction]]
+    ) -> "Graph":
+
+        validate_route(from_node=from_node, to_nodes=to_nodes)
+
+        print(
+            f"Adding router {from_node.name} to nodes {[node.name for node in to_nodes]}"
+        )
+        self.add_node(from_node)
+        for node in to_nodes:
+            self.add_node(node)
+            self.routers[from_node.name].append(node.name)
+        return self
+
+    def serialize(self):
+        return cloudpickle.dumps(self)
+
+    @staticmethod
+    def deserialize(graph: bytes) -> "Graph":
+        return cloudpickle.loads(graph)
+
+    @staticmethod
+    def from_path(path: str) -> "Graph":
+        with open(path, "rb") as f:
+            return cloudpickle.load(f)
+
+    def add_edge(
+        self,
+        from_node: Type[IndexifyFunction],
+        to_node: Union[Type[IndexifyFunction], RouterFn],
+    ) -> "Graph":
+        self.add_edges(from_node, [to_node])
+        return self
+
+    def invoke_fn_ser(
+        self, name: str, input: IndexifyData, acc: Optional[Any] = None
+    ) -> List[IndexifyData]:
+        fn_wrapper = self.get_function(name)
+        input = self.deserialize_input(name, input)
+        serializer = get_serializer(fn_wrapper.indexify_function.payload_encoder)
+        if acc is not None:
+            acc = fn_wrapper.indexify_function.accumulate.model_validate(
+                serializer.deserialize(acc.payload)
+            )
+        if acc is None and fn_wrapper.indexify_function.accumulate is not None:
+            acc = fn_wrapper.indexify_function.accumulate.model_validate(
+                self.accumulator_zero_values[name]
+            )
+        outputs: List[Any] = fn_wrapper.run_fn(input, acc=acc)
+        return [
+            IndexifyData(payload=serializer.serialize(output)) for output in outputs
+        ]
+
+    def invoke_router(self, name: str, input: IndexifyData) -> Optional[RouterOutput]:
+        fn_wrapper = self.get_function(name)
+        input = self.deserialize_input(name, input)
+        return RouterOutput(edges=fn_wrapper.run_router(input))
+
+    def deserialize_input(self, compute_fn: str, indexify_data: IndexifyData) -> Any:
+        compute_fn = self.nodes[compute_fn]
+        if not compute_fn:
+            raise ValueError(f"Compute function {compute_fn} not found in graph")
+        if compute_fn.payload_encoder == "cloudpickle":
+            return CloudPickleSerializer.deserialize(indexify_data.payload)
+        payload = msgpack.unpackb(indexify_data.payload)
+        signature = inspect.signature(compute_fn.run)
+        arg_types = {}
+        for name, param in signature.parameters.items():
+            if (
+                param.annotation != inspect.Parameter.empty
+                and param.annotation != getattr(compute_fn, "accumulate", None)
+            ):
+                arg_types[name] = param.annotation
+        if len(arg_types) > 1:
+            raise ValueError(
+                f"Compute function {compute_fn} has multiple arguments, but only one is supported"
+            )
+        elif len(arg_types) == 0:
+            raise ValueError(f"Compute function {compute_fn} has no arguments")
+        arg_name, arg_type = next(iter(arg_types.items()))
+        if arg_type is None:
+            raise ValueError(f"Argument {arg_name} has no type annotation")
+        if is_pydantic_model_from_annotation(arg_type):
+            if len(payload.keys()) == 1 and isinstance(list(payload.values())[0], dict):
+                payload = list(payload.values())[0]
+            return arg_type.model_validate(payload)
+        return payload
+
+    def add_edges(
+        self,
+        from_node: Union[Type[IndexifyFunction], Type[IndexifyRouter]],
+        to_node: List[Union[Type[IndexifyFunction], Type[IndexifyRouter]]],
+    ) -> "Graph":
+        self.add_node(from_node)
+        from_node_name = from_node.name
+        for node in to_node:
+            self.add_node(node)
+            self.edges[from_node_name].append(node.name)
+        return self
+
+    def definition(self) -> ComputeGraphMetadata:
+        start_node = self.nodes[self._start_node]
+        start_node = FunctionMetadata(
+            name=start_node.name,
+            fn_name=start_node.fn_name,
+            description=start_node.description,
+            reducer=start_node.accumulate is not None,
+        )
+        metadata_edges = self.edges.copy()
+        metadata_nodes = {}
+        for node_name, node in self.nodes.items():
+            if node_name in self.routers:
+                metadata_nodes[node_name] = NodeMetadata(
+                    dynamic_router=RouterMetadata(
+                        name=node_name,
+                        description=node.description or "",
+                        source_fn=node_name,
+                        target_fns=self.routers[node_name],
+                        payload_encoder=node.payload_encoder,
+                    )
+                )
+            else:
+                metadata_nodes[node_name] = NodeMetadata(
+                    compute_fn=FunctionMetadata(
+                        name=node_name,
+                        fn_name=node.fn_name,
+                        description=node.description,
+                        reducer=node.accumulate is not None,
+                    )
+                )
+        return ComputeGraphMetadata(
+            name=self.name,
+            description=self.description or "",
+            start_node=NodeMetadata(compute_fn=start_node),
+            nodes=metadata_nodes,
+            edges=metadata_edges,
+        )

indexify/functions_sdk/graph_validation.py

@@ -0,0 +1,69 @@
+import inspect
+import re
+from typing import List, Type, Union
+
+from .indexify_functions import IndexifyFunction, IndexifyRouter
+
+
+def validate_node(indexify_fn: Union[Type[IndexifyFunction], Type[IndexifyRouter]]):
+    if inspect.isfunction(indexify_fn):
+        raise Exception(
+            f"Unable to add node of type `{type(indexify_fn)}`. "
+            f"Required, `IndexifyFunction` or `IndexifyRouter`"
+        )
+    if not (
+        issubclass(indexify_fn, IndexifyFunction)
+        or issubclass(indexify_fn, IndexifyRouter)
+    ):
+        raise Exception(
+            f"Unable to add node of type `{indexify_fn.__name__}`. "
+            f"Required, `IndexifyFunction` or `IndexifyRouter`"
+        )
+
+    signature = inspect.signature(indexify_fn.run)
+
+    for param in signature.parameters.values():
+        if param.name == "self":
+            continue
+        if param.annotation == inspect.Parameter.empty:
+            raise Exception(
+                f"Input param {param.name} in {indexify_fn.name} has empty"
+                f" type annotation"
+            )
+
+    if signature.return_annotation == inspect.Signature.empty:
+        raise Exception(f"Function {indexify_fn.name} has empty return type annotation")
+
+
+def validate_route(
+    from_node: Type[IndexifyRouter], to_nodes: List[Type[IndexifyFunction]]
+):
+    signature = inspect.signature(from_node.run)
+
+    if signature.return_annotation == inspect.Signature.empty:
+        raise Exception(f"Function {from_node.name} has empty return type annotation")
+
+    # We lose the exact type string when the object is created
+    source = inspect.getsource(from_node.run)
+
+    union_pattern = r"Union\[((?:\w+(?:,\s*)?)+)\]"
+    union_match = re.search(union_pattern, source)
+
+    src_route_nodes = None
+    if union_match:
+        # nodes = re.findall(r'\w+', match.group(1))
+        src_route_nodes = [node.strip() for node in union_match.group(1).split(",")]
+        if len(src_route_nodes) <= 1:
+            raise Exception(f"Invalid router for {from_node.name}, lte 1 route.")
+    else:
+        raise Exception(
+            f"Invalid router for {from_node.name}, cannot find output nodes"
+        )
+
+    to_node_names = [i.name for i in to_nodes]
+
+    for src_node in src_route_nodes:
+        if src_node not in to_node_names:
+            raise Exception(
+                f"Unable to find {src_node} in to_nodes " f"{to_node_names}"
+            )

indexify/functions_sdk/image.py

@@ -0,0 +1,26 @@
+class Image:
+    def __init__(self):
+        self._image_name = None
+
+        self._tag = "latest"
+
+        self._base_image = None
+
+        self._run_strs = []
+        pass
+
+    def image_name(self, image_name):
+        self._image_name = image_name
+        return self
+
+    def tag(self, tag):
+        self._tag = tag
+        return self
+
+    def base_image(self, base_image):
+        self._base_image = base_image
+        return self
+
+    def run(self, run_str):
+        self._run_strs.append(run_str)
+        return self

indexify/functions_sdk/indexify_functions.py

@@ -0,0 +1,192 @@
+from abc import ABC, abstractmethod
+from functools import update_wrapper
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Type,
+    Union,
+    get_args,
+    get_origin,
+)
+
+from pydantic import BaseModel
+from typing_extensions import get_type_hints
+
+from .data_objects import IndexifyData, RouterOutput
+from .image import Image
+
+
+class EmbeddingIndexes(BaseModel):
+    dim: int
+    distance: Optional[str] = "cosine"
+    database_url: Optional[str] = None
+
+
+class PlacementConstraints(BaseModel):
+    min_python_version: Optional[str] = "3.9"
+    max_python_version: Optional[str] = None
+    platform: Optional[str] = None
+    image_name: Optional[str] = None
+
+
+class IndexifyFunction(ABC):
+    name: str = ""
+    base_image: Optional[str] = None
+    description: str = ""
+    placement_constraints: List[PlacementConstraints] = []
+    accumulate: Optional[Type[Any]] = None
+    payload_encoder: Optional[str] = "cloudpickle"
+
+    @abstractmethod
+    def run(self, *args, **kwargs) -> Union[List[Any], Any]:
+        pass
+
+    def partial(self, **kwargs) -> Callable:
+        from functools import partial
+
+        return partial(self.run, **kwargs)
+
+
+class IndexifyRouter(ABC):
+    name: str = ""
+    description: str = ""
+    image: Image = None
+    placement_constraints: List[PlacementConstraints] = []
+    payload_encoder: Optional[str] = "cloudpickle"
+
+    @abstractmethod
+    def run(self, *args, **kwargs) -> Optional[List[IndexifyFunction]]:
+        pass
+
+
+def indexify_router(
+    name: Optional[str] = None,
+    description: Optional[str] = "",
+    image: Image = None,
+    placement_constraints: List[PlacementConstraints] = [],
+    output_encoder: Optional[str] = "cloudpickle",
+):
+    def construct(fn):
+        args = locals().copy()
+        args["name"] = args["name"] if args.get("name", None) else fn.__name__
+        args["fn_name"] = fn.__name__
+        args["description"] = (
+            args["description"]
+            if args.get("description", None)
+            else (fn.__doc__ or "").strip().replace("\n", "")
+        )
+
+        class IndexifyRo(IndexifyRouter):
+            def run(self, *args, **kwargs) -> Optional[List[IndexifyFunction]]:
+                return fn(*args, **kwargs)
+
+            update_wrapper(run, fn)
+
+        for key, value in args.items():
+            if key != "fn" and key != "self":
+                setattr(IndexifyRo, key, value)
+
+        IndexifyRo.image = image
+        IndexifyRo.payload_encoder = output_encoder
+        return IndexifyRo
+
+    return construct
+
+
+def indexify_function(
+    name: Optional[str] = None,
+    description: Optional[str] = "",
+    image: Image = None,
+    accumulate: Optional[Type[BaseModel]] = None,
+    min_batch_size: Optional[int] = None,
+    max_batch_size: Optional[int] = None,
+    output_encoder: Optional[str] = "cloudpickle",
+    placement_constraints: List[PlacementConstraints] = [],
+):
+    def construct(fn):
+        args = locals().copy()
+        args["name"] = args["name"] if args.get("name", None) else fn.__name__
+        args["fn_name"] = fn.__name__
+        args["description"] = (
+            args["description"]
+            if args.get("description", None)
+            else (fn.__doc__ or "").strip().replace("\n", "")
+        )
+
+        class IndexifyFn(IndexifyFunction):
+            def run(self, *args, **kwargs) -> Union[List[Any], Any]:
+                return fn(*args, **kwargs)
+
+            update_wrapper(run, fn)
+
+        for key, value in args.items():
+            if key != "fn" and key != "self":
+                setattr(IndexifyFn, key, value)
+
+        IndexifyFn.image = image
+        IndexifyFn.accumulate = accumulate
+        IndexifyFn.min_batch_size = min_batch_size
+        IndexifyFn.max_batch_size = max_batch_size
+        IndexifyFn.payload_encoder = output_encoder
+        return IndexifyFn
+
+    return construct
+
+
+class IndexifyFunctionWrapper:
+    def __init__(self, indexify_function: Union[IndexifyFunction, IndexifyRouter]):
+        self.indexify_function: Union[
+            IndexifyFunction, IndexifyRouter
+        ] = indexify_function()
+
+    def get_output_model(self) -> Any:
+        if not isinstance(self.indexify_function, IndexifyFunction):
+            raise TypeError("Input must be an instance of IndexifyFunction")
+
+        extract_method = self.indexify_function.run
+        type_hints = get_type_hints(extract_method)
+        return_type = type_hints.get("return", Any)
+        if get_origin(return_type) is list:
+            return_type = get_args(return_type)[0]
+        elif get_origin(return_type) is Union:
+            inner_types = get_args(return_type)
+            if len(inner_types) == 2 and type(None) in inner_types:
+                return_type = (
+                    inner_types[0] if inner_types[1] is type(None) else inner_types[1]
+                )
+        return return_type
+
+    def run_router(self, input: Union[Dict, Type[BaseModel]]) -> List[str]:
+        kwargs = input if isinstance(input, dict) else {"input": input}
+        args = []
+        kwargs = {}
+        if isinstance(input, dict):
+            kwargs = input
+        else:
+            args.append(input)
+        extracted_data = self.indexify_function.run(*args, **kwargs)
+        if not isinstance(extracted_data, list) and extracted_data is not None:
+            return [extracted_data.name]
+        edges = []
+        for fn in extracted_data or []:
+            edges.append(fn.name)
+        return edges
+
+    def run_fn(
+        self, input: Union[Dict, Type[BaseModel]], acc: Type[Any] = None
+    ) -> List[IndexifyData]:
+        args = []
+        kwargs = {}
+        if acc is not None:
+            args.append(acc)
+        if isinstance(input, dict):
+            kwargs = input
+        else:
+            args.append(input)
+
+        extracted_data = self.indexify_function.run(*args, **kwargs)
+
+        return extracted_data if isinstance(extracted_data, list) else [extracted_data]

indexify/functions_sdk/local_cache.py

@@ -0,0 +1,46 @@
+import os
+from hashlib import sha256
+from typing import List, Optional
+
+
+class CacheAwareFunctionWrapper:
+    def __init__(self, cache_dir: str):
+        self._cache_dir = cache_dir
+        if not os.path.exists(cache_dir):
+            os.makedirs(cache_dir)
+
+    def _get_key(self, input: bytes) -> str:
+        h = sha256()
+        h.update(input)
+        return h.hexdigest()
+
+    def get(self, graph: str, node_name: str, input: bytes) -> Optional[List[bytes]]:
+        key = self._get_key(input)
+        dir_path = os.path.join(self._cache_dir, graph, node_name, key)
+        if not os.path.exists(dir_path):
+            return None
+
+        files = os.listdir(dir_path)
+        outputs = []
+        for file in files:
+            with open(os.path.join(dir_path, file), "rb") as f:
+                return f.read()
+
+        return outputs
+
+    def set(
+        self,
+        graph: str,
+        node_name: str,
+        input: bytes,
+        output: List[bytes],
+    ):
+        key = self._get_key(input)
+        dir_path = os.path.join(self._cache_dir, graph, node_name, key)
+        if not os.path.exists(dir_path):
+            os.makedirs(dir_path)
+
+        for i, output_item in enumerate(output):
+            file_path = os.path.join(dir_path, f"{i}.cbor")
+            with open(file_path, "wb") as f:
+                f.write(output_item)

indexify/functions_sdk/object_serializer.py

@@ -0,0 +1,61 @@
+from typing import Any, List
+
+import cloudpickle
+import msgpack
+from pydantic import BaseModel
+
+from .data_objects import IndexifyData
+
+
+def get_serializer(serializer_type: str) -> Any:
+    if serializer_type == "cloudpickle":
+        return CloudPickleSerializer()
+    elif serializer_type == "msgpack":
+        return MsgPackSerializer()
+    else:
+        raise ValueError(f"Unknown serializer type: {serializer_type}")
+
+
+class CloudPickleSerializer:
+    @staticmethod
+    def serialize(data: Any) -> bytes:
+        return cloudpickle.dumps(data)
+
+    @staticmethod
+    def deserialize(data: bytes) -> Any:
+        return cloudpickle.loads(data)
+
+    @staticmethod
+    def serialize_list(data: List[Any]) -> bytes:
+        return cloudpickle.dumps(data)
+
+    @staticmethod
+    def deserialize_list(data: bytes) -> List[Any]:
+        return cloudpickle.loads(data)
+
+
+class MsgPackSerializer:
+    @staticmethod
+    def serialize(data: Any) -> bytes:
+        if (
+            isinstance(data, type)
+            and issubclass(data, BaseModel)
+            or isinstance(data, BaseModel)
+        ):
+            return msgpack.packb(data.model_dump())
+        return msgpack.packb(data)
+
+    @staticmethod
+    def deserialize(data: bytes) -> IndexifyData:
+        cached_output = msgpack.unpackb(data)
+        print(cached_output)
+        return IndexifyData(**cached_output)
+
+    @staticmethod
+    def serialize_list(data: List[IndexifyData]) -> bytes:
+        data = [item.model_dump() for item in data]
+        return msgpack.packb(data)
+
+    @staticmethod
+    def deserialize_list(data: bytes) -> List[IndexifyData]:
+        return [IndexifyData(**item) for item in msgpack.unpackb(data)]