Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 36 additions & 0 deletions api/core/rag/embedding/embedding_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,3 +31,39 @@ async def aembed_documents(self, texts: list[str]) -> list[list[float]]:
async def aembed_query(self, text: str) -> list[float]:
"""Asynchronous Embed query text."""
raise NotImplementedError

def validate_embedding_input(self, texts: list[str]) -> None:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Instance method without `self` usage wastes resources


The validate_embedding_input method is defined as an instance method but does not utilize the self parameter, so it unnecessarily creates a bound method for every class instance. This wastes memory and adds minor computational overhead during method calls.

Decorate validate_embedding_input with @staticmethod to avoid binding it to instances. This improves efficiency and clarifies that the method does not depend on instance state.

"""
Validate embedding input texts.

:param texts: list of texts
"""
if not texts:
raise ValueError("Texts list cannot be empty")

for text in texts:
if not text or len(text.strip()) == 0:
raise ValueError("All texts must be non-empty")

def compute_embedding_similarity(self, embedding1: list[float], embedding2: list[float]) -> float:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Instance method without `self` usage wastes resources


The method compute_embedding_similarity is defined as an instance method but does not use the self parameter. Python creates a bound method for every class instance, requiring additional memory and CPU cycles when this is not needed.

Add the @staticmethod decorator above compute_embedding_similarity to remove the implicit self and avoid bound method instantiation, improving resource efficiency for all class instances.

"""
Compute cosine similarity between two embeddings.

:param embedding1: first embedding
:param embedding2: second embedding
:return: similarity score
"""
if not embedding1 or not embedding2:
return 0.0

if len(embedding1) != len(embedding2):
raise ValueError("Embeddings must have same length")

dot_product = sum(a * b for a, b in zip(embedding1, embedding2))
norm1 = sum(a * a for a in embedding1) ** 0.5
norm2 = sum(b * b for b in embedding2) ** 0.5

if norm1 == 0 or norm2 == 0:
return 0.0

return dot_product / (norm1 * norm2)
29 changes: 29 additions & 0 deletions api/core/rag/index_processor/index_processor_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,35 @@ def index(self, dataset: Dataset, document: DatasetDocument, chunks: Any):
def format_preview(self, chunks: Any) -> Mapping[str, Any]:
raise NotImplementedError

def validate_index_params(self, dataset: Dataset, top_k: int, score_threshold: float) -> None:
"""
Validate indexing parameters.

:param dataset: Dataset object
:param top_k: Top K value
:param score_threshold: Score threshold
"""
if not dataset:
raise ValueError("Dataset cannot be None")

if top_k <= 0:
raise ValueError("Top K must be positive")

if score_threshold < 0.0 or score_threshold > 1.0:
raise ValueError("Score threshold must be between 0.0 and 1.0")

def estimate_index_size(self, documents: list[Document]) -> int:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Method lacks `@staticmethod`, causing unnecessary binding


The method estimate_index_size defined on line 110 does not use the self parameter within its body, meaning it does not need instance access. This causes Python to create a bound method every time an instance is made, incurring unnecessary memory overhead.

Add the @staticmethod decorator to estimate_index_size to avoid binding it to class instances, improving performance and reducing memory usage.

"""
Estimate the total size of documents for indexing.

:param documents: List of documents
:return: Estimated size in characters
"""
if not documents:
return 0

return sum(len(doc.page_content) for doc in documents if doc.page_content)

@abstractmethod
def retrieve(
self,
Expand Down
28 changes: 27 additions & 1 deletion api/core/rag/models/document.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,8 +50,34 @@ class Document(BaseModel):

attachments: list[AttachmentDocument] | None = None

def validate_document(self) -> None:
"""
Validate the document content and metadata.

:raises ValueError: If validation fails
"""
if not self.page_content or len(self.page_content.strip()) == 0:
raise ValueError("Document page content cannot be empty")

if self.metadata and not isinstance(self.metadata, dict):
raise ValueError("Metadata must be a dictionary")

class GeneralChunk(BaseModel):
def compute_content_length(self) -> int:
"""
Compute the length of the document content.

:return: Content length
"""
return len(self.page_content) if self.page_content else 0

def get_metadata_value(self, key: str) -> Any:
"""
Get a value from metadata.

:param key: Metadata key
:return: Value or None
"""
return self.metadata.get(key) if self.metadata else None
"""

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

String statement has no effect


The string statement has not been assigned to anything. This is pointless and should be removed if not necessary. In case this is supposed to describe what's happening in the code, it is recommended to use comments or docstrings instead.

General Chunk.
"""
Expand Down
28 changes: 28 additions & 0 deletions api/core/rag/pipeline_validator.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
from typing import Any, Dict, List
from core.rag.entities.context_entities import DocumentContext

class RagPipelineValidator:
@staticmethod
def validate_pipeline_config(config: Dict[str, Any]) -> None:
if not config:
raise ValueError("Pipeline config cannot be empty")

if 'top_k' in config and config['top_k'] <= 0:
raise ValueError("Top K must be positive")

if 'score_threshold' in config and (config['score_threshold'] < 0 or config['score_threshold'] > 1):
raise ValueError("Score threshold must be between 0 and 1")

@staticmethod
def compute_document_similarity(doc1: DocumentContext, doc2: DocumentContext) -> float:
if not doc1 or not doc2:
return 0.0

# Simple similarity based on content length difference
len1 = len(doc1.content) if doc1.content else 0
len2 = len(doc2.content) if doc2.content else 0
return 1.0 / (1.0 + abs(len1 - len2))
Comment on lines +17 to +24

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

`compute_document_similarity` uses only content length difference


The compute_document_similarity function provides a misleading implementation of document similarity. It only compares the lengths of document contents, which is not a reliable measure of semantic similarity and can lead to incorrect logic if a caller expects a meaningful comparison.

Consider replacing this with a standard text similarity algorithm like Jaccard similarity or cosine similarity on TF-IDF vectors. If this simple length-based metric is intentional, rename the function to compute_similarity_by_length to accurately reflect its behavior.

Comment on lines +17 to +24

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

`compute_document_similarity` compares document lengths, not content


The compute_document_similarity function's name is misleading as it suggests a content-based comparison, while its logic is solely based on the difference in content length. This can lead to incorrect behavior in features that rely on this function for semantic similarity.

To fix this, either rename the function to more accurately reflect its behavior (e.g., compute_length_based_similarity), or replace the current implementation with a proper content similarity algorithm like Jaccard similarity.


@staticmethod
def filter_high_quality_documents(documents: List[DocumentContext], min_quality: float) -> List[DocumentContext]:
return [doc for doc in documents if doc.score >= min_quality]
34 changes: 34 additions & 0 deletions api/core/rag/rerank/rerank_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,3 +25,37 @@ def run(
:return:
"""
raise NotImplementedError

def validate_rerank_params(self, query: str, documents: list[Document], top_n: int | None) -> None:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Instance method without `self` usage wastes memory


The method validate_rerank_params is defined with self but does not use it, causing Python to create bound method objects for each instance. This wastes memory and adds overhead in method calls.

Decorate validate_rerank_params with @staticmethod to avoid binding to class instances and improve performance.

"""
Validate rerank parameters.

:param query: search query
:param documents: documents
:param top_n: top n
"""
if not query:
raise ValueError("Query cannot be empty")

if not documents:
raise ValueError("Documents list cannot be empty")

if top_n is not None and top_n <= 0:
raise ValueError("Top N must be positive")

def compute_rerank_score(self, query: str, document: Document) -> float:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Instance method ignores `self`, wasting binding overhead


The method compute_rerank_score is defined as an instance method but does not reference self. This causes Python to create a bound method for each instance, increasing memory and computation unnecessarily.

Add the @staticmethod decorator before the method to define it as a static method, avoiding binding overhead and clarifying its usage.

"""
Compute rerank score for a document.

:param query: search query
:param document: document
:return: score
"""
if not query or not document:
return 0.0

# Simple scoring based on term overlap
query_terms = set(query.lower().split())
doc_terms = set(document.page_content.lower().split()) if document.page_content else set()
overlap = len(query_terms & doc_terms)
return overlap / len(query_terms) if query_terms else 0.0
97 changes: 97 additions & 0 deletions api/core/rag/retrieval/dataset_retrieval.py
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,97 @@ def _record_usage(self, usage: LLMUsage | None) -> None:
else:
self._llm_usage = self._llm_usage.plus(usage)

def validate_retrieval_config(self, config: DatasetEntity) -> None:
"""
Validate the dataset retrieval configuration for enhanced RAG pipeline reliability.

:param config: The dataset configuration to validate
:raises ValueError: If configuration is invalid
"""
if not config:
raise ValueError("Dataset configuration cannot be None")

if not config.dataset_ids:
raise ValueError("At least one dataset ID must be provided")

retrieve_config = config.retrieve_config
if not retrieve_config:
raise ValueError("Retrieve configuration is required")

# Validate top_k is reasonable
if retrieve_config.top_k <= 0:
raise ValueError("Top K must be greater than 0")

if retrieve_config.top_k > 100:
raise ValueError("Top K cannot exceed 100 to prevent excessive resource usage")

# Validate score threshold if enabled
if retrieve_config.score_threshold_enabled:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nested `if` statements can be combined using `and`


The code contains nested if statements that can be merged into a single condition using the and operator, reducing unnecessary nesting and simplifying logic flows. This enhances readability by making it clear all conditions must be true for the block to execute.

Merge the nested if conditions into one line using the and operator to combine expressions into a single if statement.

if retrieve_config.score_threshold < 0.0 or retrieve_config.score_threshold > 1.0:
raise ValueError("Score threshold must be between 0.0 and 1.0")

# Validate reranking configuration
if retrieve_config.reranking_enable:
rerank_model = retrieve_config.reranking_model
if not rerank_model.reranking_provider_name or not rerank_model.reranking_model_name:
raise ValueError("Reranking model must be fully specified when reranking is enabled")

def _calculate_relevance_score(self, documents: list[Document], query: str) -> float:
"""
Calculate average relevance score for documents.

:param documents: List of documents
:param query: Query string
:return: Average score
"""
if not documents:
return 0.0

total_score = 0.0
for doc in documents:
# Simple scoring based on query term presence
score = len(query.split()) / len(doc.page_content.split()) if doc.page_content else 0.0

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

`len(doc.page_content.split())` can be zero, causing `ZeroDivisionError`


The code len(doc.page_content.split()) is used as a divisor without checking if it's non-zero. If doc.page_content is a string containing only whitespace (e.g., " "), split() returns an empty list, which will cause a ZeroDivisionError and crash the request, leading to a potential denial-of-service.

To fix this, add a check to ensure doc.page_content contains non-whitespace characters before performing the division.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

`len(doc.page_content.split())` can be zero, causing a crash


The scoring logic len(query.split()) / len(doc.page_content.split()) is vulnerable to a division-by-zero error. If doc.page_content is a string that results in an empty list when split (e.g., an empty string or only whitespace), the denominator becomes zero. This leads to an unhandled ZeroDivisionError and a potential denial of service.

To fix this, add a check to ensure the list of words from doc.page_content is not empty before performing the division. Assign a score of 0.0 in cases where there is no content.

total_score += score

return total_score / len(documents)

def _validate_query_input(self, query: str, inputs: Mapping[str, Any] | None) -> None:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Instance method unused, wastes bound method overhead


The method _validate_query_input defines self as its first parameter but does not use it anywhere, creating an unneeded bound method. This causes Python to allocate extra resources during instantiation, affecting performance when creating class instances.

Add the @staticmethod decorator before _validate_query_input to remove the implicit self parameter, improving memory usage and call overhead for this utility method.

"""
Validate query and inputs for retrieval.

:param query: Query string
:param inputs: Additional inputs
"""
if not query or len(query.strip()) == 0:
raise ValueError("Query cannot be empty")

if inputs:
for key, value in inputs.items():
if value is None:
raise ValueError(f"Input '{key}' cannot be None")

def _filter_documents_by_score(self, documents: list[Document], min_score: float) -> list[Document]:
"""
Filter documents by minimum score.

:param documents: List of documents
:param min_score: Minimum score threshold
:return: Filtered documents
"""
return [doc for doc in documents if doc.score >= min_score]

def _compute_average_score(self, documents: list[Document]) -> float:
"""
Compute average score of documents.

:param documents: List of documents
:return: Average score
"""
if not documents:
return 0.0
total = sum(doc.score for doc in documents)
return total / len(documents)

def retrieve(
self,
app_id: str,
Expand Down Expand Up @@ -123,6 +214,12 @@ def retrieve(
:param inputs: inputs
:return:
"""
# Validate configuration for enhanced reliability
self.validate_retrieval_config(config)

# Validate query inputs
self._validate_query_input(query, inputs)

dataset_ids = config.dataset_ids
if len(dataset_ids) == 0:
return None, []
Expand Down
40 changes: 40 additions & 0 deletions api/core/rag/splitter/text_splitter.py
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,46 @@ def _merge_splits(self, splits: Iterable[str], separator: str, lengths: list[int
docs.append(doc)
return docs

def validate_split_params(self, text: str, chunk_size: int, chunk_overlap: int) -> None:
"""
Validate text splitting parameters.

:param text: text to split
:param chunk_size: chunk size
:param chunk_overlap: chunk overlap
"""
if not text:
raise ValueError("Text cannot be empty")

if chunk_size <= 0:
raise ValueError("Chunk size must be positive")

if chunk_overlap < 0:
raise ValueError("Chunk overlap cannot be negative")

if chunk_overlap >= chunk_size:
raise ValueError("Chunk overlap must be less than chunk size")

def estimate_chunks_count(self, text: str) -> int:
"""
Estimate the number of chunks for given text.

:param text: text to split
:return: estimated chunk count
"""
if not text:
return 0

text_length = len(text)
chunk_size = self._chunk_size
overlap = self._chunk_overlap

if text_length <= chunk_size:
return 1

effective_size = chunk_size - overlap
return (text_length - overlap) // effective_size + 1
Comment on lines +176 to +177

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

`chunk_size - overlap` can be zero, causing `ZeroDivisionError`


The calculation of effective_size does not handle the case where self._chunk_overlap is greater than or equal to self._chunk_size. This can lead to a ZeroDivisionError if effective_size is zero, crashing the application. An attacker might exploit this by providing configuration that leads to this state, causing a denial of service.

Add a validation check before the division to ensure effective_size is positive. For example, raise a ValueError if chunk_size &lt;= overlap.


@classmethod
def from_huggingface_tokenizer(cls, tokenizer: Any, **kwargs: Any) -> TextSplitter:
"""Text splitter that uses HuggingFace tokenizer to count length."""
Expand Down
47 changes: 41 additions & 6 deletions api/services/dataset_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -4024,11 +4024,46 @@
if set(local_member_list) != set(request_member_list):
raise ValueError("Dataset operators cannot change the dataset permissions.")

@classmethod
def clear_partial_member_list(cls, dataset_id):
try:
db.session.query(DatasetPermission).where(DatasetPermission.dataset_id == dataset_id).delete()
db.session.commit()
except Exception as e:
db.session.rollback()
raise e

Check failure on line 4028 in api/services/dataset_service.py

View workflow job for this annotation

GitHub Actions / Style Check / Python Style

"e" is not defined (reportUndefinedVariable)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Undefined variable 'e'


The variable name is not defined where it is used.
This will lead to an error during the runtime.
Make sure there is no typo. If the name was supposed to be imported, verify that you've actually imported the name.



def validate_dataset_config(dataset_config: dict) -> None:
"""
Validate dataset configuration for RAG.

:param dataset_config: dataset config dict
"""
if not dataset_config:
raise ValueError("Dataset config cannot be empty")

if 'retrieval_model' in dataset_config:
retrieval = dataset_config['retrieval_model']
if retrieval.get('top_k', 0) <= 0:
raise ValueError("Top K must be positive")

if retrieval.get('score_threshold', 0) < 0 or retrieval.get('score_threshold', 0) > 1:
raise ValueError("Score threshold must be between 0 and 1")


def compute_dataset_stats(dataset_id: str) -> dict:
"""
Compute statistics for a dataset.

:param dataset_id: dataset ID
:return: stats dict
"""
if not dataset_id:
raise ValueError("Dataset ID cannot be empty")

# Query document count
document_count = db.session.query(func.count(Document.id)).where(Document.dataset_id == dataset_id).scalar()

# Query segment count
segment_count = db.session.query(func.count(DocumentSegment.id)).join(Document).where(Document.dataset_id == dataset_id).scalar()

return {
'document_count': document_count,
'segment_count': segment_count,
'avg_segments_per_doc': segment_count / document_count if document_count > 0 else 0
}
Loading
Loading