indexify 0.2.24__tar.gz → 0.2.26__tar.gz

{indexify-0.2.24 → indexify-0.2.26}/PKG-INFO

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

{indexify-0.2.24 → indexify-0.2.26}/indexify/executor/agent.py

@@ -418,7 +418,7 @@ class ExtractorAgent:
                         headers={"Content-Type": "application/json"},
                     ) as event_source:
                         if not event_source.response.is_success:
-                            resp = await event_source.response.aread().decode('utf-8')
+                            resp = await event_source.response.aread().decode("utf-8")
                             console.print(f"failed to register: {str(resp)}")
                             await asyncio.sleep(5)
                             continue

{indexify-0.2.24 → indexify-0.2.26}/indexify/executor/function_worker.py

@@ -17,6 +17,8 @@ from indexify.functions_sdk.indexify_functions import (
     GraphInvocationContext,
     IndexifyFunctionWrapper,
     RouterCallResult,
+    IndexifyRouter,
+    IndexifyFunction,
 )
 
 function_wrapper_map: Dict[str, IndexifyFunctionWrapper] = {}
@@ -169,7 +171,7 @@ def _run_function(
             fn = function_wrapper_map[key]
             if (
                 str(type(fn.indexify_function))
-                == "<class 'indexify.functions_sdk.indexify_functions.IndexifyRo'>"
+                == "<class 'indexify.functions_sdk.indexify_functions.IndexifyRouter'>"
             ):
                 router_call_result: RouterCallResult = fn.invoke_router(fn_name, input)
                 router_output = RouterOutput(edges=router_call_result.edges)

{indexify-0.2.24 → indexify-0.2.26}/indexify/functions_sdk/data_objects.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Optional, Union, Literal
+from typing import Any, Dict, List, Literal, Optional, Union
 
 from pydantic import BaseModel, Json
 

{indexify-0.2.24 → indexify-0.2.26}/indexify/functions_sdk/graph.py

@@ -30,10 +30,11 @@ from .graph_definition import (
 from .graph_validation import validate_node, validate_route
 from .indexify_functions import (
     FunctionCallResult,
+    GraphInvocationContext,
     IndexifyFunction,
     IndexifyFunctionWrapper,
     IndexifyRouter,
-    GraphInvocationContext,
+    RouterCallResult,
 )
 from .local_cache import CacheAwareFunctionWrapper
 from .object_serializer import get_serializer
@@ -73,17 +74,20 @@ class Graph:
         self.accumulator_zero_values: Dict[str, Any] = {}
 
         self.add_node(start_node)
+        if issubclass(start_node, IndexifyRouter):
+            self.routers[start_node.name] = []
         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]] = {}
+        self._local_graph_ctx: Optional[GraphInvocationContext] = None
 
     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])
+        return IndexifyFunctionWrapper(self.nodes[name], self._local_graph_ctx)
 
     def get_accumulators(self) -> Dict[str, Any]:
         return self.accumulator_zero_values
@@ -153,14 +157,17 @@ class Graph:
 
     def definition(self) -> ComputeGraphMetadata:
         start_node = self.nodes[self._start_node]
+        is_reducer = False
+        if hasattr(start_node, 'accumulate'):
+            is_reducer = start_node.accumulate is not None
         start_node = FunctionMetadata(
             name=start_node.name,
             fn_name=start_node.name,
             description=start_node.description,
-            reducer=start_node.accumulate is not None,
+            reducer=is_reducer,
             image_name=start_node.image._image_name,
             image_information=start_node.image.to_image_information(),
-            payload_encoder=start_node.encoder
+            payload_encoder=start_node.encoder,
         )
         metadata_edges = self.edges.copy()
         metadata_nodes = {}
@@ -205,7 +212,11 @@ class Graph:
     def run(self, block_until_done: bool = False, **kwargs) -> str:
         start_node = self.nodes[self._start_node]
         serializer = get_serializer(start_node.encoder)
-        input = IndexifyData(id=generate(), payload=serializer.serialize(kwargs), encoder=start_node.encoder)
+        input = IndexifyData(
+            id=generate(),
+            payload=serializer.serialize(kwargs),
+            encoder=start_node.encoder,
+        )
         print(f"[bold] Invoking {self._start_node}[/bold]")
         outputs = defaultdict(list)
         self._accumulator_values[input.id] = {}
@@ -216,37 +227,35 @@ class Graph:
                 k: IndexifyData(payload=serializer.serialize(v), encoder=node.encoder)
             }
         self._results[input.id] = outputs
-        enable_cache = kwargs.get("enable_cache", True)
         ctx = GraphInvocationContext(
             invocation_id=input.id,
             graph_name=self.name,
             graph_version="1",
             indexify_client=None,
         )
-        self._run(input, outputs, enable_cache, ctx)
+        self._local_graph_ctx = ctx
+        self._run(input, outputs)
         return input.id
 
     def _run(
         self,
         initial_input: IndexifyData,
         outputs: Dict[str, List[bytes]],
-        enable_cache: bool,
-        ctx: GraphInvocationContext,
     ):
         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]
-            function_outputs: FunctionCallResult = IndexifyFunctionWrapper(
-                node, context=ctx
-            ).invoke_fn_ser(node_name, input, accumulator_values.get(node_name, None))
-            if function_outputs.traceback_msg is not None:
-                print(function_outputs.traceback_msg)
-                import os
-
-                print("exiting local execution due to error")
-                os._exit(1)
+            function_outputs: Union[
+                FunctionCallResult, RouterCallResult
+            ] = self._invoke_fn(node_name, input)
+            self._log_local_exec_tracebacks(function_outputs)
+            if isinstance(function_outputs, RouterCallResult):
+                for edge in function_outputs.edges:
+                    if edge in self.nodes:
+                        queue.append((edge, input))
+                continue
+            out_edges = self.edges.get(node_name, [])
             fn_outputs = function_outputs.ser_outputs
             print(f"ran {node_name}: num outputs: {len(fn_outputs)}")
             if accumulator_values.get(node_name, None) is not None:
@@ -260,26 +269,37 @@ class Graph:
                 )
                 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 fn_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 fn_outputs:
                     queue.append((out_edge, output))
 
-    def _route(self, node_name: str, input: IndexifyData) -> Optional[RouterOutput]:
-        router = self.nodes[node_name]
-        return IndexifyFunctionWrapper(router).invoke_router(node_name, input)
+    def _invoke_fn(
+        self, node_name: str, input: IndexifyData
+    ) -> Optional[Union[RouterCallResult, FunctionCallResult]]:
+        node = self.nodes[node_name]
+        if node_name in self.routers:
+            result = IndexifyFunctionWrapper(node, self._local_graph_ctx).invoke_router(
+                node_name, input
+            )
+            for dynamic_edge in result.edges:
+                if dynamic_edge in self.nodes:
+                    print(f"[bold]dynamic router returned node: {dynamic_edge}[/bold]")
+            return result
+
+        acc_value = self._accumulator_values.get(node_name, None)
+        return IndexifyFunctionWrapper(
+            node, context=self._local_graph_ctx
+        ).invoke_fn_ser(node_name, input, acc_value)
+
+    def _log_local_exec_tracebacks(
+        self, results: Union[FunctionCallResult, RouterCallResult]
+    ):
+        if results.traceback_msg is not None:
+            print(results.traceback_msg)
+            import os
+
+            print("exiting local execution due to error")
+            os._exit(1)
 
     def output(
         self,

{indexify-0.2.24 → indexify-0.2.26}/indexify/functions_sdk/graph_validation.py

@@ -55,8 +55,7 @@ def validate_route(
                     raise Exception(
                         f"Unable to find {arg.name} in to_nodes {[node.name for node in to_nodes]}"
                     )
-
-    if (
+    elif (
         hasattr(return_annotation, "__origin__")
         and return_annotation.__origin__ is list
     ):

{indexify-0.2.24 → indexify-0.2.26}/indexify/functions_sdk/indexify_functions.py

@@ -48,37 +48,6 @@ class GraphInvocationContext(BaseModel):
         )
 
 
-def format_filtered_traceback(exc_info=None):
-    """
-    Format a traceback excluding indexify_functions.py lines.
-    Can be used in exception handlers to replace traceback.format_exc()
-    """
-    if exc_info is None:
-        exc_info = sys.exc_info()
-
-    # Get the full traceback as a string
-    full_traceback = traceback.format_exception(*exc_info)
-
-    # Filter out lines containing indexify_functions.py
-    filtered_lines = []
-    skip_next = False
-
-    for line in full_traceback:
-        if "indexify_functions.py" in line:
-            skip_next = True
-            continue
-        if skip_next:
-            if line.strip().startswith("File "):
-                skip_next = False
-            else:
-                continue
-        filtered_lines.append(line)
-
-    # Clean up any double blank lines that might have been created
-    cleaned = re.sub(r"\n\s*\n\s*\n", "\n\n", "".join(filtered_lines))
-    return cleaned
-
-
 def is_pydantic_model_from_annotation(type_annotation):
     # If it's a string representation
     if isinstance(type_annotation, str):
@@ -146,6 +115,11 @@ class IndexifyRouter:
         pass
 
 
+from inspect import signature
+
+from typing_extensions import get_type_hints
+
+
 def indexify_router(
     name: Optional[str] = None,
     description: Optional[str] = "",
@@ -154,28 +128,30 @@ def indexify_router(
     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.encoder = encoder
-        return IndexifyRo
+        # Get function signature using inspect.signature
+        fn_sig = signature(fn)
+        fn_hints = get_type_hints(fn)
+
+        # Create run method that preserves signature
+        def run(self, *args, **kwargs):
+            return fn(*args, **kwargs)
+
+        # Apply original signature and annotations to run method
+        run.__signature__ = fn_sig
+        run.__annotations__ = fn_hints
+
+        attrs = {
+            "name": name if name else fn.__name__,
+            "description": description
+            if description
+            else (fn.__doc__ or "").strip().replace("\n", ""),
+            "image": image,
+            "placement_constraints": placement_constraints,
+            "encoder": encoder,
+            "run": run,
+        }
+
+        return type("IndexifyRouter", (IndexifyRouter,), attrs)
 
     return construct
 
@@ -189,28 +165,31 @@ def indexify_function(
     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.encoder = encoder
-        return IndexifyFn
+        # Get function signature using inspect.signature
+        fn_sig = signature(fn)
+        fn_hints = get_type_hints(fn)
+
+        # Create run method that preserves signature
+        def run(self, *args, **kwargs):
+            return fn(*args, **kwargs)
+
+        # Apply original signature and annotations to run method
+        run.__signature__ = fn_sig
+        run.__annotations__ = fn_hints
+
+        attrs = {
+            "name": name if name else fn.__name__,
+            "description": description
+            if description
+            else (fn.__doc__ or "").strip().replace("\n", ""),
+            "image": image,
+            "placement_constraints": placement_constraints,
+            "accumulate": accumulate,
+            "encoder": encoder,
+            "run": run,
+        }
+
+        return type("IndexifyFunction", (IndexifyFunction,), attrs)
 
     return construct
 
@@ -266,7 +245,7 @@ class IndexifyFunctionWrapper:
         try:
             extracted_data = self.indexify_function.run(*args, **kwargs)
         except Exception as e:
-            return [], format_filtered_traceback()
+            return [], traceback.format_exc()
         if not isinstance(extracted_data, list) and extracted_data is not None:
             return [extracted_data.name], None
         edges = []
@@ -289,7 +268,7 @@ class IndexifyFunctionWrapper:
         try:
             extracted_data = self.indexify_function.run(*args, **kwargs)
         except Exception as e:
-            return [], format_filtered_traceback()
+            return [], traceback.format_exc()
         if extracted_data is None:
             return [], None
 
@@ -313,7 +292,11 @@ class IndexifyFunctionWrapper:
             )
         outputs, err = self.run_fn(input, acc=acc)
         ser_outputs = [
-            IndexifyData(payload=serializer.serialize(output), encoder=self.indexify_function.encoder) for output in outputs
+            IndexifyData(
+                payload=serializer.serialize(output),
+                encoder=self.indexify_function.encoder,
+            )
+            for output in outputs
         ]
         return FunctionCallResult(ser_outputs=ser_outputs, traceback_msg=err)
 
@@ -328,6 +311,7 @@ class IndexifyFunctionWrapper:
         serializer = get_serializer(encoder)
         return serializer.deserialize(payload)
 
+
 def get_ctx() -> GraphInvocationContext:
     frame = inspect.currentframe()
     caller_frame = frame.f_back.f_back

{indexify-0.2.24 → indexify-0.2.26}/indexify/functions_sdk/object_serializer.py

@@ -1,7 +1,7 @@
 from typing import Any, List
 
-import jsonpickle
 import cloudpickle
+import jsonpickle
 import msgpack
 from pydantic import BaseModel
 
@@ -17,6 +17,7 @@ def get_serializer(serializer_type: str) -> Any:
         return JsonSerializer()
     raise ValueError(f"Unknown serializer type: {serializer_type}")
 
+
 class JsonSerializer:
     @staticmethod
     def serialize(data: Any) -> str:

{indexify-0.2.24 → indexify-0.2.26}/indexify/http_client.py

@@ -16,6 +16,7 @@ from indexify.functions_sdk.graph import ComputeGraphMetadata, Graph
 from indexify.functions_sdk.indexify_functions import IndexifyFunction
 from indexify.settings import DEFAULT_SERVICE_URL
 
+
 class InvocationEventPayload(BaseModel):
     invocation_id: str
     fn_name: str
@@ -76,7 +77,9 @@ class IndexifyClient:
         self._fns: Dict[str, IndexifyFunction] = {}
         self._api_key = api_key
         if not self._api_key:
-            print("API key not provided. Trying to fetch from environment TENSORLAKE_API_KEY variable")
+            print(
+                "API key not provided. Trying to fetch from environment TENSORLAKE_API_KEY variable"
+            )
             self._api_key = os.getenv("TENSORLAKE_API_KEY")
 
     def _request(self, method: str, **kwargs) -> httpx.Response:
@@ -143,7 +146,7 @@ class IndexifyClient:
             verify=verify_option,
         )
         return client
-    
+
     def _add_api_key(self, kwargs):
         if self._api_key:
             kwargs["headers"] = {"Authorization": f"Bearer {self._api_key}"}
@@ -213,16 +216,16 @@ class IndexifyClient:
         self, compute_graph: str, invocation_id: str, key: str, value: Json
     ) -> None:
         response = self._post(
-                f"internal/namespaces/{self.namespace}/compute_graphs/{compute_graph}/invocations/{invocation_id}/ctx",
-                json={"key": key, "value": value},
-            )
+            f"internal/namespaces/{self.namespace}/compute_graphs/{compute_graph}/invocations/{invocation_id}/ctx",
+            json={"key": key, "value": value},
+        )
         response.raise_for_status()
 
     def get_state_key(self, compute_graph: str, invocation_id: str, key: str) -> Json:
         response = self._get(
-                f"internal/namespaces/{self.namespace}/compute_graphs/{compute_graph}/invocations/{invocation_id}/ctx",
-                json={"key": key},
-            )
+            f"internal/namespaces/{self.namespace}/compute_graphs/{compute_graph}/invocations/{invocation_id}/ctx",
+            json={"key": key},
+        )
         response.raise_for_status()
         return response.json().get("value")
 
@@ -270,7 +273,11 @@ class IndexifyClient:
     ) -> str:
         ser_input = cloudpickle.dumps(kwargs)
         params = {"block_until_finish": block_until_done}
-        kwargs = {"headers": {"Content-Type": "application/cbor"}, "data": ser_input, "params":params}
+        kwargs = {
+            "headers": {"Content-Type": "application/cbor"},
+            "data": ser_input,
+            "params": params,
+        }
         self._add_api_key(kwargs)
         with httpx.Client() as client:
             with connect_sse(

{indexify-0.2.24 → indexify-0.2.26}/pyproject.toml

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