Completes OPEN-5776 Introduce a StepType enum to the Python client

Author: gustavocidornelas

openlayer/llm_monitors.py

@@ -134,6 +134,7 @@ def modified_create_chat_completion(*args, **kwargs) -> str:
                         model=kwargs.get("model"),
                         model_parameters=kwargs.get("model_parameters"),
                         raw_output=response.model_dump(),
+                        provider="OpenAI",
                     )
                 # pylint: disable=broad-except
                 except Exception as e:
@@ -202,6 +203,7 @@ def stream_chunks():
                                         else None
                                     )
                                 },
+                                provider="OpenAI",
                             )
                         # pylint: disable=broad-except
                         except Exception as e:
@@ -248,6 +250,7 @@ def modified_create_completion(*args, **kwargs):
                         model=kwargs.get("model"),
                         model_parameters=kwargs.get("model_parameters"),
                         raw_output=response.model_dump(),
+                        provider="OpenAI",
                     )
             # pylint: disable=broad-except
             except Exception as e:
@@ -331,6 +334,7 @@ def monitor_thread_run(self, run: openai.types.beta.threads.run.Run) -> None:
                 output=prompt[-1]["content"],
                 **run_step_vars,
                 metadata=metadata,
+                provider="OpenAI",
             )
 
         # pylint: disable=broad-except

openlayer/tracing/enums.py

@@ -0,0 +1,6 @@
+import enum
+
+
+class StepType(enum.Enum):
+    USER_CALL = "user_call"
+    CHAT_COMPLETION = "chat_completion"

openlayer/tracing/steps.py

@@ -3,6 +3,8 @@
 import time
 from typing import Any, Dict, Optional
 
+from . import enums
+
 
 class Step:
     def __init__(
@@ -17,7 +19,7 @@ def __init__(
         self.output = output
         self.metadata = metadata
 
-        self.step_type = None
+        self.step_type: enums.StepType = None
         self.start_time = time.time()
         self.end_time = None
         self.ground_truth = None
@@ -39,7 +41,7 @@ def to_dict(self) -> Dict[str, Any]:
         """Dictionary representation of the Step."""
         return {
             "name": self.name,
-            "type": self.step_type,
+            "type": self.step_type.value,
             "inputs": self.inputs,
             "output": self.output,
             "groundTruth": self.ground_truth,
@@ -60,10 +62,10 @@ def __init__(
         metadata: Dict[str, any] = {},
     ) -> None:
         super().__init__(name=name, inputs=inputs, output=output, metadata=metadata)
-        self.step_type = "user_call"
+        self.step_type = enums.StepType.USER_CALL
 
 
-class OpenAIChatCompletionStep(Step):
+class ChatCompletionStep(Step):
     def __init__(
         self,
         name: str,
@@ -73,7 +75,8 @@ def __init__(
     ) -> None:
         super().__init__(name=name, inputs=inputs, output=output, metadata=metadata)
 
-        self.step_type = "openai_chat_completion"
+        self.step_type = enums.StepType.CHAT_COMPLETION
+        self.provider: str = None
         self.prompt_tokens: int = None
         self.completion_tokens: int = None
         self.tokens: int = None
@@ -83,10 +86,11 @@ def __init__(
         self.raw_output: str = None
 
     def to_dict(self) -> Dict[str, Any]:
-        """Dictionary representation of the OpenAIChatCompletionStep."""
+        """Dictionary representation of the ChatCompletionStep."""
         step_dict = super().to_dict()
         step_dict.update(
             {
+                "provider": self.provider,
                 "promptTokens": self.prompt_tokens,
                 "completionTokens": self.completion_tokens,
                 "tokens": self.tokens,
@@ -100,12 +104,12 @@ def to_dict(self) -> Dict[str, Any]:
 
 
 # ----------------------------- Factory function ----------------------------- #
-def step_factory(step_type: str, *args, **kwargs) -> Step:
+def step_factory(step_type: enums.StepType, *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.")
+    if step_type.value not in [item.value for item in enums.StepType]:
+        raise ValueError(f"Step type {step_type.value} not recognized.")
     step_type_mapping = {
-        "user_call": UserCallStep,
-        "openai_chat_completion": OpenAIChatCompletionStep,
+        enums.StepType.USER_CALL: UserCallStep,
+        enums.StepType.CHAT_COMPLETION: ChatCompletionStep,
     }
     return step_type_mapping[step_type](*args, **kwargs)

openlayer/tracing/tracer.py

@@ -9,7 +9,7 @@
 from typing import Any, Dict, Generator, List, Optional, Tuple
 
 from ..services import data_streamer
-from . import steps, traces
+from . import enums, steps, traces
 
 logger = logging.getLogger(__name__)
 
@@ -31,7 +31,7 @@
 @contextmanager
 def create_step(
     name: str,
-    step_type: str = "user_call",
+    step_type: enums.StepType = enums.StepType.USER_CALL,
     inputs: Optional[Any] = None,
     output: Optional[Any] = None,
     metadata: Dict[str, any] = {},
@@ -71,7 +71,7 @@ def create_step(
                 "costColumnName": "cost",
                 "numOfTokenColumnName": "tokens",
             }
-            if isinstance(new_step, steps.OpenAIChatCompletionStep):
+            if isinstance(new_step, steps.ChatCompletionStep):
                 config.update(
                     {
                         "prompt": new_step.inputs.get("prompt"),
@@ -92,7 +92,7 @@ def create_step(
 def add_openai_chat_completion_step_to_trace(**kwargs) -> None:
     """Adds an OpenAI chat completion step to the trace."""
     with create_step(
-        step_type="openai_chat_completion", name="chat_completion"
+        step_type=enums.StepType.CHAT_COMPLETION, name="OpenAI Chat Completion"
     ) as step:
         step.log(
             **kwargs,