feat: validate class names in dynamic adapters (#1408)

RinZ27 · Tomas2D · web-flow · commit 23d1e52fec74 · 2026-04-08T09:49:41.000+02:00
Signed-off-by: RinZ27 &lt;222222878+RinZ27@users.noreply.github.com&gt;
Co-authored-by: Tomáš Dvořák &lt;toomas2d@gmail.com&gt;
diff --git a/python/beeai_framework/adapters/agentstack/backend/vector_store.py b/python/beeai_framework/adapters/agentstack/backend/vector_store.py
@@ -8,6 +8,7 @@
 from beeai_framework.backend.embedding import EmbeddingModel
 from beeai_framework.backend.types import Document, DocumentWithScore
 from beeai_framework.backend.vector_store import QueryLike, VectorStore
+from beeai_framework.utils.strings import validate_class_name
 
 try:
     from agentstack_sdk.platform import VectorStore as AgentStackSDKVectorStore
@@ -22,6 +23,8 @@ class AgentStackVectorStore(VectorStore, ABC):
     @classmethod
     def _class_from_name(cls, class_name: str, embedding_model: EmbeddingModel, **kwargs: Any) -> VectorStore:
         """Create an instance from class name (required by VectorStore base class)."""
+        validate_class_name(class_name)
+
         # Get the current module to look for classes
         import sys
 
diff --git a/python/beeai_framework/adapters/beeai/backend/document_processor.py b/python/beeai_framework/adapters/beeai/backend/document_processor.py
@@ -14,6 +14,7 @@
 from beeai_framework.backend.chat import ChatModel
 from beeai_framework.backend.document_processor import DocumentProcessor
 from beeai_framework.backend.types import DocumentWithScore
+from beeai_framework.utils.strings import validate_class_name
 
 try:
     from llama_index.core.postprocessor.llm_rerank import LLMRerank
@@ -27,6 +28,8 @@ class BeeAIDocumentProcessor(DocumentProcessor, ABC):
     @classmethod
     def _class_from_name(cls, class_name: str, **kwargs: Any) -> BeeAIDocumentProcessor:
         """Create an instance from class name (required by DocumentProcessor base class)."""
+        validate_class_name(class_name)
+
         # Get the current module to look for classes
         import sys
 
diff --git a/python/beeai_framework/adapters/beeai/backend/vector_store.py b/python/beeai_framework/adapters/beeai/backend/vector_store.py
@@ -9,6 +9,7 @@
 from beeai_framework.backend.embedding import EmbeddingModel
 from beeai_framework.backend.types import Document, DocumentWithScore
 from beeai_framework.backend.vector_store import QueryLike, VectorStore
+from beeai_framework.utils.strings import validate_class_name
 
 try:
     from langchain_core.documents import Document as LCDocument
@@ -26,6 +27,8 @@ class BeeAIVectorStore(VectorStore, ABC):
     @classmethod
     def _class_from_name(cls, class_name: str, embedding_model: EmbeddingModel, **kwargs: Any) -> BeeAIVectorStore:
         """Create an instance from class name (required by VectorStore base class)."""
+        validate_class_name(class_name)
+
         # Get the current module to look for classes
         import sys
 
diff --git a/python/beeai_framework/adapters/langchain/backend/document_loader.py b/python/beeai_framework/adapters/langchain/backend/document_loader.py
@@ -16,6 +16,7 @@
 from beeai_framework.adapters.langchain.mappers.documents import lc_document_to_document
 from beeai_framework.backend.document_loader import DocumentLoader
 from beeai_framework.backend.types import Document
+from beeai_framework.utils.strings import validate_class_name
 from beeai_framework.logger import Logger
 
 logger = Logger(__name__)
@@ -31,6 +32,8 @@ def _class_from_name(cls, class_name: str, **kwargs: Any) -> LangChainDocumentLo
         """
         Dynamically imports and instantiates `class_name` from LangChain document loaders
         """
+        validate_class_name(class_name)
+
         lc_document_loader = None
 
         # Try importing from langchain_community first (most common document loaders)
diff --git a/python/beeai_framework/adapters/langchain/backend/text_splitter.py b/python/beeai_framework/adapters/langchain/backend/text_splitter.py
@@ -16,6 +16,7 @@
 from beeai_framework.adapters.langchain.mappers.documents import document_to_lc_document, lc_document_to_document
 from beeai_framework.backend.text_splitter import TextSplitter
 from beeai_framework.backend.types import Document
+from beeai_framework.utils.strings import validate_class_name
 from beeai_framework.logger import Logger
 
 logger = Logger(__name__)
@@ -31,6 +32,8 @@ def _class_from_name(cls, class_name: str, **kwargs: Any) -> LangChainTextSplitt
         """
         Dynamically imports and instantiates `class_name` from LangChain text splitters
         """
+        validate_class_name(class_name)
+
         lc_text_splitter = None
 
         # Try importing from langchain_text_splitters first (most common text splitters)
diff --git a/python/beeai_framework/adapters/langchain/backend/vector_store.py b/python/beeai_framework/adapters/langchain/backend/vector_store.py
@@ -20,6 +20,7 @@
 from beeai_framework.adapters.langchain.mappers.documents import document_to_lc_document, lc_document_to_document
 from beeai_framework.backend.embedding import EmbeddingModel
 from beeai_framework.backend.vector_store import QueryLike, VectorStore
+from beeai_framework.utils.strings import validate_class_name
 from beeai_framework.logger import Logger
 
 logger = Logger(__name__)
@@ -50,6 +51,8 @@ def _class_from_name(cls, class_name: str, embedding_model: EmbeddingModel, **kw
         """
         Dynamically imports and instantiates `class_name` from all vector store paths on LangChain
         """
+        validate_class_name(class_name)
+
         lc_vector_store = None
         # Convert BeeAI embedding to LangChain embedding
         lc_embedding = LangChainBeeAIEmbeddingModel(embedding_model)
diff --git a/python/beeai_framework/backend/document_loader.py b/python/beeai_framework/backend/document_loader.py
@@ -8,6 +8,7 @@
 
 from beeai_framework.backend.types import Document
 from beeai_framework.backend.utils import load_module, parse_module
+from beeai_framework.utils.strings import validate_class_name
 
 __all__ = ["DocumentLoader"]
 
@@ -85,6 +86,7 @@ def from_name(cls, name: str, **kwargs: Any) -> DocumentLoader:
             If the specified class cannot be found in any known integration package.
         """
         parsed_module = parse_module(name)
+        validate_class_name(parsed_module.entity_id)
         TargetDocumentLoader = load_module(parsed_module.provider_id, "document_loader")  # type: ignore # noqa: N806
         return TargetDocumentLoader._class_from_name(  # type: ignore[no-any-return]
             class_name=parsed_module.entity_id, **kwargs
diff --git a/python/beeai_framework/backend/document_processor.py b/python/beeai_framework/backend/document_processor.py
@@ -8,6 +8,7 @@
 
 from beeai_framework.backend.types import DocumentWithScore
 from beeai_framework.backend.utils import load_module, parse_module
+from beeai_framework.utils.strings import validate_class_name
 
 __all__ = ["DocumentProcessor"]
 
@@ -44,6 +45,7 @@ def from_name(cls, name: str, **kwargs: Any) -> DocumentProcessor:
             If the specified class cannot be found in any known integration package.
         """
         parsed_module = parse_module(name)
+        validate_class_name(parsed_module.entity_id)
         TargetDocumentProcessor = load_module(parsed_module.provider_id, "document_processor")  # type: ignore # noqa: N806
         return TargetDocumentProcessor._class_from_name(  # type: ignore[no-any-return]
             class_name=parsed_module.entity_id, **kwargs
diff --git a/python/beeai_framework/backend/text_splitter.py b/python/beeai_framework/backend/text_splitter.py
@@ -8,6 +8,7 @@
 
 from beeai_framework.backend.types import Document
 from beeai_framework.backend.utils import load_module, parse_module
+from beeai_framework.utils.strings import validate_class_name
 
 __all__ = ["TextSplitter"]
 
@@ -82,6 +83,8 @@ def from_name(cls, name: str, **kwargs: Any) -> TextSplitter:
                 f"Only provider {parsed_module.provider_id} was specified. Text Splitter name was not specified."
             )
 
+        validate_class_name(parsed_module.entity_id)
+
         target: type[TextSplitter] = load_module(parsed_module.provider_id, "text_splitter")
         return target._class_from_name(class_name=parsed_module.entity_id, **kwargs)
 
diff --git a/python/beeai_framework/backend/vector_store.py b/python/beeai_framework/backend/vector_store.py
@@ -9,6 +9,7 @@
 from beeai_framework.backend.embedding import EmbeddingModel
 from beeai_framework.backend.types import Document, DocumentWithScore
 from beeai_framework.backend.utils import load_module, parse_module
+from beeai_framework.utils.strings import validate_class_name
 
 
 class QueryLike(Protocol):
@@ -104,6 +105,8 @@ def from_name(cls, name: str, *, embedding_model: EmbeddingModel, **kwargs: Any)
                 f"Only provider {parsed_module.provider_id} was specified. Vector Store name was not specified."
             )
 
+        validate_class_name(parsed_module.entity_id)
+
         target: type[VectorStore] = load_module(parsed_module.provider_id, "vector_store")
         return target._class_from_name(
             class_name=parsed_module.entity_id,
diff --git a/python/beeai_framework/utils/strings.py b/python/beeai_framework/utils/strings.py
@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import json
+import keyword
 import random
 import re
 import string
@@ -161,3 +162,11 @@ def is_valid_unicode_escape_sequence(s: str) -> bool:
         return True
     except UnicodeError:
         return False
+
+
+def validate_class_name(class_name: str) -> None:
+    """Validate that a class name is a valid Python identifier and not a keyword to prevent injection."""
+    if not class_name or not class_name.isidentifier() or keyword.iskeyword(class_name):
+        raise ValueError(
+            f"Invalid class name: {class_name}. Must be a valid Python identifier and not a reserved keyword."
+        )
diff --git a/python/tests/utils/test_strings.py b/python/tests/utils/test_strings.py
@@ -0,0 +1,48 @@
+# Copyright 2025 © BeeAI a Series of LF Projects, LLC
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+from beeai_framework.utils.strings import validate_class_name
+
+
+def test_validate_class_name_valid():
+    # Valid identifiers
+    validate_class_name("MyClass")
+    validate_class_name("my_function")
+    validate_class_name("_InternalClass")
+    validate_class_name("Class123")
+
+
+def test_validate_class_name_invalid():
+    # Invalid identifiers
+    invalid_names = [
+        "",
+        "123Class",
+        "My-Class",
+        "My Class",
+        "My.Class",
+        "class",  # reserved word is allowed as identifier in isidentifier() but we might want to be stricter?
+                  # actually isidentifier() returns True for keywords.
+                  # but you can't import a class named 'class'.
+        "import",
+        "def",
+        "!",
+        "@",
+        "#",
+        "import os; os.system('ls')",
+        "__import__('os').system('ls')",
+    ]
+    for name in invalid_names:
+        # Some of these are actually valid identifiers (like 'class', 'import')
+        # but most are not.
+        if not name.isidentifier() or not name:
+            with pytest.raises(ValueError, match="Invalid class name"):
+                validate_class_name(name)
+
+def test_validate_class_name_keywords():
+    # Python keywords are identifiers but we shouldn't allow them as class names for safety?
+    # Actually, the main goal is to prevent module probing or injection like 'module.submodule'.
+    # A single keyword like 'class' is not an injection in the same way.
+    # But let's see if we want to block them.
+    # For now, validate_class_name only uses isidentifier().
+    pass
