diff --git a/basalt/_version.py b/basalt/_version.py index d31c31e..493f741 100644 --- a/basalt/_version.py +++ b/basalt/_version.py @@ -1 +1 @@ -__version__ = "0.2.3" +__version__ = "0.3.0" diff --git a/basalt/endpoints/monitor/send_trace.py b/basalt/endpoints/monitor/send_trace.py index 456831d..ab0999b 100644 --- a/basalt/endpoints/monitor/send_trace.py +++ b/basalt/endpoints/monitor/send_trace.py @@ -54,7 +54,8 @@ def prepare_request(self, dto: Optional[Input] = None) -> Dict[str, Any]: "input": dict_log["input"] if "input" in dict_log else None, "output": dict_log["output"] if "output" in dict_log else None, "prompt": dict_log["prompt"] if "prompt" in dict_log else None, - "evaluators": dict_log["evaluators"] if "evaluators" in dict_log else None + "evaluators": dict_log["evaluators"] if "evaluators" in dict_log else None, + "idealOutput": dict_log["ideal_output"] if "ideal_output" in dict_log else None } logs.append(log_data) @@ -87,6 +88,11 @@ def prepare_request(self, dto: Optional[Input] = None) -> Dict[str, Any]: else: processed_log["parentId"] = None + # Rename ideal output + if "ideal_output" in processed_log: + processed_log["idealOutput"] = processed_log["ideal_output"] + del processed_log["ideal_output"] + processed_logs.append(processed_log) # Create the request body @@ -96,6 +102,7 @@ def prepare_request(self, dto: Optional[Input] = None) -> Dict[str, Any]: "experiment": {"id": trace_data.get("experiment", {}).id} if trace_data.get("experiment") else None, "input": trace_data.get("input"), "output": trace_data.get("output"), + "idealOutput": trace_data.get("ideal_output"), "metadata": trace_data.get("metadata"), "organization": trace_data.get("organization"), "user": trace_data.get("user"), diff --git a/basalt/objects/base_log.py b/basalt/objects/base_log.py index a856512..d6f8878 100644 --- a/basalt/objects/base_log.py +++ b/basalt/objects/base_log.py @@ -21,6 +21,7 @@ def __init__(self, params: BaseLogParams): self._trace = params.get("trace") self._parent = params.get("parent") self._evaluators = params.get("evaluators") + self._ideal_output = params.get("ideal_output") # Add to trace's logs list if trace exists if self._trace: @@ -51,6 +52,11 @@ def name(self) -> str: """Get the log name.""" return self._name + @property + def ideal_output(self) -> Optional[str]: + """Get the ideal output.""" + return self._ideal_output + @property def start_time(self) -> datetime: """Get the start time.""" @@ -98,6 +104,11 @@ def add_evaluator(self, evaluator: Evaluator) -> 'BaseLog': self._evaluators.append(evaluator) return self + def set_ideal_output(self, ideal_output: str) -> 'BaseLog': + """Set the ideal output.""" + self._ideal_output = ideal_output + return self + def update(self, params: Dict[str, Any]) -> 'BaseLog': """Update the log.""" self._name = params.get("name", self._name) @@ -122,6 +133,7 @@ def to_dict(self) -> Dict[str, Any]: "id": self._id, "type": self._type, "name": self._name, + "ideal_output": self._ideal_output, "start_time": self._start_time, "end_time": self._end_time, "metadata": self._metadata, diff --git a/basalt/objects/trace.py b/basalt/objects/trace.py index 5c932ae..cab5bef 100644 --- a/basalt/objects/trace.py +++ b/basalt/objects/trace.py @@ -19,6 +19,7 @@ def __init__(self, feature_slug: str, params: TraceParams, flusher: 'Flusher', l self._input = params.get("input") self._output = params.get("output") + self._ideal_output = params.get("ideal_output") self._name = params.get("name") self._start_time = params.get("start_time", datetime.now()) self._end_time = params.get("end_time") @@ -61,6 +62,11 @@ def output(self) -> Optional[str]: """Get the trace output.""" return self._output + @property + def ideal_output(self) -> Optional[str]: + """Get the trace ideal_output.""" + return self._ideal_output + @property def start_time(self) -> datetime: """Get the start time.""" @@ -132,6 +138,11 @@ def start(self, input: Optional[str] = None) -> 'Trace': self._start_time = datetime.now() return self + def set_ideal_output(self, ideal_output: str) -> 'Trace': + """Sets the ideal output for the trace.""" + self._ideal_output = ideal_output + return self + def identify(self, params: Dict[str, Any]) -> 'Trace': """ Set identification information for the trace. @@ -319,6 +330,7 @@ def to_dict(self) -> Dict[str, Any]: "feature_slug": self._feature_slug, "input": self._input, "output": self._output, + "ideal_output": self._ideal_output, "name": self._name, "start_time": self._start_time, "end_time": self._end_time, diff --git a/basalt/ressources/monitor/base_log_types.py b/basalt/ressources/monitor/base_log_types.py index 95e41c2..024302d 100644 --- a/basalt/ressources/monitor/base_log_types.py +++ b/basalt/ressources/monitor/base_log_types.py @@ -29,6 +29,7 @@ class BaseLogParams: evaluators: The evaluators to attach to the log. """ name: str + ideal_output: Optional[str] = None start_time: Optional[Union[datetime, str]] = None end_time: Optional[Union[datetime, str]] = None metadata: Optional[Dict[str, Any]] = None diff --git a/basalt/ressources/monitor/log_types.py b/basalt/ressources/monitor/log_types.py index 652559c..90caf8c 100644 --- a/basalt/ressources/monitor/log_types.py +++ b/basalt/ressources/monitor/log_types.py @@ -67,6 +67,7 @@ class Log(BaseLog): """ input: Optional[str] = None output: Optional[str] = None + ideal_output: Optional[str] = None def start(self, input: Optional[str] = None) -> 'Log': """Marks the log as started and sets the input if provided. diff --git a/basalt/ressources/monitor/trace_types.py b/basalt/ressources/monitor/trace_types.py index 41cdfdc..6d659e0 100644 --- a/basalt/ressources/monitor/trace_types.py +++ b/basalt/ressources/monitor/trace_types.py @@ -28,6 +28,7 @@ class TraceParams: name: Optional[str] = None input: Optional[str] = None output: Optional[str] = None + ideal_output: Optional[str] = None start_time: Optional[datetime] = None end_time: Optional[datetime] = None user: Optional[User] = None @@ -97,6 +98,10 @@ def start(self, input: Optional[str] = None) -> 'Trace': ``` """ ... + + def set_ideal_output(self, ideal_output: str) -> 'Trace': + """Sets the ideal output for the trace.""" + ... def set_metadata(self, metadata: Dict[str, Any]) -> 'Trace': """Sets or updates the metadata for this trace. diff --git a/basalt/sdk/monitorsdk.py b/basalt/sdk/monitorsdk.py index dd97189..fd0f3e9 100644 --- a/basalt/sdk/monitorsdk.py +++ b/basalt/sdk/monitorsdk.py @@ -246,6 +246,7 @@ def _create_trace( "input": params.input, "output": params.output, "name": params.name, + "ideal_output": params.ideal_output, "start_time": params.start_time, "end_time": params.end_time, "user": params.user, diff --git a/basalt/utils/flusher.py b/basalt/utils/flusher.py index e25e23d..ddcd730 100644 --- a/basalt/utils/flusher.py +++ b/basalt/utils/flusher.py @@ -33,6 +33,7 @@ def _trace_to_dict(self, trace: 'Trace') -> Dict[str, Any]: "feature_slug": trace.feature_slug, "input": trace.input, "output": output, + "ideal_output": trace.ideal_output, "name": trace._name, "start_time": trace.start_time.isoformat() if trace.start_time else None, "end_time": trace.end_time.isoformat() if trace.end_time else None, @@ -63,6 +64,7 @@ def _log_to_dict(self, log: Any) -> Dict[str, Any]: base_dict = { "id": log.id, "type": log.type, + "ideal_output": log.ideal_output, "name": log._name, "input": log.input, "output": output,