Added ruff linting checks and applied fixes

Author: tomusher

pyproject.toml

@@ -28,6 +28,10 @@ classifiers = [
   "Framework :: Django :: 6.0",
 ]
 
+[project.urls]
+Home = "https://github.com/wagtail/django-ai-core"
+
+
 [dependency-groups]
 dev = [
     "django>=5.2.3",
@@ -57,6 +61,22 @@ build-backend = "uv_build"
 [tool.uv]
 default-groups = ["dev", "test"]
 
+[tool.ruff.lint]
+select = [
+    "B",   # flake8-bugbear
+    "C4",  # flake8-comprehensions
+    "C90", # mccabe
+    "DJ",  # flake8-django
+    "E",   # pycodestyle errors
+    "F",   # pyflakes
+    "I",   # isort
+    "RUF", # ruff-specific rules
+    "TRY", # tryceratops
+    "PT",  # flake8-pytest-style
+    "SIM", # flake8-simplify
+]
+ignore = ["TRY003", "E501", "RUF012"]
+
 [tool.pytest.ini_options]
 DJANGO_SETTINGS_MODULE = "testapp.settings"
 pythonpath = ["./src", "./tests"]

src/django_ai_core/contrib/agents/base.py

@@ -1,10 +1,10 @@
 import inspect
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import get_args, get_origin, Annotated
+from typing import Annotated, get_args, get_origin
 
-from django.core.validators import validate_slug
 from django.core.exceptions import ValidationError
+from django.core.validators import validate_slug
 
 from .permissions import BasePermission
 from .views import AgentExecutionView
@@ -68,10 +68,10 @@ def __init_subclass__(cls, **kwargs):
         if hasattr(cls, "slug"):
             try:
                 validate_slug(cls.slug)
-            except ValidationError:
+            except ValidationError as e:
                 raise ValueError(
                     f"Agent {cls.__name__} has an invalid slug: {cls.slug}. Use a valid “slug” consisting of letters, numbers, underscores or hyphens."
-                )
+                ) from e
 
 
 class AgentRegistry:
@@ -82,7 +82,7 @@ def register(self):
         """Decorator to register an agent."""
 
         def decorator(cls: type[Agent]) -> type[Agent]:
-            agent_slug = getattr(cls, "slug")
+            agent_slug = cls.slug
             self._agents[agent_slug] = cls
             return cls
 

src/django_ai_core/contrib/agents/views.py

@@ -1,10 +1,10 @@
-from typing import Any, TYPE_CHECKING
 import json
+from typing import TYPE_CHECKING, Any
 
 from django.http import JsonResponse
+from django.utils.decorators import method_decorator
 from django.views import View
 from django.views.decorators.csrf import csrf_exempt
-from django.utils.decorators import method_decorator
 
 from .permissions import AllowAny
 
@@ -97,8 +97,8 @@ def _get_agent(self) -> "Agent":
 
         try:
             return registry.get(self.agent_slug)()
-        except KeyError:
-            raise AgentNotFound
+        except KeyError as e:
+            raise AgentNotFound from e
 
     def _execute_agent(self, agent: "Agent", arguments: dict[str, Any]) -> Any:
         return agent.execute(**arguments)

src/django_ai_core/contrib/index/__init__.py

@@ -1,13 +1,13 @@
 from .base import (
-    VectorIndex,
     IndexRegistry,
+    VectorIndex,
     registry,
 )
 from .chunking import (
     ChunkTransformer,
-    SimpleChunkTransformer,
-    SentenceChunkTransformer,
     ParagraphChunkTransformer,
+    SentenceChunkTransformer,
+    SimpleChunkTransformer,
 )
 from .embedding import (
     CoreEmbeddingTransformer,
@@ -22,18 +22,17 @@
     StorageProvider,
 )
 
-
 __all__ = [
+    "CachedEmbeddingTransformer",
     "ChunkTransformer",
-    "ModelSource",
+    "CoreEmbeddingTransformer",
     "IndexRegistry",
+    "ModelSource",
+    "ModelSource",
     "ParagraphChunkTransformer",
     "SentenceChunkTransformer",
     "SimpleChunkTransformer",
-    "CoreEmbeddingTransformer",
-    "CachedEmbeddingTransformer",
     "StorageProvider",
     "VectorIndex",
-    "ModelSource",
     "registry",
 ]

src/django_ai_core/contrib/index/base.py

@@ -1,15 +1,16 @@
+import logging
 from typing import TYPE_CHECKING, Any, ClassVar, Iterable
+
 from django.utils.text import slugify
+
 from .source import HasPostIndexUpdateHook
-import logging
 
 if TYPE_CHECKING:
-    from .schema import Document
-    from .storage.base import StorageProvider
-    from .query import ResultQuerySetMixin
     from .embedding import EmbeddingTransformer
+    from .query import QueryHandler, ResultQuerySetMixin
+    from .schema import Document
     from .source import Source
-    from .query import QueryHandler
+    from .storage.base import StorageProvider
 
 
 logger = logging.getLogger(__name__)

src/django_ai_core/contrib/index/chunking.py

@@ -52,8 +52,8 @@ def __init__(self, chunk_size: int = 1024, chunk_overlap: int = 20):
         self.chunk_overlap = chunk_overlap
 
     def transform(self, text: str) -> list[str]:
-        from llama_index.core.node_parser import SentenceSplitter
         from llama_index.core import Document as LlamaDocument
+        from llama_index.core.node_parser import SentenceSplitter
 
         splitter = SentenceSplitter(
             chunk_size=self.chunk_size, chunk_overlap=self.chunk_overlap
@@ -63,7 +63,7 @@ def transform(self, text: str) -> list[str]:
         nodes = splitter.get_nodes_from_documents([llama_doc])
 
         chunks = []
-        for i, node in enumerate(nodes):
+        for _, node in enumerate(nodes):
             chunks.append(
                 node.get_content(),
             )

src/django_ai_core/contrib/index/embedding.py

@@ -1,9 +1,10 @@
 from abc import ABC, abstractmethod
 from typing import Iterable
-from .schema import EmbeddedDocument, Document
 
 from django_ai_core.llm import LLMService
 
+from .schema import Document, EmbeddedDocument
+
 
 class EmbeddingTransformer(ABC):
     """Base class for embedding transformers which turn Documents into EmbeddedDocuments."""
@@ -13,6 +14,7 @@ def transformer_id(self) -> str:
         """Get unique identifier for this transformer."""
         return self.__class__.__name__
 
+    @abstractmethod
     def embed_string(self, text: str) -> list[float] | None:
         """Embed a string using the transformer."""
         pass
@@ -68,7 +70,7 @@ def embed_documents(
             texts = [document.content for document in batch]
             embeddings = self.llm_service.embedding(texts).data
 
-            for document, embedding in zip(batch, embeddings):
+            for document, embedding in zip(batch, embeddings, strict=False):
                 embedded_documents.append(document.add_embedding(embedding.embedding))
 
         return embedded_documents

src/django_ai_core/contrib/index/embedding_cache.py

@@ -5,8 +5,8 @@
 generation to avoid expensive embedding operations for duplicate content.
 """
 
-from abc import ABC, abstractmethod
 import logging
+from abc import ABC, abstractmethod
 
 from django_ai_core.contrib.index.schema import Document, EmbeddedDocument
 
@@ -61,13 +61,13 @@ def _get_cache_model(self):
         """Get the cache model."""
         try:
             from .models import EmbeddingCache
-
-            return EmbeddingCache
-        except ImportError:
+        except ImportError as e:
             raise ImportError(
                 "Django is not properly configured. Make sure Django settings are loaded "
                 "and 'django_ai_core.contrib.index' is in INSTALLED_APPS."
-            )
+            ) from e
+        else:
+            return EmbeddingCache
 
     def get_embedding(self, content: str, transformer_id: str) -> list[float] | None:
         """Get cached embedding for content and model."""
@@ -202,7 +202,7 @@ def embed_documents(
             result[i] = document
 
         # Place newly embedded documents
-        for i, document in zip(uncached_indices, embedded_documents):
+        for i, document in zip(uncached_indices, embedded_documents, strict=False):
             result[i] = document
 
         return [document for document in result if document is not None]

src/django_ai_core/contrib/index/management/commands/rebuild_indexes.py

@@ -4,11 +4,12 @@
 This command rebuilds all registered VectorIndex instances.
 """
 
-from django.core.management.base import BaseCommand, CommandError
-from django.utils import timezone
 import logging
 import time
 
+from django.core.management.base import BaseCommand, CommandError
+from django.utils import timezone
+
 from django_ai_core.contrib.index.base import registry
 
 logger = logging.getLogger(__name__)
@@ -108,8 +109,6 @@ def _rebuild_sequential(
             try:
                 start_time = time.time()
                 index_cls = registry.get(index_name)
-                if not index_cls:
-                    raise ValueError(f"Index '{index_name}' not found in registry")
 
                 self.stdout.write(
                     f"\n[{i}/{len(indexes_to_rebuild)}] Rebuilding index: {index_name}"

src/django_ai_core/contrib/index/models.py

@@ -1,8 +1,9 @@
 import hashlib
 from typing import TYPE_CHECKING
-from django.db import models
+
 from django.contrib.contenttypes.fields import GenericForeignKey
 from django.contrib.contenttypes.models import ContentType
+from django.db import models
 
 from .base import registry
 
@@ -79,7 +80,6 @@ def get_or_create_embedding(
                 content_hash=content_hash,
                 embedding_transformer_id=embedding_transformer_id,
             )
-            return cache_entry, False
         except cls.DoesNotExist:
             # Create new cache entry
             cache_entry = cls.objects.create(
@@ -90,6 +90,8 @@ def get_or_create_embedding(
                 embedding_dimensions=len(embedding_vector),
             )
             return cache_entry, True
+        else:
+            return cache_entry, False
 
     @classmethod
     def get_cached_embedding(
@@ -105,9 +107,10 @@ def get_cached_embedding(
                 content_hash=content_hash,
                 embedding_transformer_id=embedding_transformer_id,
             )
-            return cache_entry.embedding_vector
         except cls.DoesNotExist:
             return None
+        else:
+            return cache_entry.embedding_vector
 
 
 class DocumentEmbedding(models.Model):
@@ -164,6 +167,9 @@ class Meta:
             ("content_type", "object_id", "index_name", "source_id"),
         ]
 
+    def __str__(self):
+        return f"Model Source Index: {self.object_id} in {self.index_name}"
+
     @classmethod
     def register(cls, obj, index_name, source_id):
         """Register an object as being indexed in the specified index."""