chore(internal): codegen related update (#396)

Author: Stainless Bot

examples/tracing/groq/groq_tracing.ipynb

@@ -95,14 +95,11 @@
    "source": [
     "chat_completion = groq_client.chat.completions.create(\n",
     "    messages=[\n",
-    "        {\n",
-    "            \"role\": \"system\",\n",
-    "            \"content\": \"You are a helpful assistant.\"\n",
-    "        },\n",
+    "        {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n",
     "        {\n",
     "            \"role\": \"user\",\n",
     "            \"content\": \"Explain the importance of fast language models\",\n",
-    "        }\n",
+    "        },\n",
     "    ],\n",
     "    model=\"llama3-8b-8192\",\n",
     ")"

examples/tracing/mistral/mistral_tracing.ipynb

@@ -92,12 +92,12 @@
    "source": [
     "response = mistral_client.chat.complete(\n",
     "    model=\"mistral-large-latest\",\n",
-    "    messages = [\n",
+    "    messages=[\n",
     "        {\n",
     "            \"role\": \"user\",\n",
     "            \"content\": \"What is the best French cheese?\",\n",
     "        },\n",
-    "    ]\n",
+    "    ],\n",
     ")"
    ]
   },
@@ -109,14 +109,14 @@
    "outputs": [],
    "source": [
     "stream_response = mistral_client.chat.stream(\n",
-    "    model = \"mistral-large-latest\",\n",
-    "    messages = [\n",
+    "    model=\"mistral-large-latest\",\n",
+    "    messages=[\n",
     "        {\n",
     "            \"role\": \"user\",\n",
     "            \"content\": \"What's the meaning of life?\",\n",
     "        },\n",
-    "    ]\n",
-    ")    "
+    "    ],\n",
+    ")"
    ]
   },
   {

examples/tracing/ollama/ollama_tracing.ipynb

@@ -94,10 +94,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "chat = ChatOllama(\n",
-    "    model=\"llama3.1\",\n",
-    "    callbacks=[openlayer_handler]\n",
-    ")"
+    "chat = ChatOllama(model=\"llama3.1\", callbacks=[openlayer_handler])"
    ]
   },
   {

examples/tracing/vertex-ai/vertex_ai_tracing.ipynb

@@ -97,10 +97,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "chat = ChatVertexAI(\n",
-    "    model=\"gemini-1.5-flash-001\",\n",
-    "    callbacks=[openlayer_handler]\n",
-    ")"
+    "chat = ChatVertexAI(model=\"gemini-1.5-flash-001\", callbacks=[openlayer_handler])"
    ]
   },
   {

requirements-dev.lock

@@ -19,6 +19,9 @@ argcomplete==3.1.2
 certifi==2023.7.22
     # via httpcore
     # via httpx
+    # via requests
+charset-normalizer==3.4.0
+    # via requests
 colorlog==6.7.0
     # via nox
 dirty-equals==0.6.0
@@ -41,6 +44,7 @@ httpx==0.25.2
 idna==3.4
     # via anyio
     # via httpx
+    # via requests
 importlib-metadata==7.0.0
 iniconfig==2.0.0
     # via pytest
@@ -56,7 +60,9 @@ nodeenv==1.8.0
     # via pyright
 nox==2023.4.22
 numpy==1.26.4
+    # via openlayer
     # via pandas
+    # via pyarrow
 packaging==23.2
     # via nox
     # via pytest
@@ -66,13 +72,15 @@ platformdirs==3.11.0
     # via virtualenv
 pluggy==1.5.0
     # via pytest
+pyarrow==14.0.1
+    # via openlayer
 pydantic==2.9.2
     # via openlayer
 pydantic-core==2.23.4
     # via pydantic
 pygments==2.18.0
     # via rich
-pyright==1.1.380
+pyright==1.1.389
 pytest==8.3.3
     # via pytest-asyncio
 pytest-asyncio==0.24.0
@@ -82,6 +90,12 @@ python-dateutil==2.8.2
 pytz==2023.3.post1
     # via dirty-equals
     # via pandas
+pyyaml==6.0.2
+    # via openlayer
+requests==2.32.3
+    # via requests-toolbelt
+requests-toolbelt==1.0.0
+    # via openlayer
 respx==0.20.2
 rich==13.7.1
 ruff==0.6.9
@@ -97,14 +111,19 @@ time-machine==2.9.0
 tomli==2.0.2
     # via mypy
     # via pytest
+tqdm==4.67.1
+    # via openlayer
 typing-extensions==4.12.2
     # via anyio
     # via mypy
     # via openlayer
     # via pydantic
     # via pydantic-core
+    # via pyright
 tzdata==2024.1
     # via pandas
+urllib3==2.2.3
+    # via requests
 virtualenv==20.24.5
     # via nox
 zipp==3.17.0

requirements.lock

@@ -17,6 +17,9 @@ anyio==4.4.0
 certifi==2023.7.22
     # via httpcore
     # via httpx
+    # via requests
+charset-normalizer==3.4.0
+    # via requests
 distro==1.8.0
     # via openlayer
 exceptiongroup==1.2.2
@@ -30,10 +33,15 @@ httpx==0.25.2
 idna==3.4
     # via anyio
     # via httpx
+    # via requests
 numpy==1.26.4
+    # via openlayer
     # via pandas
+    # via pyarrow
 pandas==2.2.2
     # via openlayer
+pyarrow==14.0.1
+    # via openlayer
 pydantic==2.9.2
     # via openlayer
 pydantic-core==2.23.4
@@ -42,16 +50,26 @@ python-dateutil==2.9.0.post0
     # via pandas
 pytz==2024.1
     # via pandas
+pyyaml==6.0.2
+    # via openlayer
+requests==2.32.3
+    # via requests-toolbelt
+requests-toolbelt==1.0.0
+    # via openlayer
 six==1.16.0
     # via python-dateutil
 sniffio==1.3.0
     # via anyio
     # via httpx
     # via openlayer
+tqdm==4.67.1
+    # via openlayer
 typing-extensions==4.12.2
     # via anyio
     # via openlayer
     # via pydantic
     # via pydantic-core
 tzdata==2024.1
     # via pandas
+urllib3==2.2.3
+    # via requests

src/openlayer/lib/__init__.py

@@ -1,5 +1,4 @@
-"""Openlayer lib.
-"""
+"""Openlayer lib."""
 
 __all__ = [
     "trace",

src/openlayer/lib/core/metrics.py

@@ -100,8 +100,7 @@ def _parse_args(self) -> None:
             type=str,
             required=False,
             default="",
-            help="The name of the dataset to compute the metric on. Runs on all "
-            "datasets if not provided.",
+            help="The name of the dataset to compute the metric on. Runs on all " "datasets if not provided.",
         )
 
         # Parse the arguments
@@ -133,9 +132,7 @@ def _load_datasets(self) -> None:
             dataset_names = [dataset["name"] for dataset in datasets_list]
             if self.dataset_name:
                 if self.dataset_name not in dataset_names:
-                    raise ValueError(
-                        f"Dataset {self.dataset_name} not found in the openlayer.json."
-                    )
+                    raise ValueError(f"Dataset {self.dataset_name} not found in the openlayer.json.")
                 dataset_names = [self.dataset_name]
             output_directory = model["outputDirectory"]
             # Read the outputs directory for dataset folders. For each, load
@@ -152,11 +149,7 @@ def _load_datasets(self) -> None:
                     dataset_config = json.load(f)
                     # Merge with the dataset fields from the openlayer.json
                     dataset_dict = next(
-                        (
-                            item
-                            for item in datasets_list
-                            if item["name"] == dataset_folder
-                        ),
+                        (item for item in datasets_list if item["name"] == dataset_folder),
                         None,
                     )
                     dataset_config = {**dataset_dict, **dataset_config}
@@ -166,9 +159,7 @@ def _load_datasets(self) -> None:
                     dataset_df = pd.read_csv(os.path.join(dataset_path, "dataset.csv"))
                     data_format = "csv"
                 elif os.path.exists(os.path.join(dataset_path, "dataset.json")):
-                    dataset_df = pd.read_json(
-                        os.path.join(dataset_path, "dataset.json"), orient="records"
-                    )
+                    dataset_df = pd.read_json(os.path.join(dataset_path, "dataset.json"), orient="records")
                     data_format = "json"
                 else:
                     raise ValueError(f"No dataset found in {dataset_folder}.")
@@ -183,14 +174,10 @@ def _load_datasets(self) -> None:
                     )
                 )
         else:
-            raise ValueError(
-                "No model found in the openlayer.json file. Cannot compute metric."
-            )
+            raise ValueError("No model found in the openlayer.json file. Cannot compute metric.")
 
         if not datasets:
-            raise ValueError(
-                "No datasets found in the openlayer.json file. Cannot compute metric."
-            )
+            raise ValueError("No datasets found in the openlayer.json file. Cannot compute metric.")
 
         self.datasets = datasets
 
@@ -243,13 +230,8 @@ def compute(self, datasets: List[Dataset]) -> None:
         """Compute the metric on the model outputs."""
         for dataset in datasets:
             # Check if the metric has already been computed
-            if os.path.exists(
-                os.path.join(dataset.output_path, "metrics", f"{self.key}.json")
-            ):
-                print(
-                    f"Metric ({self.key}) already computed on {dataset.name}. "
-                    "Skipping."
-                )
+            if os.path.exists(os.path.join(dataset.output_path, "metrics", f"{self.key}.json")):
+                print(f"Metric ({self.key}) already computed on {dataset.name}. " "Skipping.")
                 continue
 
             try:
@@ -276,9 +258,7 @@ def compute_on_dataset(self, dataset: Dataset) -> MetricReturn:
         """Compute the metric on a specific dataset."""
         pass
 
-    def _write_metric_return_to_file(
-        self, metric_return: MetricReturn, output_dir: str
-    ) -> None:
+    def _write_metric_return_to_file(self, metric_return: MetricReturn, output_dir: str) -> None:
         """Write the metric return to a file."""
 
         # Create the directory if it doesn't exist
@@ -289,9 +269,7 @@ def _write_metric_return_to_file(
         # Convert the set to a list
         metric_return_dict["added_cols"] = list(metric_return.added_cols)
 
-        with open(
-            os.path.join(output_dir, f"{self.key}.json"), "w", encoding="utf-8"
-        ) as f:
+        with open(os.path.join(output_dir, f"{self.key}.json"), "w", encoding="utf-8") as f:
             json.dump(metric_return_dict, f, indent=4)
         print(f"Metric ({self.key}) value written to {output_dir}/{self.key}.json")
 

src/openlayer/lib/data/_upload.py

@@ -105,9 +105,7 @@ def upload_blob_s3(
                 fields = presigned_url_response.fields
                 fields["file"] = (object_name, f, "application/x-tar")
                 e = MultipartEncoder(fields=fields)
-                m = MultipartEncoderMonitor(
-                    e, lambda monitor: t.update(min(t.total, monitor.bytes_read) - t.n)
-                )
+                m = MultipartEncoderMonitor(e, lambda monitor: t.update(min(t.total, monitor.bytes_read) - t.n))
                 headers = {"Content-Type": m.content_type}
                 res = requests.post(
                     presigned_url_response.url,
@@ -118,9 +116,7 @@ def upload_blob_s3(
                 )
         return res
 
-    def upload_blob_gcs(
-        self, file_path: str, presigned_url_response: PresignedURLCreateResponse
-    ):
+    def upload_blob_gcs(self, file_path: str, presigned_url_response: PresignedURLCreateResponse):
         """Generic method to upload data to Google Cloud Storage and create the
         appropriate resource in the backend.
         """
@@ -141,9 +137,7 @@ def upload_blob_gcs(
                 )
         return res
 
-    def upload_blob_azure(
-        self, file_path: str, presigned_url_response: PresignedURLCreateResponse
-    ):
+    def upload_blob_azure(self, file_path: str, presigned_url_response: PresignedURLCreateResponse):
         """Generic method to upload data to Azure Blob Storage and create the
         appropriate resource in the backend.
         """
@@ -186,9 +180,7 @@ def upload_blob_local(
             with open(file_path, "rb") as f:
                 fields = {"file": (object_name, f, "application/x-tar")}
                 e = MultipartEncoder(fields=fields)
-                m = MultipartEncoderMonitor(
-                    e, lambda monitor: t.update(min(t.total, monitor.bytes_read) - t.n)
-                )
+                m = MultipartEncoderMonitor(e, lambda monitor: t.update(min(t.total, monitor.bytes_read) - t.n))
                 headers = {"Content-Type": m.content_type}
                 res = requests.post(
                     presigned_url_response.url,