chore: switch to ruff (#87)

Co-authored-by: Alan Konarski <[email protected]>

.github/workflows/ci.yml

@@ -41,8 +41,19 @@ jobs:
           source .venv/bin/activate
           pre-commit run --all-files --show-diff-on-failure --color always
 
+      - name: Run ruff formatter
+        run: uvx ruff format --check
+
+      - name: Run ruff linter
+        run: uvx ruff check
+
+      - name: Run mypy
+        run: uv run mypy .
+
       - name: Run Trivy vulnerability scanner
         uses: aquasecurity/trivy-action@master
+        env:
+          TRIVY_DB_REPOSITORY: ghcr.io/aquasecurity/trivy-db:2,public.ecr.aws/aquasecurity/trivy-db:2
         with:
           scan-type: "fs"
           ignore-unfixed: true

.pre-commit-config.yaml

@@ -1,3 +1,5 @@
+default_language_version:
+    python: python3.10
 repos:
     - repo: https://github.com/pre-commit/pre-commit-hooks
       rev: v4.5.0
@@ -11,83 +13,3 @@ repos:
           - id: check-toml
           - id: check-json
           - id: check-yaml
-
-      # PEP 8 compliant opinionated formatter.
-    - repo: https://github.com/psf/black
-      rev: 23.10.1
-      hooks:
-          - id: black
-            exclude: (docs/|notebooks/)
-            args: [--config, pyproject.toml]
-          - id: black-jupyter
-            files: \.ipynb$
-
-      # Cleaning unused imports.
-    - repo: https://github.com/hadialqattan/pycln
-      rev: v2.3.0
-      hooks:
-          - id: pycln
-            args: ["-a"]
-            exclude: (docs/|notebooks/)
-
-    # Modernizes python code and upgrade syntax for newer versions of the language
-    - repo: https://github.com/asottile/pyupgrade
-      rev: v3.15.0
-      hooks:
-          - id: pyupgrade
-            args: [--py38-plus]
-
-      # Used to have proper type annotations for library code.
-    - repo: https://github.com/pre-commit/mirrors-mypy
-      rev: v1.11.2
-      hooks:
-          - id: mypy
-            # You can add additional plugins for mypy below
-            # such as types-python-dateutil
-            additional_dependencies: [pydantic>=2.8.2, types-pyyaml>=6.0.12]
-            exclude: (/test_|setup.py|/tests/|docs/)
-
-      # Sort imports alphabetically, and automatically separated into sections and by type.
-    - repo: https://github.com/timothycrosley/isort
-      rev: 5.13.2
-      hooks:
-          - id: isort
-            args: ["--profile", "black"]
-            exclude: (docs/|notebooks/)
-
-      # Checks Python source files for errors.
-    - repo: https://github.com/PyCQA/flake8
-      rev: 7.1.1
-      hooks:
-          - id: flake8
-            name: flake8
-            entry: flake8
-            language: python
-            types: [python]
-            args: [--config, .flake8]
-            exclude: (docs/)
-
-      # Enforces a coding standard, looks for code smells, and can make suggestions about how the code could be refactored.
-    - repo: https://github.com/pycqa/pylint
-      rev: v3.2.6
-      hooks:
-          - id: pylint
-            exclude: (/test_|tests/|docs/)
-            # # You can add additional plugins for pylint here,
-            # here is an example for pydantic, remember to enable it in pyproject.toml
-            # additional_dependencies:
-            #   - 'pylint_pydantic'
-            # args:
-            #   # pylint can have issue with python libraries based on C
-            #   # if it fails to find some objects likely you need to add them
-            #   # here:
-            #   ["--extension-pkg-whitelist=pydantic"]
-
-      # Finds common security issues in Python code.
-    - repo: https://github.com/PyCQA/bandit
-      rev: 1.7.5
-      hooks:
-          - id: bandit
-            args: [-c, pyproject.toml, --recursive, packages/]
-            additional_dependencies: [".[toml]"] # required for pyproject.toml support
-            exclude:  (notebooks/)

CONTRIBUTING.md

@@ -11,9 +11,29 @@ To build and run Ragbits from the source code:
 $ source ./setup_dev_env.sh
 ```
 
+## Linting and formatting
+We use `ruff` for linting and formatting our code. To format your code, run:
+
+```bash
+$ uvx ruff format
+```
+
+To lint the code, run:
+```bash
+$ uvx ruff check --fix
+```
+
+## Type checking
+We use `mypy` for type checking. To perform type checking, simply run:
+
+```bash
+$ uv run mypy .
+```
+
+
 ## Install pre-commit
 
-To ensure code quality we use pre-commit hook with several checks. Setup it by:
+We also run some checks through a pre-commit hook. To set it up, follow these steps:
 
 ```
 pre-commit install

examples/apps/documents_chat.py

@@ -6,8 +6,8 @@
 #     "ragbits-core[chromadb, litellm]",
 # ]
 # ///
+from collections.abc import AsyncIterator
 from pathlib import Path
-from typing import AsyncIterator
 
 import chromadb
 import gradio as gr
@@ -104,7 +104,7 @@ def _prepare_document_search(self, database_path: str, index_name: str) -> None:
             chroma_client=chroma_client,
             embedding_function=embedding_client,
         )
-        self.document_search = DocumentSearch(embedder=vector_store.embedding_function, vector_store=vector_store)
+        self.document_search = DocumentSearch(embedder=embedding_client, vector_store=vector_store)
 
     async def _create_database(self, document_paths: list[str]) -> str:
         for path in document_paths:
@@ -118,7 +118,9 @@ def _load_database(self, database_path: str) -> str:
         return self.DATABASE_LOADED_MESSAGE
 
     async def _handle_message(
-        self, message: str, history: list[dict]  # pylint: disable=unused-argument
+        self,
+        message: str,
+        history: list[dict],  # pylint: disable=unused-argument
     ) -> AsyncIterator[str]:
         if not self._documents_ingested:
             yield self.NO_DOCUMENTS_INGESTED_MESSAGE
@@ -134,32 +136,31 @@ def prepare_layout(self) -> gr.Blocks:
         Returns:
             gradio layout
         """
-        with gr.Blocks(fill_height=True, fill_width=True) as app:
-            with gr.Row():
-                with gr.Column(scale=self._columns_ratios[0]):
-                    with gr.Group():
-                        documents_picker = gr.File(file_count="multiple", label=self.DOCUMENT_PICKER_LABEL)
-                        create_btn = gr.Button(self.DATABASE_CREATE_BUTTON_LABEL)
-                        creating_status_display = gr.Textbox(
-                            label=self.DATABASE_CREATION_STATUS_LABEL,
-                            interactive=False,
-                            placeholder=self.DATABASE_CREATION_STATUS_PLACEHOLDER,
-                        )
-
-                    with gr.Group():
-                        database_path = gr.Textbox(label=self.DATABASE_TEXT_BOX_LABEL)
-                        load_btn = gr.Button(self.DATABASE_LOAD_BUTTON_LABEL)
-                        loading_status_display = gr.Textbox(
-                            label=self.DATABASE_LOADING_STATUS_LABEL,
-                            interactive=False,
-                            placeholder=self.DATABASE_LOADING_STATUS_PLACEHOLDER,
-                        )
-                    load_btn.click(fn=self._load_database, inputs=database_path, outputs=loading_status_display)
-                    create_btn.click(fn=self._create_database, inputs=documents_picker, outputs=creating_status_display)
-
-                with gr.Column(scale=self._columns_ratios[1]):
-                    chat_interface = gr.ChatInterface(self._handle_message, type="messages")
-                    chat_interface.chatbot.height = f"{self._chatbot_height_vh}vh"
+        with gr.Blocks(fill_height=True, fill_width=True) as app, gr.Row():
+            with gr.Column(scale=self._columns_ratios[0]):
+                with gr.Group():
+                    documents_picker = gr.File(file_count="multiple", label=self.DOCUMENT_PICKER_LABEL)
+                    create_btn = gr.Button(self.DATABASE_CREATE_BUTTON_LABEL)
+                    creating_status_display = gr.Textbox(
+                        label=self.DATABASE_CREATION_STATUS_LABEL,
+                        interactive=False,
+                        placeholder=self.DATABASE_CREATION_STATUS_PLACEHOLDER,
+                    )
+
+                with gr.Group():
+                    database_path = gr.Textbox(label=self.DATABASE_TEXT_BOX_LABEL)
+                    load_btn = gr.Button(self.DATABASE_LOAD_BUTTON_LABEL)
+                    loading_status_display = gr.Textbox(
+                        label=self.DATABASE_LOADING_STATUS_LABEL,
+                        interactive=False,
+                        placeholder=self.DATABASE_LOADING_STATUS_PLACEHOLDER,
+                    )
+                load_btn.click(fn=self._load_database, inputs=database_path, outputs=loading_status_display)
+                create_btn.click(fn=self._create_database, inputs=documents_picker, outputs=creating_status_display)
+
+            with gr.Column(scale=self._columns_ratios[1]):
+                chat_interface = gr.ChatInterface(self._handle_message, type="messages")
+                chat_interface.chatbot.height = f"{self._chatbot_height_vh}vh"
         return app
 
 

examples/core/llm.py

@@ -36,8 +36,9 @@ class JokePrompt(Prompt[LoremPromptInput, LoremPromptOutput]):
     """
 
     system_prompt = """
-    You are a joke generator. The jokes you generate should be funny and not offensive. {% if not pun_allowed %}Also, make sure
-    that the jokes do not contain any puns.{% else %}You can use any type of joke, even if it contains puns.{% endif %}
+    You are a joke generator. The jokes you generate should be funny and not offensive.
+    {% if not pun_allowed %}Also, make sure that the jokes do not contain any puns.
+    {% else %}You can use any type of joke, even if it contains puns.{% endif %}
 
     Respond as json with two fields: joke and joke_category.
     """
@@ -47,7 +48,7 @@ class JokePrompt(Prompt[LoremPromptInput, LoremPromptOutput]):
     """
 
 
-async def main():
+async def main() -> None:
     """
     Example of using the LiteLLM client with a Prompt class. Requires the OPENAI_API_KEY environment variable to be set.
     """

examples/document-search/basic.py

@@ -23,9 +23,10 @@
 ]
 
 
-async def main():
-    """Run the example."""
-
+async def main() -> None:
+    """
+    Run the example.
+    """
     document_search = DocumentSearch(embedder=LiteLLMEmbeddings(), vector_store=InMemoryVectorStore())
 
     await document_search.ingest(documents)

examples/document-search/chroma.py

@@ -23,9 +23,10 @@
 ]
 
 
-async def main():
-    """Run the example."""
-
+async def main() -> None:
+    """
+    Run the example.
+    """
     chroma_client = chromadb.PersistentClient(path="chroma")
     embedding_client = LiteLLMEmbeddings()
 
@@ -34,7 +35,7 @@ async def main():
         chroma_client=chroma_client,
         embedding_function=embedding_client,
     )
-    document_search = DocumentSearch(embedder=vector_store.embedding_function, vector_store=vector_store)
+    document_search = DocumentSearch(embedder=embedding_client, vector_store=vector_store)
 
     await document_search.ingest(documents)
 

examples/document-search/from_config.py

@@ -47,9 +47,8 @@
 }
 
 
-async def main():
+async def main() -> None:
     """Run the example."""
-
     document_search = DocumentSearch.from_config(config)
 
     await document_search.ingest(documents)

packages/ragbits-cli/src/ragbits/cli/__init__.py

@@ -19,13 +19,10 @@ def main() -> None:
         - if found it imports the `register` function from the `cli` module and calls it with the `app` object
         - register function should add the CLI commands to the `app` object
     """
-
     cli_enabled_modules = [
         module
         for module in pkgutil.iter_modules(ragbits.__path__)
-        if module.ispkg
-        and module.name != "cli"
-        and (Path(module.module_finder.path) / module.name / "cli.py").exists()  # type: ignore
+        if module.ispkg and module.name != "cli" and (Path(module.module_finder.path) / module.name / "cli.py").exists()  # type: ignore
     ]
 
     for module in cli_enabled_modules:

packages/ragbits-core/src/ragbits/core/cli.py

@@ -33,7 +33,7 @@ def lab(
     @prompts_app.command()
     def generate_promptfoo_configs(
         file_pattern: str = core_config.prompt_path_pattern,
-        root_path: Path = Path.cwd(),
+        root_path: Path = Path.cwd(),  # noqa: B008
         target_path: Path = Path("promptfooconfigs"),
     ) -> None:
         """

packages/ragbits-core/src/ragbits/core/embeddings/litellm.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 try:
     import litellm
 
@@ -8,7 +6,11 @@
     HAS_LITELLM = False
 
 from ragbits.core.embeddings import Embeddings
-from ragbits.core.embeddings.exceptions import EmbeddingConnectionError, EmbeddingResponseError, EmbeddingStatusError
+from ragbits.core.embeddings.exceptions import (
+    EmbeddingConnectionError,
+    EmbeddingResponseError,
+    EmbeddingStatusError,
+)
 
 
 class LiteLLMEmbeddings(Embeddings):
@@ -19,10 +21,10 @@ class LiteLLMEmbeddings(Embeddings):
     def __init__(
         self,
         model: str = "text-embedding-3-small",
-        options: Optional[dict] = None,
-        api_base: Optional[str] = None,
-        api_key: Optional[str] = None,
-        api_version: Optional[str] = None,
+        options: dict | None = None,
+        api_base: str | None = None,
+        api_key: str | None = None,
+        api_version: str | None = None,
     ) -> None:
         """
         Constructs the LiteLLMEmbeddingClient.
@@ -65,7 +67,6 @@ async def embed_text(self, data: list[str]) -> list[list[float]]:
             EmbeddingStatusError: If the embedding API returns an error status code.
             EmbeddingResponseError: If the embedding API response is invalid.
         """
-
         try:
             response = await litellm.aembedding(
                 input=data,