indexify 0.2.3__tar.gz → 0.2.4__tar.gz

{indexify-0.2.3 → indexify-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: indexify
-Version: 0.2.3
+Version: 0.2.4
 Summary: Python Client for Indexify
 Home-page: https://github.com/tensorlakeai/indexify
 License: Apache 2.0

{indexify-0.2.3 → indexify-0.2.4}/indexify/__init__.py

@@ -1,23 +1,21 @@
 from . import data_loaders
-from .client import create_client
 from .functions_sdk.graph import Graph
 from .functions_sdk.image import Image
 from .functions_sdk.indexify_functions import (
     indexify_function,
     indexify_router,
 )
-from .local_client import LocalClient
-from .remote_client import RemoteClient
+from .http_client import IndexifyClient
+from .remote_graph import RemoteGraph
 from .settings import DEFAULT_SERVICE_URL
 
 __all__ = [
     "data_loaders",
     "Graph",
+    "RemoteGraph",
     "Image",
     "indexify_function",
     "indexify_router",
     "DEFAULT_SERVICE_URL",
-    "RemoteClient",
-    "LocalClient",
-    "create_client",
+    "IndexifyClient",
 ]

{indexify-0.2.3 → indexify-0.2.4}/indexify/cli.py

@@ -19,7 +19,7 @@ from rich.theme import Theme
 
 from indexify.executor.agent import ExtractorAgent
 from indexify.executor.function_worker import FunctionWorker
-from indexify.functions_sdk.image import Image, DEFAULT_IMAGE
+from indexify.functions_sdk.image import DEFAULT_IMAGE, Image
 
 custom_theme = Theme(
     {

{indexify-0.2.3 → indexify-0.2.4}/indexify/functions_sdk/graph.py

@@ -1,5 +1,7 @@
 import inspect
+import sys
 from collections import defaultdict
+from queue import deque
 from typing import (
     Annotated,
     Any,
@@ -7,7 +9,6 @@ from typing import (
     Dict,
     List,
     Optional,
-    Set,
     Type,
     Union,
     get_args,
@@ -16,16 +17,25 @@ from typing import (
 
 import cloudpickle
 import msgpack
+from nanoid import generate
 from pydantic import BaseModel
 from typing_extensions import get_args, get_origin
 
 from .data_objects import IndexifyData, RouterOutput
+from .graph_definition import (
+    ComputeGraphMetadata,
+    FunctionMetadata,
+    NodeMetadata,
+    RouterMetadata,
+    RuntimeInformation,
+)
 from .graph_validation import validate_node, validate_route
 from .indexify_functions import (
     IndexifyFunction,
     IndexifyFunctionWrapper,
     IndexifyRouter,
 )
+from .local_cache import CacheAwareFunctionWrapper
 from .object_serializer import CloudPickleSerializer, get_serializer
 
 RouterFn = Annotated[
@@ -61,41 +71,6 @@ def is_pydantic_model_from_annotation(type_annotation):
     return False
 
 
-class FunctionMetadata(BaseModel):
-    name: str
-    fn_name: str
-    description: str
-    reducer: bool = False
-    image_name: str
-    payload_encoder: str = "cloudpickle"
-
-
-class RouterMetadata(BaseModel):
-    name: str
-    description: str
-    source_fn: str
-    target_fns: List[str]
-    image_name: 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
@@ -110,6 +85,11 @@ class Graph:
         self.add_node(start_node)
         self._start_node: str = start_node.name
 
+        # Storage for local execution
+        self._results: Dict[str, Dict[str, List[IndexifyData]]] = {}
+        self._cache = CacheAwareFunctionWrapper("./indexify_local_runner_cache")
+        self._accumulator_values: Dict[str, Dict[str, IndexifyData]] = {}
+
     def get_function(self, name: str) -> IndexifyFunctionWrapper:
         if name not in self.nodes:
             raise ValueError(f"Function {name} not found in graph")
@@ -272,10 +252,123 @@ class Graph:
                         image_name=node.image._image_name,
                     )
                 )
+
         return ComputeGraphMetadata(
             name=self.name,
             description=self.description or "",
             start_node=NodeMetadata(compute_fn=start_node),
             nodes=metadata_nodes,
             edges=metadata_edges,
+            runtime_information=RuntimeInformation(
+                major_version=sys.version_info.major,
+                minor_version=sys.version_info.minor,
+            ),
         )
+
+    def run(self, block_until_done: bool = False, **kwargs) -> str:
+        start_node = self.nodes[self._start_node]
+        serializer = get_serializer(start_node.payload_encoder)
+        input = IndexifyData(id=generate(), payload=serializer.serialize(kwargs))
+        print(f"[bold] Invoking {self._start_node}[/bold]")
+        outputs = defaultdict(list)
+        self._accumulator_values[input.id] = {}
+        for k, v in self.accumulator_zero_values.items():
+            node = self.nodes[k]
+            serializer = get_serializer(node.payload_encoder)
+            self._accumulator_values[input.id] = {
+                k: IndexifyData(payload=serializer.serialize(v))
+            }
+        self._results[input.id] = outputs
+        self._run(input, outputs)
+        return input.id
+
+    def _run(
+        self,
+        initial_input: IndexifyData,
+        outputs: Dict[str, List[bytes]],
+    ):
+        accumulator_values = self._accumulator_values[initial_input.id]
+        queue = deque([(self._start_node, initial_input)])
+        while queue:
+            node_name, input = queue.popleft()
+            node = self.nodes[node_name]
+            serializer = get_serializer(node.payload_encoder)
+            input_bytes = serializer.serialize(input)
+            cached_output_bytes: Optional[bytes] = self._cache.get(
+                self.name, node_name, input_bytes
+            )
+            if cached_output_bytes is not None:
+                print(
+                    f"ran {node_name}: num outputs: {len(cached_output_bytes)} (cache hit)"
+                )
+                function_outputs: List[IndexifyData] = []
+                cached_output_list = serializer.deserialize_list(cached_output_bytes)
+                if accumulator_values.get(node_name, None) is not None:
+                    accumulator_values[node_name] = cached_output_list[-1].model_copy()
+                    outputs[node_name] = []
+                function_outputs.extend(cached_output_list)
+                outputs[node_name].extend(cached_output_list)
+            else:
+                function_outputs: List[IndexifyData] = self.invoke_fn_ser(
+                    node_name, input, accumulator_values.get(node_name, None)
+                )
+                print(f"ran {node_name}: num outputs: {len(function_outputs)}")
+                if accumulator_values.get(node_name, None) is not None:
+                    accumulator_values[node_name] = function_outputs[-1].model_copy()
+                    outputs[node_name] = []
+                outputs[node_name].extend(function_outputs)
+                function_outputs_bytes: List[bytes] = [
+                    serializer.serialize_list(function_outputs)
+                ]
+                self._cache.set(
+                    self.name,
+                    node_name,
+                    input_bytes,
+                    function_outputs_bytes,
+                )
+            if accumulator_values.get(node_name, None) is not None and queue:
+                print(
+                    f"accumulator not none for {node_name}, continuing, len queue: {len(queue)}"
+                )
+                continue
+
+            out_edges = self.edges.get(node_name, [])
+            # Figure out if there are any routers for this node
+            for i, edge in enumerate(out_edges):
+                if edge in self.routers:
+                    out_edges.remove(edge)
+                    for output in function_outputs:
+                        dynamic_edges = self._route(edge, output) or []
+                        for dynamic_edge in dynamic_edges.edges:
+                            if dynamic_edge in self.nodes:
+                                print(
+                                    f"[bold]dynamic router returned node: {dynamic_edge}[/bold]"
+                                )
+                                out_edges.append(dynamic_edge)
+            for out_edge in out_edges:
+                for output in function_outputs:
+                    queue.append((out_edge, output))
+
+    def _route(self, node_name: str, input: IndexifyData) -> Optional[RouterOutput]:
+        return self.invoke_router(node_name, input)
+
+    def get_output(
+        self,
+        invocation_id: str,
+        fn_name: str,
+    ) -> List[Any]:
+        results = self._results[invocation_id]
+        if fn_name not in results:
+            raise ValueError(f"no results found for fn {fn_name} on graph {self.name}")
+        fn = self.nodes[fn_name]
+        fn_model = self.get_function(fn_name).get_output_model()
+        serializer = get_serializer(fn.payload_encoder)
+        outputs = []
+        for result in results[fn_name]:
+            payload_dict = serializer.deserialize(result.payload)
+            if issubclass(fn_model, BaseModel) and isinstance(payload_dict, dict):
+                payload = fn_model.model_validate(payload_dict)
+            else:
+                payload = payload_dict
+            outputs.append(payload)
+        return outputs

indexify-0.2.4/indexify/functions_sdk/graph_definition.py

@@ -0,0 +1,47 @@
+from typing import Dict, List, Optional
+
+from pydantic import BaseModel
+
+from .object_serializer import get_serializer
+
+
+class FunctionMetadata(BaseModel):
+    name: str
+    fn_name: str
+    description: str
+    reducer: bool = False
+    image_name: str
+    payload_encoder: str = "cloudpickle"
+
+
+class RouterMetadata(BaseModel):
+    name: str
+    description: str
+    source_fn: str
+    target_fns: List[str]
+    image_name: str
+    payload_encoder: str = "cloudpickle"
+
+
+class NodeMetadata(BaseModel):
+    dynamic_router: Optional[RouterMetadata] = None
+    compute_fn: Optional[FunctionMetadata] = None
+
+
+# RuntimeInformation is a class that holds data about the environment in which the graph should run.
+class RuntimeInformation(BaseModel):
+    major_version: int
+    minor_version: int
+
+
+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] = {}
+    runtime_information: RuntimeInformation
+
+    def get_input_payload_serializer(self):
+        return get_serializer(self.start_node.compute_fn.payload_encoder)

{indexify-0.2.3 → indexify-0.2.4}/indexify/functions_sdk/graph_validation.py

@@ -42,25 +42,30 @@ def validate_route(
 
     if signature.return_annotation == inspect.Signature.empty:
         raise Exception(f"Function {from_node.name} has empty return type annotation")
-    
+
     return_annotation = signature.return_annotation
 
-    if hasattr(return_annotation, '__origin__') and return_annotation.__origin__ is Union:
+    if (
+        hasattr(return_annotation, "__origin__")
+        and return_annotation.__origin__ is Union
+    ):
         for arg in return_annotation.__args__:
-            if hasattr(arg, 'name'):
+            if hasattr(arg, "name"):
                 if arg not in to_nodes:
                     raise Exception(
                         f"Unable to find {arg.name} in to_nodes {[node.name for node in to_nodes]}"
                     )
 
-    if hasattr(return_annotation, '__origin__') and return_annotation.__origin__ is list:
+    if (
+        hasattr(return_annotation, "__origin__")
+        and return_annotation.__origin__ is list
+    ):
         union_args = return_annotation.__args__[0].__args__
         for arg in union_args:
-            if hasattr(arg, 'name'):
+            if hasattr(arg, "name"):
                 if arg not in to_nodes:
                     raise Exception(
                         f"Unable to find {arg.name} in to_nodes {[node.name for node in to_nodes]}"
                     )
     else:
         raise Exception(f"Return type of {from_node.name} is not a Union")
-

indexify-0.2.3/indexify/remote_client.py → indexify-0.2.4/indexify/http_client.py

@@ -10,7 +10,6 @@ from httpx_sse import connect_sse
 from pydantic import BaseModel, Json
 from rich import print
 
-from indexify.base_client import IndexifyClient
 from indexify.error import ApiException
 from indexify.functions_sdk.data_objects import IndexifyData
 from indexify.functions_sdk.graph import ComputeGraphMetadata, Graph
@@ -39,7 +38,7 @@ class GraphOutputs(BaseModel):
     outputs: List[GraphOutputMetadata]
 
 
-class RemoteClient(IndexifyClient):
+class IndexifyClient:
     def __init__(
         self,
         service_url: str = DEFAULT_SERVICE_URL,
@@ -167,7 +166,6 @@ class RemoteClient(IndexifyClient):
             files={"code": serialized_code},
             data={"compute_graph": graph_metadata.model_dump_json(exclude_none=True)},
         )
-        print(response.content.decode("utf-8"))
         response.raise_for_status()
         self._graphs[graph.name] = graph
 

indexify-0.2.4/indexify/remote_graph.py

@@ -0,0 +1,74 @@
+from typing import Any, List, Optional
+
+from indexify.functions_sdk.graph import Graph
+
+from .http_client import IndexifyClient
+from .settings import DEFAULT_SERVICE_URL
+
+
+class RemoteGraph:
+    def __init__(
+        self,
+        name: str,
+        server_url: Optional[str] = DEFAULT_SERVICE_URL,
+    ):
+        self._name = name
+        self._client = IndexifyClient(service_url=server_url)
+
+    def run(self, block_until_done: bool = False, **kwargs) -> str:
+        """
+        Run the graph with the given inputs. The input is for the start function of the graph.
+        :param block_until_done: If True, the function will block until the graph execution is complete.
+        :param kwargs: The input to the start function of the graph. Pass the input as keyword arguments.
+        :return: The invocation ID of the graph execution.
+
+        Example:
+        @indexify_function()
+        def foo(x: int) -> int:
+            return x + 1
+        remote_graph = RemoteGraph.by_name("test")
+        invocation_id = remote_graph.run(x=1)
+        """
+        return self._client.invoke_graph_with_object(
+            self._name, block_until_done, **kwargs
+        )
+
+    @classmethod
+    def deploy(cls, g: Graph, server_url: Optional[str] = "http://localhost:8090"):
+        """
+        Create a new RemoteGraph from a local Graph object.
+        :param g: The local Graph object.
+        :param server_url: The URL of the server where the graph will be registered.
+        """
+        client = IndexifyClient(service_url=server_url)
+        client.register_compute_graph(g)
+        return cls(name=g.name, server_url=server_url)
+
+    @classmethod
+    def by_name(cls, name: str, server_url: Optional[str] = "http://localhost:8090"):
+        """
+        Create a handle to call a RemoteGraph by name.
+        :param name: The name of the graph.
+        :param server_url: The URL of the server where the graph is registered.
+        :return: A RemoteGraph object.
+        """
+        return cls(name=name, server_url=server_url)
+
+    def get_output(
+        self,
+        invocation_id: str,
+        fn_name: str,
+    ) -> List[Any]:
+        """
+        Returns the extracted objects by a graph for an ingested object. If the extractor name is provided, only the objects extracted by that extractor are returned.
+        If the extractor name is not provided, all the extracted objects are returned for the input object.
+        invocation_id: str: The ID of the ingested object
+        fn_name: Optional[str]: The name of the function whose output is to be returned if provided
+        return: List[Any]: Output of the function.
+        """
+
+        return self._client.graph_outputs(
+            graph=self._name,
+            invocation_id=invocation_id,
+            fn_name=fn_name,
+        )

{indexify-0.2.3 → indexify-0.2.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "indexify"
-version = "0.2.3"
+version = "0.2.4"
 description = "Python Client for Indexify"
 authors = ["Tensorlake Inc. <support@tensorlake.ai>"]
 license = "Apache 2.0"

indexify-0.2.3/indexify/base_client.py

@@ -1,98 +0,0 @@
-from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Optional, Union
-
-from pydantic import Json
-
-from indexify.functions_sdk.graph import Graph
-
-
-class IndexifyClient(ABC):
-
-    ### Operational APIs
-    @abstractmethod
-    def register_compute_graph(self, graph: Graph):
-        """
-        Register a compute graph.
-        graph: Graph: The graph to be registered
-        """
-        pass
-
-    @abstractmethod
-    def graphs(self) -> List[str]:
-        """
-        Get the graphs.
-        return: List[str]: The graphs
-        """
-        pass
-
-    @abstractmethod
-    def namespaces(self) -> List[str]:
-        """
-        Get the namespaces.
-        return: List[str]: The namespaces
-        """
-        pass
-
-    @abstractmethod
-    def create_namespace(self, namespace: str):
-        """
-        Create a namespace.
-        namespace: str: The name of the namespace to be created
-        """
-        pass
-
-    ### Ingestion APIs
-    @abstractmethod
-    def invoke_graph_with_object(
-        self, graph: str, block_until_done: bool = False, **kwargs
-    ) -> str:
-        """
-        Invokes a graph with an input object.
-        graph: str: The name of the graph to invoke
-        kwargs: Any: Named arguments to be passed to the graph. Example: url="https://www.google.com", web_page_text="Hello world!"
-        return: str: The ID of the ingested object
-        """
-        pass
-
-    @abstractmethod
-    def invoke_graph_with_file(
-        self,
-        graph: str,
-        path: str,
-        metadata: Optional[Dict[str, Json]] = None,
-        block_until_done: bool = False,
-    ) -> str:
-        """
-        Invokes a graph with an input file. The file's mimetype is appropriately detected.
-        graph: str: The name of the graph to invoke
-        path: str: The path to the file to be ingested
-        return: str: The ID of the ingested object
-        """
-        pass
-
-    @abstractmethod
-    def rerun_graph(self, graph: str):
-        """
-        Rerun a graph.
-        graph: str: The name of the graph to rerun
-        """
-        pass
-
-    ### Retrieval APIs
-    @abstractmethod
-    def graph_outputs(
-        self,
-        graph: str,
-        invocation_id: str,
-        fn_name: Optional[str],
-    ) -> Union[Dict[str, List[Any]], List[Any]]:
-        """
-        Returns the extracted objects by a graph for an ingested object. If the extractor name is provided, only the objects extracted by that extractor are returned.
-        If the extractor name is not provided, all the extracted objects are returned for the input object.
-        graph: str: The name of the graph
-        ingested_object_id: str: The ID of the ingested object
-        extractor_name: Optional[str]: The name of the extractor whose output is to be returned if provided
-        block_until_done: bool = True: If True, the method will block until the extraction is done. If False, the method will return immediately.
-        return: Union[Dict[str, List[Any]], List[Any]]: The extracted objects. If the extractor name is provided, the output is a list of extracted objects by the extractor. If the extractor name is not provided, the output is a dictionary with the extractor name as the key and the extracted objects as the value. If no objects are found, an empty list is returned.
-        """
-        pass

indexify-0.2.3/indexify/client.py

@@ -1,18 +0,0 @@
-from typing import Optional
-
-from .base_client import IndexifyClient
-from .local_client import LocalClient
-from .remote_client import RemoteClient
-from .settings import DEFAULT_SERVICE_URL
-
-
-def create_client(
-    service_url: str = DEFAULT_SERVICE_URL,
-    config_path: Optional[str] = None,
-    in_process: bool = False,
-    *args,
-    **kwargs,
-) -> IndexifyClient:
-    if in_process:
-        return LocalClient()
-    return RemoteClient(config_path=config_path, service_url=service_url, **kwargs)

indexify-0.2.3/indexify/foo

@@ -1,72 +0,0 @@
-./functions_sdk/local_cache.py:        outputs = []
-./functions_sdk/local_cache.py:                outputs.append(f.read())
-./functions_sdk/local_cache.py:        return outputs
-./functions_sdk/graph.py:        outputs: List[Any] = fn_wrapper.run_fn(input, acc=acc)
-./functions_sdk/graph.py:            IndexifyData(payload=CborSerializer.serialize(output)) for output in outputs
-./functions_sdk/output_serializer.py:class OutputSerializer:
-./functions_sdk/output_serializer.py:        normalized_outputs = []
-./functions_sdk/output_serializer.py:            normalized_outputs.append(
-./functions_sdk/output_serializer.py:        return normalized_outputs
-./base_client.py:    def graph_outputs(
-./remote_client.py:class GraphOutputs(BaseModel):
-./remote_client.py:    outputs: List[GraphOutputMetadata]
-./remote_client.py:    def graph_outputs(
-./remote_client.py:            f"namespaces/{self.namespace}/compute_graphs/{graph}/invocations/{invocation_id}/outputs",
-./remote_client.py:        graph_outputs = GraphOutputs(**response.json())
-./remote_client.py:        outputs = []
-./remote_client.py:        for output in graph_outputs.outputs:
-./remote_client.py:                outputs.append(output)
-./remote_client.py:        return outputs
-./executor/task_reporter.py:        fn_outputs = []
-./executor/task_reporter.py:            f"[bold]task-reporter[/bold] uploading output of size: {len(completed_task.outputs)}"
-./executor/task_reporter.py:        for output in completed_task.outputs: 
-./executor/task_reporter.py:            fn_outputs.append(
-./executor/task_reporter.py:                ("node_outputs", (nanoid.generate(), io.BytesIO(output_bytes)))
-./executor/task_reporter.py:            fn_outputs.append(
-./executor/task_reporter.py:            fn_outputs.append(
-./executor/task_reporter.py:            fn_outputs.append(
-./executor/task_reporter.py:        if fn_outputs and len(fn_outputs) > 0:
-./executor/task_reporter.py:            kwargs["files"] = fn_outputs
-./executor/task_store.py:    outputs: List[IndexifyData]
-./executor/task_store.py:                    task_id=task_id, task_outcome="failed", outputs=[]
-./executor/downloader.py:            url = f"{self.base_url}/internal/fn_outputs/{task.input_key}"
-./executor/function_worker.py:    outputs: Union[List[IndexifyData], RouterOutput]
-./executor/function_worker.py:            indexify_data=result.outputs,
-./executor/function_worker.py:    return FunctionOutput(outputs=output, reducer=is_reducer)
-./executor/agent.py:                        f"Outputs: {len(task_outcome.outputs)}",
-./executor/agent.py:                            outputs=[],
-./executor/agent.py:                            outputs=[],
-./executor/agent.py:                            outputs=[],
-./executor/agent.py:                        outputs: FunctionWorkerOutput = await async_task
-./executor/agent.py:                            outputs if isinstance(outputs, RouterOutput) else None
-./executor/agent.py:                        if outputs.exception:
-./executor/agent.py:                            fn_outputs = []
-./executor/agent.py:                            fn_outputs = (
-./executor/agent.py:                                outputs.indexify_data if not isinstance(outputs, RouterOutput) else []
-./executor/agent.py:                            outputs=fn_outputs,
-./executor/agent.py:                            errors=outputs.exception,
-./executor/agent.py:                            stdout=outputs.stdout,
-./executor/agent.py:                            stderr=outputs.stderr,
-./executor/agent.py:                            reducer=outputs.reducer,
-./executor/agent.py:                            outputs=[],
-./local_client.py:# Holds the outputs of a
-./local_client.py:    outputs: Dict[str, List[IndexifyData]]
-./local_client.py:        outputs = defaultdict(list)
-./local_client.py:        self._results[input.id] = outputs
-./local_client.py:        self._run(g, input, outputs)
-./local_client.py:        outputs: Dict[str, List[bytes]],
-./local_client.py:                    f"ran {node_name}: num outputs: {len(cached_output_bytes)} (cache hit)"
-./local_client.py:                function_outputs: List[IndexifyData] = []
-./local_client.py:                    function_outputs.append(output)
-./local_client.py:                    outputs[node_name].append(output)
-./local_client.py:                function_outputs: List[IndexifyData] = g.invoke_fn_ser(
-./local_client.py:                print(f"ran {node_name}: num outputs: {len(function_outputs)}")
-./local_client.py:                    self._accumulators[node_name] = function_outputs[-1].model_copy()
-./local_client.py:                    outputs[node_name] = []
-./local_client.py:                outputs[node_name].extend(function_outputs)
-./local_client.py:                function_outputs_bytes: List[bytes] = [
-./local_client.py:                    for function_output in function_outputs
-./local_client.py:                    function_outputs_bytes,
-./local_client.py:                    for output in function_outputs:
-./local_client.py:                for output in function_outputs:
-./local_client.py:    def graph_outputs(

indexify-0.2.3/indexify/local_client.py

@@ -1,183 +0,0 @@
-from collections import defaultdict
-from queue import deque
-from typing import Any, Dict, List, Optional, Type, Union
-
-from nanoid import generate
-from pydantic import BaseModel, Json
-from rich import print
-
-from indexify.base_client import IndexifyClient
-from indexify.functions_sdk.data_objects import (
-    File,
-    IndexifyData,
-    RouterOutput,
-)
-from indexify.functions_sdk.graph import Graph
-from indexify.functions_sdk.local_cache import CacheAwareFunctionWrapper
-from indexify.functions_sdk.object_serializer import get_serializer
-
-
-# Holds the outputs of a
-class ContentTree(BaseModel):
-    id: str
-    outputs: Dict[str, List[IndexifyData]]
-
-
-class LocalClient(IndexifyClient):
-    def __init__(self, cache_dir: str = "./indexify_local_runner_cache"):
-        self._cache_dir = cache_dir
-        self._graphs: Dict[str, Graph] = {}
-        self._results: Dict[str, Dict[str, List[IndexifyData]]] = {}
-        self._cache = CacheAwareFunctionWrapper(self._cache_dir)
-        self._accumulators: Dict[str, Dict[str, IndexifyData]] = {}
-
-    def register_compute_graph(self, graph: Graph):
-        self._graphs[graph.name] = graph
-
-    def run_from_serialized_code(self, code: bytes, **kwargs):
-        g = Graph.deserialize(graph=code)
-        self.run(g, **kwargs)
-
-    def run(self, g: Graph, **kwargs):
-        serializer = get_serializer(
-            g.get_function(g._start_node).indexify_function.payload_encoder
-        )
-        input = IndexifyData(id=generate(), payload=serializer.serialize(kwargs))
-        print(f"[bold] Invoking {g._start_node}[/bold]")
-        outputs = defaultdict(list)
-        for k, v in g.get_accumulators().items():
-            serializer = get_serializer(
-                g.get_function(k).indexify_function.payload_encoder
-            )
-            self._accumulators[k] = IndexifyData(payload=serializer.serialize(v))
-        self._results[input.id] = outputs
-        self._run(g, input, outputs)
-        return input.id
-
-    def _run(
-        self,
-        g: Graph,
-        initial_input: bytes,
-        outputs: Dict[str, List[bytes]],
-    ):
-        queue = deque([(g._start_node, initial_input)])
-        while queue:
-            node_name, input = queue.popleft()
-            serializer = get_serializer(
-                g.get_function(node_name).indexify_function.payload_encoder
-            )
-            input_bytes = serializer.serialize(input)
-            cached_output_bytes: Optional[bytes] = self._cache.get(
-                g.name, node_name, input_bytes
-            )
-            if cached_output_bytes is not None:
-                print(
-                    f"ran {node_name}: num outputs: {len(cached_output_bytes)} (cache hit)"
-                )
-                function_outputs: List[IndexifyData] = []
-                cached_output_list = serializer.deserialize_list(cached_output_bytes)
-                if self._accumulators.get(node_name, None) is not None:
-                    self._accumulators[node_name] = cached_output_list[-1].model_copy()
-                    outputs[node_name] = []
-                function_outputs.extend(cached_output_list)
-                outputs[node_name].extend(cached_output_list)
-            else:
-                function_outputs: List[IndexifyData] = g.invoke_fn_ser(
-                    node_name, input, self._accumulators.get(node_name, None)
-                )
-                print(f"ran {node_name}: num outputs: {len(function_outputs)}")
-                if self._accumulators.get(node_name, None) is not None:
-                    self._accumulators[node_name] = function_outputs[-1].model_copy()
-                    outputs[node_name] = []
-                outputs[node_name].extend(function_outputs)
-                function_outputs_bytes: List[bytes] = [
-                    serializer.serialize_list(function_outputs)
-                ]
-                self._cache.set(
-                    g.name,
-                    node_name,
-                    input_bytes,
-                    function_outputs_bytes,
-                )
-            if self._accumulators.get(node_name, None) is not None and queue:
-                print(
-                    f"accumulator not none for {node_name}, continuing, len queue: {len(queue)}"
-                )
-                continue
-
-            out_edges = g.edges.get(node_name, [])
-            # Figure out if there are any routers for this node
-            for i, edge in enumerate(out_edges):
-                if edge in g.routers:
-                    out_edges.remove(edge)
-                    for output in function_outputs:
-                        dynamic_edges = self._route(g, edge, output) or []
-                        for dynamic_edge in dynamic_edges.edges:
-                            if dynamic_edge in g.nodes:
-                                print(
-                                    f"[bold]dynamic router returned node: {dynamic_edge}[/bold]"
-                                )
-                                out_edges.append(dynamic_edge)
-            for out_edge in out_edges:
-                for output in function_outputs:
-                    queue.append((out_edge, output))
-
-    def _route(
-        self, g: Graph, node_name: str, input: IndexifyData
-    ) -> Optional[RouterOutput]:
-        return g.invoke_router(node_name, input)
-
-    def graphs(self) -> str:
-        return list(self._graphs.keys())
-
-    def namespaces(self) -> str:
-        return "local"
-
-    def create_namespace(self, namespace: str):
-        pass
-
-    def rerun_graph(self, graph: str):
-        return super().rerun_graph(graph)
-
-    def invoke_graph_with_object(
-        self, graph: str, block_until_done: bool = False, **kwargs
-    ) -> str:
-        graph: Graph = self._graphs[graph]
-        return self.run(graph, **kwargs)
-
-    def invoke_graph_with_file(
-        self,
-        graph: str,
-        path: str,
-        metadata: Optional[Dict[str, Json]] = None,
-        block_until_done: bool = False,
-    ) -> str:
-        graph = self._graphs[graph]
-        with open(path, "rb") as f:
-            data = f.read()
-            file = File(data=data, metadata=metadata).model_dump()
-        return self.run(graph, file=file)
-
-    def graph_outputs(
-        self,
-        graph: str,
-        invocation_id: str,
-        fn_name: str,
-    ) -> Union[Dict[str, List[Any]], List[Any]]:
-        if invocation_id not in self._results:
-            raise ValueError(f"no results found for graph {graph}")
-        if fn_name not in self._results[invocation_id]:
-            raise ValueError(f"no results found for fn {fn_name} on graph {graph}")
-        results = []
-        fn_model = self._graphs[graph].get_function(fn_name).get_output_model()
-        serializer = get_serializer(
-            self._graphs[graph].get_function(fn_name).indexify_function.payload_encoder
-        )
-        for result in self._results[invocation_id][fn_name]:
-            payload_dict = serializer.deserialize(result.payload)
-            if issubclass(fn_model, BaseModel) and isinstance(payload_dict, dict):
-                payload = fn_model.model_validate(payload_dict)
-            else:
-                payload = payload_dict
-            results.append(payload)
-        return results