Closes OPEN-5836 Tracing improvements to get Openlayer Assistant traced

whoseoyster · whoseoyster · commit 8ec1410b3b41 · 2024-03-18T22:03:47.000-07:00
diff --git a/openlayer/__init__.py b/openlayer/__init__.py
@@ -21,6 +21,7 @@
     project.status()
     project.push()
 """
+
 import os
 import shutil
 import tarfile
diff --git a/openlayer/api.py b/openlayer/api.py
@@ -19,6 +19,7 @@
     project_data = self.api.post_request(endpoint, body=payload)
 
 """
+
 import os
 import shutil
 from enum import Enum
@@ -174,9 +175,9 @@ def post_request(
         return self._api_request(
             method="POST",
             endpoint=endpoint,
-            headers=self._headers
-            if files is None
-            else self._headers_multipart_form_data,
+            headers=(
+                self._headers if files is None else self._headers_multipart_form_data
+            ),
             body=body,
             files=files,
             data=data,
@@ -188,9 +189,9 @@ def put_request(self, endpoint: str, body=None, files=None, data=None):
         return self._api_request(
             "PUT",
             endpoint,
-            headers=self._headers
-            if files is None
-            else self._headers_multipart_form_data,
+            headers=(
+                self._headers if files is None else self._headers_multipart_form_data
+            ),
             body=body,
             files=files,
             data=data,
diff --git a/openlayer/constants.py b/openlayer/constants.py
@@ -1,5 +1,6 @@
 """Module for storing constants used throughout the OpenLayer Python Client.
 """
+
 import os
 
 import marshmallow as ma
diff --git a/openlayer/exceptions.py b/openlayer/exceptions.py
@@ -5,6 +5,7 @@
   if project is None:
     raise errors.OpenlayerResourceNotFound(f"Project {project_id} does not exist")
 """
+
 from typing import Dict
 
 
diff --git a/openlayer/llm_monitors.py b/openlayer/llm_monitors.py
@@ -1,5 +1,6 @@
 """Module with classes for monitoring calls to LLMs."""
 
+import json
 import logging
 import time
 import warnings
@@ -149,17 +150,18 @@ def modified_create_chat_completion(*args, **kwargs) -> str:
                         output_data = output_content.strip()
                     elif output_function_call or output_tool_calls:
                         if output_function_call:
-                            function_call = dict(output_function_call)
+                            function_call = {
+                                "name": output_function_call.name,
+                                "arguments": json.loads(output_function_call.arguments),
+                            }
                         else:
-                            function_call = dict(output_tool_calls[0].function)
-                        metadata = {
-                            "function_call_name": function_call.get("name"),
-                            "function_call_arguments": function_call.get("arguments"),
-                        }
-                        output_data = str(function_call)
+                            function_call = {
+                                "name": output_tool_calls[0].name,
+                                "arguments": json.loads(output_function_call.arguments),
+                            }
+                        output_data = function_call
                     else:
                         output_data = None
-                    output_data = response.choices[0].message.content
                     cost = self.get_cost_estimate(
                         model=kwargs.get("model"),
                         num_input_tokens=response.usage.prompt_tokens,
@@ -181,7 +183,6 @@ def modified_create_chat_completion(*args, **kwargs) -> str:
                         model_parameters=kwargs.get("model_parameters"),
                         raw_output=response.model_dump(),
                         provider="OpenAI",
-                        metadata=metadata,
                     )
                 # pylint: disable=broad-except
                 except Exception as e:
diff --git a/openlayer/model_runners/prediction_jobs/classification_prediction_job.py b/openlayer/model_runners/prediction_jobs/classification_prediction_job.py
@@ -8,6 +8,7 @@
 Example usage:
     python classification_prediction_job.py --input /path/to/input.csv --output /path/to/output.csv
 """
+
 import argparse
 import logging
 
diff --git a/openlayer/model_runners/prediction_jobs/regression_prediction_job.py b/openlayer/model_runners/prediction_jobs/regression_prediction_job.py
@@ -8,6 +8,7 @@
 Example usage:
     python regression_prediction_job.py --input /path/to/input.csv --output /path/to/output.csv
 """
+
 import argparse
 import logging
 
diff --git a/openlayer/model_runners/tests/test_llm_runners.py b/openlayer/model_runners/tests/test_llm_runners.py
@@ -4,6 +4,7 @@
 
     pytest test_llm_runners.py
 """
+
 from typing import Dict
 
 import anthropic
diff --git a/openlayer/services/data_streamer.py b/openlayer/services/data_streamer.py
@@ -29,28 +29,67 @@ def __init__(
         openlayer_inference_pipeline_id: Optional[str] = None,
         publish: bool = False,
     ) -> None:
-        self.openlayer_api_key = openlayer_api_key or utils.get_env_variable(
+        self._openlayer_api_key = openlayer_api_key or utils.get_env_variable(
             "OPENLAYER_API_KEY"
         )
-        self.openlayer_project_name = openlayer_project_name or utils.get_env_variable(
+        self._openlayer_project_name = openlayer_project_name or utils.get_env_variable(
             "OPENLAYER_PROJECT_NAME"
         )
-        self.openlayer_inference_pipeline_name = (
+        self._openlayer_inference_pipeline_name = (
             openlayer_inference_pipeline_name
             or utils.get_env_variable("OPENLAYER_INFERENCE_PIPELINE_NAME")
             or "production"
         )
-        self.openlayer_inference_pipeline_id = (
+        self._openlayer_inference_pipeline_id = (
             openlayer_inference_pipeline_id
             or utils.get_env_variable("OPENLAYER_INFERENCE_PIPELINE_ID")
         )
         self.publish = publish
 
-        self._validate_attributes()
-
         # Lazy load the inference pipeline
         self.inference_pipeline = None
 
+    @property
+    def openlayer_api_key(self) -> Optional[str]:
+        """The Openlayer API key."""
+        return self._get_openlayer_attribute("_openlayer_api_key", "OPENLAYER_API_KEY")
+
+    @property
+    def openlayer_project_name(self) -> Optional[str]:
+        """The name of the project on Openlayer."""
+        return self._get_openlayer_attribute(
+            "_openlayer_project_name", "OPENLAYER_PROJECT_NAME"
+        )
+
+    @property
+    def openlayer_inference_pipeline_name(self) -> Optional[str]:
+        """The name of the inference pipeline on Openlayer."""
+        return self._get_openlayer_attribute(
+            "_openlayer_inference_pipeline_name", "OPENLAYER_INFERENCE_PIPELINE_NAME"
+        )
+
+    @property
+    def openlayer_inference_pipeline_id(self) -> Optional[str]:
+        """The id of the inference pipeline on Openlayer."""
+        return self._get_openlayer_attribute(
+            "_openlayer_inference_pipeline_id", "OPENLAYER_INFERENCE_PIPELINE_ID"
+        )
+
+    def _get_openlayer_attribute(
+        self, attribute_name: str, env_variable: str
+    ) -> Optional[str]:
+        """A helper method to fetch an Openlayer attribute value.
+
+        Args:
+            attribute_name: The name of the attribute in this class.
+            env_variable: The name of the environment variable to fetch.
+        """
+        attribute_value = getattr(self, attribute_name, None)
+        if not attribute_value:
+            attribute_value = utils.get_env_variable(env_variable)
+            setattr(self, attribute_name, attribute_value)
+        return attribute_value
+
     def _validate_attributes(self) -> None:
         """Granular validation of the arguments."""
         if self.publish:
@@ -97,6 +136,7 @@ def stream_data(self, data: Dict[str, any], config: Dict[str, any]) -> None:
             config: The configuration for the data stream.
         """
 
+        self._validate_attributes()
         self._check_inference_pipeline_ready()
         self.inference_pipeline.stream_data(stream_data=data, stream_config=config)
         logger.info("Data streamed to Openlayer.")
diff --git a/openlayer/tracing/steps.py b/openlayer/tracing/steps.py
@@ -3,6 +3,7 @@
 import time
 from typing import Any, Dict, Optional
 
+from .. import utils
 from . import enums
 
 
@@ -40,6 +41,7 @@ def add_nested_step(self, nested_step: "Step") -> None:
 
     def log(self, **kwargs: Any) -> None:
         """Logs step data."""
+        kwargs = utils.json_serialize(kwargs)
         for key, value in kwargs.items():
             if hasattr(self, key):
                 setattr(self, key, value)
diff --git a/openlayer/tracing/tracer.py b/openlayer/tracing/tracer.py
@@ -41,6 +41,7 @@ def create_step(
     new_step: steps.Step = steps.step_factory(
         step_type=step_type, name=name, inputs=inputs, output=output, metadata=metadata
     )
+    new_step.start_time = time.time()
 
     parent_step: Optional[steps.Step] = _current_step.get(None)
     is_root_step: bool = parent_step is None
@@ -59,6 +60,12 @@ def create_step(
     try:
         yield new_step
     finally:
+        if new_step.end_time is None:
+            new_step.end_time = time.time()
+        if new_step.latency is None:
+            latency = (new_step.end_time - new_step.start_time) * 1000  # in ms
+            new_step.latency = latency
+
         _current_step.reset(token)
         if is_root_step:
             logger.debug("Ending the trace...")
@@ -95,9 +102,7 @@ def add_openai_chat_completion_step_to_trace(**kwargs) -> None:
     with create_step(
         step_type=enums.StepType.CHAT_COMPLETION, name="OpenAI Chat Completion"
     ) as step:
-        step.log(
-            **kwargs,
-        )
+        step.log(**kwargs)
 
 
 # ----------------------------- Tracing decorator ---------------------------- #
diff --git a/openlayer/utils.py b/openlayer/utils.py
@@ -1,7 +1,9 @@
 """Series of helper functions and classes that are used throughout the
 OpenLayer Python client.
 """
+
 import io
+import json
 import logging
 import os
 import re
@@ -250,3 +252,24 @@ def write_yaml(dictionary: dict, filename: str):
     """
     with open(filename, "w", encoding="UTF-8") as stream:
         yaml.dump(dictionary, stream)
+
+
+def json_serialize(data):
+    """
+    Recursively attempts to convert data into JSON-serializable formats.
+    """
+    if isinstance(data, (str, int, float, bool, type(None))):
+        return data  # Already JSON-serializable
+    elif isinstance(data, dict):
+        return {k: json_serialize(v) for k, v in data.items()}
+    elif isinstance(data, list):
+        return [json_serialize(item) for item in data]
+    elif isinstance(data, tuple):
+        return tuple(json_serialize(item) for item in data)
+    else:
+        # Fallback: Convert to string if not serializable
+        try:
+            json.dumps(data)
+            return data  # Data was serializable
+        except TypeError:
+            return str(data)  # Not serializable, convert to string
diff --git a/openlayer/validators/base_validator.py b/openlayer/validators/base_validator.py
@@ -3,6 +3,7 @@
 The entry point for all validators. This is the interface that all validators
 must implement.
 """
+
 import logging
 from abc import ABC, abstractmethod
 from typing import List
diff --git a/openlayer/validators/baseline_model_validators.py b/openlayer/validators/baseline_model_validators.py
@@ -1,5 +1,6 @@
 """Implements the baseline model specific validation classes.
 """
+
 import logging
 import os
 from typing import Dict, List, Optional
diff --git a/openlayer/validators/commit_validators.py b/openlayer/validators/commit_validators.py
@@ -1,5 +1,6 @@
 """Implements the commit bundle specific validation class.
 """
+
 import logging
 from abc import ABC, abstractmethod
 from typing import Dict, List, Optional
@@ -66,16 +67,16 @@ def __init__(
             else {}
         )
         if "training" in self._bundle_resources:
-            self.training_dataset_config: Dict[
-                str, any
-            ] = utils.load_dataset_config_from_bundle(
-                bundle_path=bundle_path, label="training"
+            self.training_dataset_config: Dict[str, any] = (
+                utils.load_dataset_config_from_bundle(
+                    bundle_path=bundle_path, label="training"
+                )
             )
         elif "fine-tuning" in self._bundle_resources:
-            self.training_dataset_config: Dict[
-                str, any
-            ] = utils.load_dataset_config_from_bundle(
-                bundle_path=bundle_path, label="fine-tuning"
+            self.training_dataset_config: Dict[str, any] = (
+                utils.load_dataset_config_from_bundle(
+                    bundle_path=bundle_path, label="fine-tuning"
+                )
             )
         else:
             self.training_dataset_config = {}
diff --git a/openlayer/validators/inference_pipeline_validators.py b/openlayer/validators/inference_pipeline_validators.py
@@ -1,5 +1,6 @@
 """Implements the inference pipeline validation class.
 """
+
 from typing import Dict
 
 import marshmallow as ma
diff --git a/openlayer/validators/project_validators.py b/openlayer/validators/project_validators.py
@@ -1,5 +1,6 @@
 """Implements the project specific validation class.
 """
+
 from typing import Dict
 
 import marshmallow as ma
diff --git a/tests/test_openlayer.py b/tests/test_openlayer.py
@@ -1,6 +1,7 @@
 """
 Module with sample openlayer test
 """
+
 import openlayer
 
 
