From f50b491f56412d14aace7625d0e9a4c54cebd9b1 Mon Sep 17 00:00:00 2001
From: Vladimir Blagojevic <dovlex@gmail.com>
Date: Tue, 20 Aug 2024 20:27:15 +0200
Subject: [PATCH] Lint

---
 .../anthropic/example/prompt_caching.py       | 28 ++++++++++++-------
 integrations/anthropic/pyproject.toml         |  1 +
 2 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/integrations/anthropic/example/prompt_caching.py b/integrations/anthropic/example/prompt_caching.py
index 495f292e6..6ec71c520 100644
--- a/integrations/anthropic/example/prompt_caching.py
+++ b/integrations/anthropic/example/prompt_caching.py
@@ -41,25 +41,33 @@
 
 # Build QA pipeline
 qa_pipeline = Pipeline()
-qa_pipeline.add_component("llm", AnthropicChatGenerator(
-    api_key=Secret.from_env_var("ANTHROPIC_API_KEY"),
-    streaming_callback=print_streaming_chunk,
-    generation_kwargs={"extra_headers": {"anthropic-beta": "prompt-caching-2024-07-31"}},
-))
+qa_pipeline.add_component(
+    "llm",
+    AnthropicChatGenerator(
+        api_key=Secret.from_env_var("ANTHROPIC_API_KEY"),
+        streaming_callback=print_streaming_chunk,
+        generation_kwargs={"extra_headers": {"anthropic-beta": "prompt-caching-2024-07-31"}},
+    ),
+)
 
-questions = ["Why is Monte-Carlo Tree Search used in LATS",
-             "Summarize LATS selection, expansion, evaluation, simulation, backpropagation, and reflection"]
+questions = [
+    "Why is Monte-Carlo Tree Search used in LATS",
+    "Summarize LATS selection, expansion, evaluation, simulation, backpropagation, and reflection",
+]
 
 # Answer the questions using prompt caching (i.e. the entire document is cached, we run the question against it)
 for question in questions:
     print("Question: " + question)
     result = qa_pipeline.run(
         data={
-            "llm": {"messages": [final_prompt_msg,
-                                 ChatMessage.from_user("Given these documents, answer the question:" + question)]},
+            "llm": {
+                "messages": [
+                    final_prompt_msg,
+                    ChatMessage.from_user("Given these documents, answer the question:" + question),
+                ]
+            },
         }
     )
 
     print("\n\nChecking cache usage:", result["llm"]["replies"][0].meta.get("usage"))
     print("\n")
-
diff --git a/integrations/anthropic/pyproject.toml b/integrations/anthropic/pyproject.toml
index 3f8c9812b..e1d3fa867 100644
--- a/integrations/anthropic/pyproject.toml
+++ b/integrations/anthropic/pyproject.toml
@@ -106,6 +106,7 @@ select = [
   "YTT",
 ]
 ignore = [
+  "T201", # print statements
   # Allow non-abstract empty methods in abstract base classes
   "B027",
   # Ignore checks for possible passwords