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

for more information, see https://pre-commit.ci
jupyterlab · Jul 14, 2023 · 9138ed0 · 9138ed0
1 parent aab694f
commit 9138ed0
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 10 deletions.
diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/learn.py
@@ -7,11 +7,11 @@
 from jupyter_ai.document_loaders.directory import get_embeddings, split
 from jupyter_ai.document_loaders.splitter import ExtensionSplitter, NotebookSplitter
 from jupyter_ai.models import (
-    DEFAULT_CHUNK_OVERLAP, 
-    DEFAULT_CHUNK_SIZE, 
-    HumanChatMessage, 
-    IndexedDir, 
-    IndexMetadata
+    DEFAULT_CHUNK_OVERLAP,
+    DEFAULT_CHUNK_SIZE,
+    HumanChatMessage,
+    IndexedDir,
+    IndexMetadata,
 )
 from jupyter_core.paths import jupyter_data_dir
 from langchain import FAISS
@@ -40,8 +40,16 @@ def __init__(
         self.parser.add_argument("-v", "--verbose", action="store_true")
         self.parser.add_argument("-d", "--delete", action="store_true")
         self.parser.add_argument("-l", "--list", action="store_true")
-        self.parser.add_argument("-c", "--chunk-size", action="store", default=DEFAULT_CHUNK_SIZE, type=int)
-        self.parser.add_argument("-o", "--chunk-overlap", action="store", default=DEFAULT_CHUNK_OVERLAP, type=int)
+        self.parser.add_argument(
+            "-c", "--chunk-size", action="store", default=DEFAULT_CHUNK_SIZE, type=int
+        )
+        self.parser.add_argument(
+            "-o",
+            "--chunk-overlap",
+            action="store",
+            default=DEFAULT_CHUNK_OVERLAP,
+            type=int,
+        )
         self.parser.add_argument("path", nargs=argparse.REMAINDER)
         self.index_name = "default"
         self.index = None
@@ -125,7 +133,9 @@ def _build_list_response(self):
         {dir_list}"""
         return message
 
-    async def learn_dir(self, path: str, chunk_size: int =2000, chunk_overlap: int = 100):
+    async def learn_dir(
+        self, path: str, chunk_size: int = 2000, chunk_overlap: int = 100
+    ):
         dask_client = await self.dask_client_future
         splitters = {
             ".py": PythonCodeTextSplitter(
@@ -162,7 +172,11 @@ def _add_dir_to_metadata(self, path: str, chunk_size: int, chunk_overlap: int):
         dirs = self.metadata.dirs
         index = next((i for i, dir in enumerate(dirs) if dir.path == path), None)
         if not index:
-            dirs.append(IndexedDir(path=path, chunk_size=chunk_size, chunk_overlap=chunk_overlap))
+            dirs.append(
+                IndexedDir(
+                    path=path, chunk_size=chunk_size, chunk_overlap=chunk_overlap
+                )
+            )
         self.metadata.dirs = dirs
 
     async def delete_and_relearn(self):

diff --git a/packages/jupyter-ai/jupyter_ai/models.py b/packages/jupyter-ai/jupyter_ai/models.py
@@ -3,10 +3,10 @@
 from jupyter_ai_magics.providers import AuthStrategy, Field
 from pydantic import BaseModel
 
-
 DEFAULT_CHUNK_SIZE = 2000
 DEFAULT_CHUNK_OVERLAP = 100
 
+
 # the type of message used to chat with the agent
 class ChatRequest(BaseModel):
     prompt: str