Completes OPEN-5973 Allow inference_id to be passed to the OpenAI monitors

Author: gustavocidornelas

openlayer/llm_monitors.py

@@ -136,6 +136,9 @@ def _get_modified_create_chat_completion(self) -> callable:
         def modified_create_chat_completion(*args, **kwargs) -> str:
             stream = kwargs.get("stream", False)
 
+            # Pop the reserved Openlayer kwargs
+            inference_id = kwargs.pop("inference_id", None)
+
             if not stream:
                 start_time = time.time()
                 response = self.create_chat_completion(*args, **kwargs)
@@ -169,21 +172,26 @@ def modified_create_chat_completion(*args, **kwargs) -> str:
                         num_input_tokens=response.usage.prompt_tokens,
                         num_output_tokens=response.usage.completion_tokens,
                     )
-
-                    self._add_to_trace(
-                        end_time=end_time,
-                        inputs={
+                    trace_args = {
+                        "end_time": end_time,
+                        "inputs": {
                             "prompt": kwargs["messages"],
                         },
-                        output=output_data,
-                        latency=(end_time - start_time) * 1000,
-                        tokens=response.usage.total_tokens,
-                        cost=cost,
-                        prompt_tokens=response.usage.prompt_tokens,
-                        completion_tokens=response.usage.completion_tokens,
-                        model=response.model,
-                        model_parameters=kwargs.get("model_parameters"),
-                        raw_output=response.model_dump(),
+                        "output": output_data,
+                        "latency": (end_time - start_time) * 1000,
+                        "tokens": response.usage.total_tokens,
+                        "cost": cost,
+                        "prompt_tokens": response.usage.prompt_tokens,
+                        "completion_tokens": response.usage.completion_tokens,
+                        "model": response.model,
+                        "model_parameters": kwargs.get("model_parameters"),
+                        "raw_output": response.model_dump(),
+                    }
+                    if inference_id:
+                        trace_args["id"] = str(inference_id)
+
+                    self._add_to_trace(
+                        **trace_args,
                     )
                 # pylint: disable=broad-except
                 except Exception as e:
@@ -267,28 +275,33 @@ def stream_chunks():
                                     else 0
                                 ),
                             )
-
-                            self._add_to_trace(
-                                end_time=end_time,
-                                inputs={
+                            trace_args = {
+                                "end_time": end_time,
+                                "inputs": {
                                     "prompt": kwargs["messages"],
                                 },
-                                output=output_data,
-                                latency=latency,
-                                tokens=num_of_completion_tokens,
-                                cost=completion_cost,
-                                prompt_tokens=None,
-                                completion_tokens=num_of_completion_tokens,
-                                model=kwargs.get("model"),
-                                model_parameters=kwargs.get("model_parameters"),
-                                raw_output=raw_outputs,
-                                metadata={
+                                "output": output_data,
+                                "latency": latency,
+                                "tokens": num_of_completion_tokens,
+                                "cost": completion_cost,
+                                "prompt_tokens": None,
+                                "completion_tokens": num_of_completion_tokens,
+                                "model": kwargs.get("model"),
+                                "model_parameters": kwargs.get("model_parameters"),
+                                "raw_output": raw_outputs,
+                                "metadata": {
                                     "timeToFirstToken": (
                                         (first_token_time - start_time) * 1000
                                         if first_token_time
                                         else None
                                     )
                                 },
+                            }
+                            if inference_id:
+                                trace_args["id"] = str(inference_id)
+
+                            self._add_to_trace(
+                                **trace_args,
                             )
                         # pylint: disable=broad-except
                         except Exception as e:

openlayer/tracing/steps.py

@@ -1,6 +1,7 @@
 """Module with the different Step classes that can be used in a trace."""
 
 import time
+import uuid
 from typing import Any, Dict, Optional
 
 from .. import utils
@@ -23,6 +24,7 @@ def __init__(
         metadata: Optional[Dict[str, any]] = None,
     ) -> None:
         self.name = name
+        self.id = uuid.uuid4()
         self.inputs = inputs
         self.output = output
         self.metadata = metadata or {}
@@ -50,6 +52,7 @@ def to_dict(self) -> Dict[str, Any]:
         """Dictionary representation of the Step."""
         return {
             "name": self.name,
+            "id": str(self.id),
             "type": self.step_type.value,
             "inputs": self.inputs,
             "output": self.output,

openlayer/tracing/tracer.py

@@ -83,6 +83,7 @@ def create_step(
                 "costColumnName": "cost",
                 "numOfTokenColumnName": "tokens",
                 "timestampColumnName": "inferenceTimestamp",
+                "inferenceIdColumnName": "inferenceId",
             }
             if isinstance(new_step, steps.ChatCompletionStep):
                 config.update(
@@ -195,6 +196,7 @@ def process_trace_for_upload(
     trace_data = {
         **input_variables,
         "inferenceTimestamp": root_step.start_time,
+        "inferenceId": str(root_step.id),
         "output": root_step.output,
         "groundTruth": root_step.ground_truth,
         "latency": root_step.latency,