ci: ignore rye linting errors for custom code

whoseoyster · whoseoyster · commit 40aa59885d21 · 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,9 +42,7 @@ 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,
@@ -85,9 +83,7 @@ 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)
@@ -137,9 +133,8 @@ def write_output_to_directory(
         """Writes the output DataFrame to a file in the specified directory based on the
         given format.
         """
-        os.makedirs(
-            output_dir, exist_ok=True
-        )  # Create the directory if it doesn't exist
+        # Create the directory if it doesn't exist
+        os.makedirs(output_dir, exist_ok=True)
 
         # Determine the filename based on the dataset name and format
         filename = f"dataset.{fmt}"
@@ -158,7 +153,7 @@ def write_output_to_directory(
         else:
             raise ValueError("Unsupported format. Please choose 'csv' or 'json'.")
 
-        print(f"Output written to {output_path}")
+        print(f"Output written to {output_path}")  # noqa: T201
 
     @abc.abstractmethod
     def run(self, **kwargs) -> RunReturn:
diff --git a/src/openlayer/lib/integrations/langchain_callback.py b/src/openlayer/lib/integrations/langchain_callback.py
@@ -17,10 +17,7 @@
 class OpenlayerHandler(BaseCallbackHandler):
     """LangChain callback handler that logs to Openlayer."""
 
-    def __init__(
-        self,
-        **kwargs: Any,
-    ) -> None:
+    def __init__(self, **kwargs: Any) -> None:
         super().__init__()
 
         self.start_time: float = None
@@ -37,14 +34,14 @@ def __init__(
         self.output: str = None
         self.metatada: Dict[str, Any] = kwargs or {}
 
-    def on_llm_start(
-        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
-    ) -> Any:
+    # noqa arg002
+    def on_llm_start(self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any) -> Any:
         """Run when LLM starts running."""
+        pass
 
     def on_chat_model_start(
         self,
-        serialized: Dict[str, Any],
+        serialized: Dict[str, Any],  # noqa: ARG002
         messages: List[List[langchain_schema.BaseMessage]],
         **kwargs: Any,
     ) -> Any:
@@ -80,44 +77,34 @@ def _langchain_messages_to_prompt(
 
     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:
+    def on_llm_end(self, response: langchain_schema.LLMResult, **kwargs: Any) -> Any:  # noqa: ARG002, E501
         """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."""
@@ -139,46 +126,42 @@ 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,16 +139,12 @@ 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()
@@ -159,22 +155,16 @@ 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,
             )
 
@@ -191,13 +181,7 @@ 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,
@@ -223,10 +207,7 @@ 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
 
 
@@ -285,12 +266,8 @@ 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(
@@ -350,9 +327,7 @@ 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
 
@@ -394,9 +369,7 @@ 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,
@@ -413,9 +386,7 @@ def trace_openai_assistant_thread_run(
         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
@@ -430,7 +401,7 @@ def trace_openai_assistant_thread_run(
 
     # pylint: disable=broad-except
     except Exception as e:
-        print(f"Failed to monitor run. {e}")
+        print(f"Failed to monitor run. {e}")  # noqa: T201
 
 
 def _type_check_run(run: "openai.types.beta.threads.run.Run") -> None:
