Merge pull request #88 from microsoft/python

corrected tracing error when usage values are unexpected
microsoft · Sep 16, 2024 · 391ae60 · 391ae60
2 parents a621bc6 + 28442f0
commit 391ae60
Showing 1 changed file with 54 additions and 35 deletions.
diff --git a/runtime/prompty/prompty/tracer.py b/runtime/prompty/prompty/tracer.py
@@ -1,6 +1,7 @@
 import os
 import json
 import inspect
+import numbers
 import traceback
 import importlib
 import contextlib
@@ -139,7 +140,11 @@ def wrapper(*args, **kwargs):
                     {
                         "exception": {
                             "type": type(e),
-                            "traceback": traceback.format_tb(tb=e.__traceback__) if e.__traceback__ else None,
+                            "traceback": (
+                                traceback.format_tb(tb=e.__traceback__)
+                                if e.__traceback__
+                                else None
+                            ),
                             "message": str(e),
                             "args": to_dict(e.args),
                         }
@@ -179,7 +184,11 @@ async def wrapper(*args, **kwargs):
                     {
                         "exception": {
                             "type": type(e),
-                            "traceback": traceback.format_tb(tb=e.__traceback__) if e.__traceback__ else None,
+                            "traceback": (
+                                traceback.format_tb(tb=e.__traceback__)
+                                if e.__traceback__
+                                else None
+                            ),
                             "message": str(e),
                             "args": to_dict(e.args),
                         }
@@ -250,11 +259,10 @@ def add(key: str, value: Any) -> None:
             # hoist usage to parent frame
             if "result" in frame and isinstance(frame["result"], dict):
                 if "usage" in frame["result"]:
-                    if "__usage" in frame:
-                        for key, value in frame["result"]["usage"].items():
-                            frame["__usage"][key] += value
-                    else:
-                        frame["__usage"] = frame["result"]["usage"]
+                    frame["__usage"] = self.hoist_item(
+                        frame["result"]["usage"],
+                        frame["__usage"] if "__usage" in frame else {},
+                    )
 
             # streamed results may have usage as well
             if "result" in frame and isinstance(frame["result"], list):
@@ -264,48 +272,59 @@ def add(key: str, value: Any) -> None:
                         and "usage" in result
                         and isinstance(result["usage"], dict)
                     ):
-                        if "__usage" not in frame:
-                            frame["__usage"] = {}
-                        for key, value in result["usage"].items():
-                            if key not in frame["__usage"]:
-                                frame["__usage"][key] = value
-                            else:
-                                frame["__usage"][key] += value
+                        frame["__usage"] = self.hoist_item(
+                            result["usage"],
+                            frame["__usage"] if "__usage" in frame else {},
+                        )
 
             # add any usage frames from below
             if "__frames" in frame:
                 for child in frame["__frames"]:
                     if "__usage" in child:
-                        if "__usage" not in frame:
-                            frame["__usage"] = {}
-                        for key, value in child["__usage"].items():
-                            if key not in frame["__usage"]:
-                                frame["__usage"][key] = value
-                            else:
-                                frame["__usage"][key] += value
+                        frame["__usage"] = self.hoist_item(
+                            child["__usage"],
+                            frame["__usage"] if "__usage" in frame else {},
+                        )
 
             # if stack is empty, dump the frame
             if len(self.stack) == 0:
-                trace_file = (
-                    self.output
-                    / f"{frame['name']}.{datetime.now().strftime('%Y%m%d.%H%M%S')}.tracy"
-                )
-
-                v = importlib.metadata.version("prompty")
-                enriched_frame = {
-                    "runtime": "python",
-                    "version": v,
-                    "trace": frame,
-                }
-
-                with open(trace_file, "w") as f:
-                    json.dump(enriched_frame, f, indent=4)
+                self.write_trace(frame)
             # otherwise, append the frame to the parent
             else:
                 if "__frames" not in self.stack[-1]:
                     self.stack[-1]["__frames"] = []
                 self.stack[-1]["__frames"].append(frame)
 
+    def hoist_item(self, src: Dict[str, Any], cur: Dict[str, Any]) -> None:
+        for key, value in src.items():
+            if value is None or isinstance(value, list) or isinstance(value, dict):
+                continue
+            try:
+                if key not in cur:
+                    cur[key] = value
+                else:
+                    cur[key] += value
+            except:
+                continue
+
+        return cur
+
+    def write_trace(self, frame: Dict[str, Any]) -> None:
+        trace_file = (
+            self.output
+            / f"{frame['name']}.{datetime.now().strftime('%Y%m%d.%H%M%S')}.tracy"
+        )
+
+        v = importlib.metadata.version("prompty")
+        enriched_frame = {
+            "runtime": "python",
+            "version": v,
+            "trace": frame,
+        }
+
+        with open(trace_file, "w") as f:
+            json.dump(enriched_frame, f, indent=4)
+
 
 @contextlib.contextmanager
 def console_tracer(name: str) -> Iterator[Callable[[str, Any], None]]: