[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
jupyterlab · Sep 23, 2024 · 50af610 · 50af610
1 parent 04a6f36
commit 50af610
Show file tree

Hide file tree

Showing 5 changed files with 16 additions and 11 deletions.
diff --git a/packages/jupyter-ai-test/jupyter_ai_test/test_llms.py b/packages/jupyter-ai-test/jupyter_ai_test/test_llms.py
@@ -51,7 +51,7 @@ def _stream(
         time.sleep(1)
         yield GenerationChunk(
             text="Hello! This is a dummy response from a test LLM. I will now count from 1 to 5.\n\n",
-            generation_info={"test_metadata_field":"foobar"}
+            generation_info={"test_metadata_field": "foobar"},
         )
         for i in range(1, 6):
             time.sleep(0.2)

diff --git a/packages/jupyter-ai/jupyter_ai/callback_handlers/__init__.py b/packages/jupyter-ai/jupyter_ai/callback_handlers/__init__.py
@@ -3,4 +3,4 @@
 Not to be confused with Jupyter AI chat handlers.
 """
 
-from .metadata import MetadataCallbackHandler
+from .metadata import MetadataCallbackHandler
diff --git a/packages/jupyter-ai/jupyter_ai/callback_handlers/metadata.py b/packages/jupyter-ai/jupyter_ai/callback_handlers/metadata.py
@@ -1,6 +1,7 @@
 from langchain_core.callbacks import BaseCallbackHandler
 from langchain_core.outputs import LLMResult
 
+
 class MetadataCallbackHandler(BaseCallbackHandler):
     """
     When passed as a callback handler, this stores the LLMResult's

diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/default.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/default.py
@@ -82,16 +82,19 @@ def _start_stream(self, human_msg: HumanChatMessage) -> str:
 
         return stream_id
 
-    def _send_stream_chunk(self, stream_id: str, content: str, complete: bool = False, metadata: Dict[str, Any] = {}):
+    def _send_stream_chunk(
+        self,
+        stream_id: str,
+        content: str,
+        complete: bool = False,
+        metadata: Dict[str, Any] = {},
+    ):
         """
         Sends an `agent-stream-chunk` message containing content that should be
         appended to an existing `agent-stream` message with ID `stream_id`.
         """
         stream_chunk_msg = AgentStreamChunkMessage(
-            id=stream_id,
-            content=content,
-            stream_complete=complete,
-            metadata=metadata
+            id=stream_id, content=content, stream_complete=complete, metadata=metadata
         )
 
         for handler in self._root_chat_handlers.values():
@@ -116,7 +119,7 @@ async def process_message(self, message: HumanChatMessage):
                 {"input": message.body},
                 config={
                     "configurable": {"last_human_msg": message},
-                    "callbacks": [metadata_handler]
+                    "callbacks": [metadata_handler],
                 },
             ):
                 if not received_first_chunk:
@@ -135,4 +138,6 @@ async def process_message(self, message: HumanChatMessage):
                     break
 
             # complete stream after all chunks have been streamed
-            self._send_stream_chunk(stream_id, "", complete=True, metadata=metadata_handler.jai_metadata)
+            self._send_stream_chunk(
+                stream_id, "", complete=True, metadata=metadata_handler.jai_metadata
+            )
diff --git a/packages/jupyter-ai/jupyter_ai/models.py b/packages/jupyter-ai/jupyter_ai/models.py
@@ -86,7 +86,7 @@ class BaseAgentMessage(BaseModel):
     The persona of the selected provider. If the selected provider is `None`,
     this defaults to a description of `JupyternautPersona`.
     """
-    
+
     metadata: Dict[str, Any] = {}
     """
     Message metadata set by a provider after fully processing an input. The
@@ -122,7 +122,6 @@ class AgentStreamChunkMessage(BaseModel):
     """
 
 
-
 class HumanChatMessage(BaseModel):
     type: Literal["human"] = "human"
     id: str