Add OpenAIChatCompletion step to tracing

Author: gustavocidornelas

openlayer/llm_monitors.py

@@ -9,6 +9,7 @@
 
 from . import constants, utils
 from .services import data_streamer
+from .tracing import tracer
 
 logger = logging.getLogger(__name__)
 
@@ -182,37 +183,36 @@ def modified_create_chat_completion(*args, **kwargs) -> str:
             if not stream:
                 start_time = time.time()
                 response = self.create_chat_completion(*args, **kwargs)
-                latency = (time.time() - start_time) * 1000
+                end_time = time.time()
+                latency = (end_time - start_time) * 1000
 
                 try:
                     # Extract data
-                    prompt, input_data = self.format_input(kwargs["messages"])
                     output_data = response.choices[0].message.content.strip()
                     num_of_tokens = response.usage.total_tokens
                     cost = self.get_cost_estimate(
                         model=kwargs.get("model"),
                         num_input_tokens=response.usage.prompt_tokens,
                         num_output_tokens=response.usage.completion_tokens,
                     )
-
-                    # Prepare config
-                    config = self.data_config.copy()
-                    config["prompt"] = prompt
-                    if not self.monitor_output_only:
-                        config.update({"inputVariableNames": list(input_data.keys())})
-
-                    self._append_row_to_df(
-                        input_data=input_data,
-                        output_data=output_data,
-                        num_of_tokens=num_of_tokens,
-                        latency=latency,
-                        cost=cost,
-                    )
-
-                    self.data_streamer.stream_data(
-                        data=self.df.tail(1).to_dict(orient="records"),
-                        config=config,
-                    )
+                    with tracer.create_step(
+                        step_type="openai_chat_completion", name="chat_completion"
+                    ) as step:
+                        step.update_data(
+                            end_time=end_time,
+                            inputs={
+                                "prompt": kwargs["messages"],
+                            },
+                            output=output_data,
+                            latency=latency,
+                            tokens=num_of_tokens,
+                            cost=cost,
+                            prompt_tokens=response.usage.prompt_tokens,
+                            completion_tokens=response.usage.completion_tokens,
+                            model=kwargs.get("model"),
+                            model_parameters=kwargs.get("model_parameters"),
+                            raw_output=response.model_dump(),
+                        )
                 # pylint: disable=broad-except
                 except Exception as e:
                     logger.error("Failed to monitor chat request. %s", e)

openlayer/tracing/steps.py

@@ -17,7 +17,7 @@ def __init__(
         self.output = output
         self.metadata = metadata
 
-        self.step_type = "user_call"
+        self.step_type = None
         self.start_time = time.time()
         self.end_time = None
         self.ground_truth = None
@@ -49,3 +49,61 @@ def to_dict(self) -> Dict[str, Any]:
             "startTime": self.start_time,
             "endTime": self.end_time,
         }
+
+
+class UserCallStep(Step):
+    def __init__(
+        self,
+        name: str,
+        inputs: Optional[Any] = None,
+        output: Optional[Any] = None,
+        metadata: Dict[str, any] = {},
+    ) -> None:
+        super().__init__(name=name, inputs=inputs, output=output, metadata=metadata)
+        self.step_type = "user_call"
+
+
+class OpenAIChatCompletionStep(Step):
+    def __init__(
+        self,
+        name: str,
+        inputs: Optional[Any] = None,
+        output: Optional[Any] = None,
+        metadata: Dict[str, any] = {},
+    ) -> None:
+        super().__init__(name=name, inputs=inputs, output=output, metadata=metadata)
+
+        self.step_type = "openai_chat_completion"
+        self.prompt_tokens: int = None
+        self.completion_tokens: int = None
+        self.cost: float = None
+        self.model: str = None
+        self.model_parameters: Dict[str, Any] = None
+        self.raw_output: str = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Dictionary representation of the OpenAIChatCompletionStep."""
+        step_dict = super().to_dict()
+        step_dict.update(
+            {
+                "promptTokens": self.prompt_tokens,
+                "completionTokens": self.completion_tokens,
+                "cost": self.cost,
+                "model": self.model,
+                "modelParameters": self.model_parameters,
+                "rawOutput": self.raw_output,
+            }
+        )
+        return step_dict
+
+
+# ----------------------------- Factory function ----------------------------- #
+def step_factory(step_type: str, *args, **kwargs) -> Step:
+    """Factory function to create a step based on the step_type."""
+    if step_type not in ["user_call", "openai_chat_completion"]:
+        raise ValueError(f"Step type {step_type} not recognized.")
+    step_type_mapping = {
+        "user_call": UserCallStep,
+        "openai_chat_completion": OpenAIChatCompletionStep,
+    }
+    return step_type_mapping[step_type](*args, **kwargs)

openlayer/tracing/tracer.py

@@ -17,13 +17,15 @@
 @contextmanager
 def create_step(
     name: str,
+    step_type: str = "user_call",
     inputs: Optional[Any] = None,
     output: Optional[Any] = None,
     metadata: Dict[str, any] = {},
 ) -> Generator[steps.Step, None, None]:
     """Starts a trace and yields a Step object."""
-    new_step = steps.Step(name=name, inputs=inputs, output=output, metadata=metadata)
-
+    new_step = steps.step_factory(
+        step_type=step_type, name=name, inputs=inputs, output=output, metadata=metadata
+    )
     parent_step = _current_step.get(None)
     is_root_step = parent_step is None
 
@@ -65,7 +67,12 @@ def wrapper(*func_args, **func_kwargs):
                 output = func(*func_args, **func_kwargs)
                 end_time = time.time()
                 latency = (end_time - step.start_time) * 1000  # in ms
-                inputs = func_signature.bind(*func_args, **func_kwargs).arguments
+
+                bound = func_signature.bind(*func_args, **func_kwargs)
+                bound.apply_defaults()
+                inputs = dict(bound.arguments)
+                inputs.pop("self", None)
+                inputs.pop("cls", None)
 
                 step.update_data(
                     inputs=inputs,