Add tool hosting object parsing

Author: javieraspuru-ax

pyproject.toml

@@ -34,6 +34,7 @@ packages = [
 Repository = 'https://github.com/axiomatic-ai/axiomatic-python-sdk'
 
 [tool.poetry.dependencies]
+dill = ">=0.3.9"
 python = "^3.8"
 httpx = ">=0.21.2"
 pydantic = ">= 1.9.2"

requirements.txt

@@ -1,3 +1,4 @@
+dill==0.3.9
 httpx>=0.21.2
 pydantic>= 1.9.2
 pydantic-core==^2.18.2

src/axiomatic/client.py

@@ -1,4 +1,6 @@
 import base64
+import dill
+import json
 import requests
 import os
 import time
@@ -129,11 +131,41 @@ def tool_exec(self, tool: str, code: str, poll_interval: int = 3, debug: bool =
                         if debug:
                             print(f"status: {result.status}")
                         if result.status == "SUCCEEDED":
-                            return result.output
+                            output = json.loads(result.output)
+                            if not output['objects']:
+                                return result.output
+                            else:
+                                return {
+                                    "messages": output['messages'],
+                                    "objects": self._load_objects_from_base64(output['objects'])
+                                }
                         else:
                             return result.error_trace
             else:
                 return output.error_trace
 
+    def load(self, job_id: str, obj_key: str):
+        result = self._ax_client.tools.status(job_id=job_id)
+        if result.status == "SUCCEEDED":
+            output = json.loads(result.output)
+            if not output['objects']:
+                return result.output
+            else:
+                return self._load_objects_from_base64(output['objects'])[obj_key]
+        else:
+            return result.error_trace
+
+    def _load_objects_from_base64(self, encoded_dict):
+        loaded_objects = {}
+        for key, encoded_str in encoded_dict.items():
+            try:
+                decoded_bytes = base64.b64decode(encoded_str)
+                loaded_obj = dill.loads(decoded_bytes)
+                loaded_objects[key] = loaded_obj
+            except Exception as e:
+                print(f"Error loading object for key '{key}': {e}")
+                loaded_objects[key] = None
+        return loaded_objects
+
 
 class AsyncAxiomatic(AsyncBaseClient): ...