chore: apply formatting to custom files

gustavocidornelas · whoseoyster · commit 3414c66705e0 · 2024-05-28T18:28:33.000-10:00
diff --git a/src/openlayer/lib/core/base_model.py b/src/openlayer/lib/core/base_model.py
@@ -42,7 +42,9 @@ class OpenlayerModel(abc.ABC):
     def run_from_cli(self) -> None:
         """Run the model from the command line."""
         parser = argparse.ArgumentParser(description="Run data through a model.")
-        parser.add_argument("--dataset-path", type=str, required=True, help="Path to the dataset")
+        parser.add_argument(
+            "--dataset-path", type=str, required=True, help="Path to the dataset"
+        )
         parser.add_argument(
             "--output-dir",
             type=str,
@@ -83,7 +85,9 @@ def run_batch_from_df(self, df: pd.DataFrame) -> Tuple[pd.DataFrame, dict]:
             # Filter row_dict to only include keys that are valid parameters
             # for the 'run' method
             row_dict = row.to_dict()
-            filtered_kwargs = {k: v for k, v in row_dict.items() if k in run_signature.parameters}
+            filtered_kwargs = {
+                k: v for k, v in row_dict.items() if k in run_signature.parameters
+            }
 
             # Call the run method with filtered kwargs
             output = self.run(**filtered_kwargs)
diff --git a/src/openlayer/lib/integrations/langchain_callback.py b/src/openlayer/lib/integrations/langchain_callback.py
@@ -35,7 +35,9 @@ def __init__(self, **kwargs: Any) -> None:
         self.metatada: Dict[str, Any] = kwargs or {}
 
     # noqa arg002
-    def on_llm_start(self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any) -> Any:
+    def on_llm_start(
+        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
+    ) -> Any:
         """Run when LLM starts running."""
         pass
 
@@ -79,32 +81,45 @@ def on_llm_new_token(self, token: str, **kwargs: Any) -> Any:
         """Run on new LLM token. Only available when streaming is enabled."""
         pass
 
-    def on_llm_end(self, response: langchain_schema.LLMResult, **kwargs: Any) -> Any:  # noqa: ARG002, E501
+    def on_llm_end(
+        self, response: langchain_schema.LLMResult, **kwargs: Any  # noqa: ARG002, E501
+    ) -> Any:
         """Run when LLM ends running."""
         self.end_time = time.time()
         self.latency = (self.end_time - self.start_time) * 1000
 
         if response.llm_output and "token_usage" in response.llm_output:
-            self.prompt_tokens = response.llm_output["token_usage"].get("prompt_tokens", 0)
-            self.completion_tokens = response.llm_output["token_usage"].get("completion_tokens", 0)
+            self.prompt_tokens = response.llm_output["token_usage"].get(
+                "prompt_tokens", 0
+            )
+            self.completion_tokens = response.llm_output["token_usage"].get(
+                "completion_tokens", 0
+            )
             self.cost = self._get_cost_estimate(
                 num_input_tokens=self.prompt_tokens,
                 num_output_tokens=self.completion_tokens,
             )
-            self.total_tokens = response.llm_output["token_usage"].get("total_tokens", 0)
+            self.total_tokens = response.llm_output["token_usage"].get(
+                "total_tokens", 0
+            )
 
         for generations in response.generations:
             for generation in generations:
                 self.output += generation.text.replace("\n", " ")
 
         self._add_to_trace()
 
-    def _get_cost_estimate(self, num_input_tokens: int, num_output_tokens: int) -> float:
+    def _get_cost_estimate(
+        self, num_input_tokens: int, num_output_tokens: int
+    ) -> float:
         """Returns the cost estimate for a given model and number of tokens."""
         if self.model not in constants.OPENAI_COST_PER_TOKEN:
             return None
         cost_per_token = constants.OPENAI_COST_PER_TOKEN[self.model]
-        return cost_per_token["input"] * num_input_tokens + cost_per_token["output"] * num_output_tokens
+        return (
+            cost_per_token["input"] * num_input_tokens
+            + cost_per_token["output"] * num_output_tokens
+        )
 
     def _add_to_trace(self) -> None:
         """Adds to the trace."""
@@ -126,42 +141,56 @@ def _add_to_trace(self) -> None:
             metadata=self.metatada,
         )
 
-    def on_llm_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> Any:
+    def on_llm_error(
+        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+    ) -> Any:
         """Run when LLM errors."""
         pass
 
-    def on_chain_start(self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any) -> Any:
+    def on_chain_start(
+        self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
+    ) -> Any:
         """Run when chain starts running."""
         pass
 
     def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> Any:
         """Run when chain ends running."""
         pass
 
-    def on_chain_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> Any:
+    def on_chain_error(
+        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+    ) -> Any:
         """Run when chain errors."""
         pass
 
-    def on_tool_start(self, serialized: Dict[str, Any], input_str: str, **kwargs: Any) -> Any:
+    def on_tool_start(
+        self, serialized: Dict[str, Any], input_str: str, **kwargs: Any
+    ) -> Any:
         """Run when tool starts running."""
         pass
 
     def on_tool_end(self, output: str, **kwargs: Any) -> Any:
         """Run when tool ends running."""
         pass
 
-    def on_tool_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> Any:
+    def on_tool_error(
+        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+    ) -> Any:
         """Run when tool errors."""
         pass
 
     def on_text(self, text: str, **kwargs: Any) -> Any:
         """Run on arbitrary text."""
         pass
 
-    def on_agent_action(self, action: langchain_schema.AgentAction, **kwargs: Any) -> Any:
+    def on_agent_action(
+        self, action: langchain_schema.AgentAction, **kwargs: Any
+    ) -> Any:
         """Run on agent action."""
         pass
 
-    def on_agent_finish(self, finish: langchain_schema.AgentFinish, **kwargs: Any) -> Any:
+    def on_agent_finish(
+        self, finish: langchain_schema.AgentFinish, **kwargs: Any
+    ) -> Any:
         """Run on agent end."""
         pass
diff --git a/src/openlayer/lib/integrations/openai_tracer.py b/src/openlayer/lib/integrations/openai_tracer.py
@@ -139,12 +139,16 @@ def stream_chunks(
                 if delta.function_call.name:
                     collected_function_call["name"] += delta.function_call.name
                 if delta.function_call.arguments:
-                    collected_function_call["arguments"] += delta.function_call.arguments
+                    collected_function_call[
+                        "arguments"
+                    ] += delta.function_call.arguments
             elif delta.tool_calls:
                 if delta.tool_calls[0].function.name:
                     collected_function_call["name"] += delta.tool_calls[0].function.name
                 if delta.tool_calls[0].function.arguments:
-                    collected_function_call["arguments"] += delta.tool_calls[0].function.arguments
+                    collected_function_call["arguments"] += delta.tool_calls[
+                        0
+                    ].function.arguments
 
             yield chunk
         end_time = time.time()
@@ -155,16 +159,22 @@ def stream_chunks(
     finally:
         # Try to add step to the trace
         try:
-            collected_output_data = [message for message in collected_output_data if message is not None]
+            collected_output_data = [
+                message for message in collected_output_data if message is not None
+            ]
             if collected_output_data:
                 output_data = "".join(collected_output_data)
             else:
-                collected_function_call["arguments"] = json.loads(collected_function_call["arguments"])
+                collected_function_call["arguments"] = json.loads(
+                    collected_function_call["arguments"]
+                )
                 output_data = collected_function_call
             completion_cost = estimate_cost(
                 model=kwargs.get("model"),
                 prompt_tokens=0,
-                completion_tokens=(num_of_completion_tokens if num_of_completion_tokens else 0),
+                completion_tokens=(
+                    num_of_completion_tokens if num_of_completion_tokens else 0
+                ),
                 is_azure_openai=is_azure_openai,
             )
 
@@ -181,7 +191,13 @@ def stream_chunks(
                 model_parameters=get_model_parameters(kwargs),
                 raw_output=raw_outputs,
                 id=inference_id,
-                metadata={"timeToFirstToken": ((first_token_time - start_time) * 1000 if first_token_time else None)},
+                metadata={
+                    "timeToFirstToken": (
+                        (first_token_time - start_time) * 1000
+                        if first_token_time
+                        else None
+                    )
+                },
             )
             add_to_trace(
                 **trace_args,
@@ -207,7 +223,10 @@ def estimate_cost(
         cost_per_token = constants.AZURE_OPENAI_COST_PER_TOKEN[model]
     elif model in constants.OPENAI_COST_PER_TOKEN:
         cost_per_token = constants.OPENAI_COST_PER_TOKEN[model]
-        return cost_per_token["input"] * prompt_tokens + cost_per_token["output"] * completion_tokens
+        return (
+            cost_per_token["input"] * prompt_tokens
+            + cost_per_token["output"] * completion_tokens
+        )
     return None
 
 
@@ -266,8 +285,12 @@ def create_trace_args(
 def add_to_trace(is_azure_openai: bool = False, **kwargs) -> None:
     """Add a chat completion step to the trace."""
     if is_azure_openai:
-        tracer.add_chat_completion_step_to_trace(**kwargs, name="Azure OpenAI Chat Completion", provider="Azure")
-    tracer.add_chat_completion_step_to_trace(**kwargs, name="OpenAI Chat Completion", provider="OpenAI")
+        tracer.add_chat_completion_step_to_trace(
+            **kwargs, name="Azure OpenAI Chat Completion", provider="Azure"
+        )
+    tracer.add_chat_completion_step_to_trace(
+        **kwargs, name="OpenAI Chat Completion", provider="OpenAI"
+    )
 
 
 def handle_non_streaming_create(
@@ -327,7 +350,9 @@ def handle_non_streaming_create(
         )
     # pylint: disable=broad-except
     except Exception as e:
-        logger.error("Failed to trace the create chat completion request with Openlayer. %s", e)
+        logger.error(
+            "Failed to trace the create chat completion request with Openlayer. %s", e
+        )
 
     return response
 
@@ -369,7 +394,9 @@ def parse_non_streaming_output_data(
 
 
 # --------------------------- OpenAI Assistants API -------------------------- #
-def trace_openai_assistant_thread_run(client: openai.OpenAI, run: "openai.types.beta.threads.run.Run") -> None:
+def trace_openai_assistant_thread_run(
+    client: openai.OpenAI, run: "openai.types.beta.threads.run.Run"
+) -> None:
     """Trace a run from an OpenAI assistant.
 
     Once the run is completed, the thread data is published to Openlayer,
@@ -386,7 +413,9 @@ def trace_openai_assistant_thread_run(client: openai.OpenAI, run: "openai.types.
         metadata = _extract_run_metadata(run)
 
         # Convert thread to prompt
-        messages = client.beta.threads.messages.list(thread_id=run.thread_id, order="asc")
+        messages = client.beta.threads.messages.list(
+            thread_id=run.thread_id, order="asc"
+        )
         prompt = _thread_messages_to_prompt(messages)
 
         # Add step to the trace
