# ------------------------------------
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
# ------------------------------------
# mypy: disable-error-code="union-attr,arg-type,misc,return-value,assignment,func-returns-value"
# pylint: disable=R,redefined-outer-name,bare-except,unspecified-encoding
import os
import json
import inspect
import traceback
import importlib
import contextlib
from pathlib import Path
from numbers import Number
from datetime import datetime
from functools import wraps, partial
from typing import Any, Callable, Dict, Iterator, List, Union
# clean up key value pairs for sensitive values
def sanitize(key: str, value: Any) -> Any:
if isinstance(value, str) and any([s in key.lower() for s in ["key", "token", "secret", "password", "credential"]]):
return len(str(value)) * "*"
if isinstance(value, dict):
return {k: sanitize(k, v) for k, v in value.items()}
return value
class Tracer:
_tracers: Dict[str, Callable[[str], Iterator[Callable[[str, Any], None]]]] = {}
@classmethod
def add(cls, name: str, tracer: Callable[[str], Iterator[Callable[[str, Any], None]]]) -> None:
cls._tracers[name] = tracer
@classmethod
def clear(cls) -> None:
cls._tracers = {}
@classmethod
@contextlib.contextmanager
def start(cls, name: str) -> Iterator[Callable[[str, Any], None]]:
with contextlib.ExitStack() as stack:
traces: List[Any] = [stack.enter_context(tracer(name)) for tracer in cls._tracers.values()] # type: ignore
yield lambda key, value: [ # type: ignore
# normalize and sanitize any trace values
trace(key, sanitize(key, to_dict(value)))
for trace in traces
]
def to_dict(obj: Any) -> Union[Dict[str, Any], List[Dict[str, Any]], str, Number, bool]:
# simple json types
if isinstance(obj, str) or isinstance(obj, Number) or isinstance(obj, bool):
return obj
# datetime
if isinstance(obj, datetime):
return obj.isoformat()
# safe Prompty obj serialization
if type(obj).__name__ == "Prompty":
return obj.to_safe_dict()
# safe PromptyStream obj serialization
if type(obj).__name__ == "PromptyStream":
return "PromptyStream"
if type(obj).__name__ == "AsyncPromptyStream":
return "AsyncPromptyStream"
# recursive list and dict
if isinstance(obj, List):
return [to_dict(item) for item in obj] # type: ignore
if isinstance(obj, Dict):
return {k: v if isinstance(v, str) else to_dict(v) for k, v in obj.items()}
if isinstance(obj, Path):
return str(obj)
# cast to string otherwise...
return str(obj)
def _name(func: Callable, args):
if hasattr(func, "__qualname__"):
signature = f"{func.__module__}.{func.__qualname__}"
else:
signature = f"{func.__module__}.{func.__name__}"
# core invoker gets special treatment prompty.invoker.Invoker
core_invoker = signature.startswith("prompty.invoker.Invoker.run")
if core_invoker:
name = type(args[0]).__name__
if signature.endswith("async"):
signature = f"{args[0].__module__}.{args[0].__class__.__name__}.invoke_async"
else:
signature = f"{args[0].__module__}.{args[0].__class__.__name__}.invoke"
else:
name = func.__name__
return name, signature
def _inputs(func: Callable, args, kwargs) -> dict:
ba = inspect.signature(func).bind(*args, **kwargs)
ba.apply_defaults()
inputs = {k: to_dict(v) for k, v in ba.arguments.items() if k != "self"}
return inputs
def _results(result: Any) -> Union[Dict, List[Dict], str, Number, bool]:
return to_dict(result) if result is not None else "None"
def _trace_sync(func: Union[Callable, None] = None, **okwargs: Any) -> Callable:
@wraps(func) # type: ignore
def wrapper(*args, **kwargs):
name, signature = _name(func, args) # type: ignore
with Tracer.start(name) as trace:
trace("signature", signature)
# support arbitrary keyword
# arguments for trace decorator
for k, v in okwargs.items():
trace(k, to_dict(v))
inputs = _inputs(func, args, kwargs) # type: ignore
trace("inputs", inputs)
try:
result = func(*args, **kwargs) # type: ignore
trace("result", _results(result))
except Exception as e:
trace(
"result",
{
"exception": {
"type": type(e),
"traceback": (traceback.format_tb(tb=e.__traceback__) if e.__traceback__ else None),
"message": str(e),
"args": to_dict(e.args),
}
},
)
raise e
return result
return wrapper
def _trace_async(func: Union[Callable, None] = None, **okwargs: Any) -> Callable:
@wraps(func) # type: ignore
async def wrapper(*args, **kwargs):
name, signature = _name(func, args) # type: ignore
with Tracer.start(name) as trace:
trace("signature", signature)
# support arbitrary keyword
# arguments for trace decorator
for k, v in okwargs.items():
trace(k, to_dict(v))
inputs = _inputs(func, args, kwargs) # type: ignore
trace("inputs", inputs)
try:
result = await func(*args, **kwargs) # type: ignore
trace("result", _results(result))
except Exception as e:
trace(
"result",
{
"exception": {
"type": type(e),
"traceback": (traceback.format_tb(tb=e.__traceback__) if e.__traceback__ else None),
"message": str(e),
"args": to_dict(e.args),
}
},
)
raise e
return result
return wrapper
def trace(func: Union[Callable, None] = None, **kwargs: Any) -> Callable:
if func is None:
return partial(trace, **kwargs)
wrapped_method = _trace_async if inspect.iscoroutinefunction(func) else _trace_sync
return wrapped_method(func, **kwargs)
class PromptyTracer:
def __init__(self, output_dir: Union[str, None] = None) -> None:
if output_dir:
self.output = Path(output_dir).resolve().absolute()
else:
self.output = Path(Path(os.getcwd()) / ".runs").resolve().absolute()
if not self.output.exists():
self.output.mkdir(parents=True, exist_ok=True)
self.stack: List[Dict[str, Any]] = []
@contextlib.contextmanager
def tracer(self, name: str) -> Iterator[Callable[[str, Any], None]]:
try:
self.stack.append({"name": name})
frame = self.stack[-1]
frame["__time"] = {
"start": datetime.now(),
}
def add(key: str, value: Any) -> None:
if key not in frame:
frame[key] = value
# multiple values creates list
else:
if isinstance(frame[key], list):
frame[key].append(value)
else:
frame[key] = [frame[key], value]
yield add
finally:
frame = self.stack.pop()
start: datetime = frame["__time"]["start"]
end: datetime = datetime.now()
# add duration to frame
frame["__time"] = {
"start": start.strftime("%Y-%m-%dT%H:%M:%S.%f"),
"end": end.strftime("%Y-%m-%dT%H:%M:%S.%f"),
"duration": int((end - start).total_seconds() * 1000),
}
# hoist usage to parent frame
if "result" in frame and isinstance(frame["result"], dict):
if "usage" in frame["result"]:
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):
for result in frame["result"]:
if isinstance(result, dict) and "usage" in result and isinstance(result["usage"], dict):
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:
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:
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]) -> Dict[str, Any]:
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") # type: ignore
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]]:
try:
print(f"Starting {name}")
yield lambda key, value: print(f"{key}:\n{json.dumps(to_dict(value), indent=4)}")
finally:
print(f"Ending {name}")
