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:
"""
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:
"""
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)
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` only compares content lengths, not content


The compute_document_similarity function bases its calculation only on the length of document content, which can be misleading. For example, two completely different documents are rated as identical if they have the same character count, potentially causing incorrect behavior in features relying on this score.

Rename the function to be more descriptive, like compute_length_based_similarity, or implement a content-based similarity algorithm such as Jaccard similarity to accurately reflect document content 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 validate_rerank_params method is defined as an instance method but does not reference the self parameter, causing Python to create a bound method for each class instance that is unnecessary. This increases memory usage and computational overhead.
Decorate the method with @staticmethod to avoid binding to class instances and optimize 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 without `self` usage wastes resources


The compute_rerank_score method is defined as an instance method with a self parameter but does not access any instance attributes or methods. This results in Python creating a bound method for each class instance, consuming additional memory and computation.

Decorate compute_rerank_score with @staticmethod to indicate it does not depend on instance state, eliminating the binding overhead and improving performance.

"""
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` can be merged using `and` for clarity


Nested if statements like if retrieve_config.score_threshold_enabled: can be collapsed into a single condition using the and operator when multiple conditions exist. This reduces nesting complexity and improves code clarity.

Merge the nested if conditions into a single line using and to enhance readability and maintainability.

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:

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` use wastes memory


The _calculate_relevance_score method doesn't use its self parameter or any instance data, resulting in a bound method created per object instance which is inefficient. This can increase memory use and call overhead unnecessarily.

Decorate _calculate_relevance_score with @staticmethod to avoid binding the method to instance objects and reduce memory and CPU overhead.

"""
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 a crash


The scoring logic at len(query.split()) / len(doc.page_content.split()) is vulnerable to a ZeroDivisionError. If doc.page_content contains only whitespace characters (e.g., " "), split() will produce an empty list, causing its length to be zero and triggering an unhandled exception, leading to a denial of service.

Add a check to ensure the length of doc.page_content.split() is not zero before performing the division. For example: page_content_words = doc.page_content.split(); score = len(query.split()) / len(page_content_words) if page_content_words else 0.0.

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 lacks `@staticmethod`, causing overhead


The method _validate_query_input does not reference the instance (self), which means it unnecessarily creates a bound method for each class instance, wasting memory and CPU cycles. This impacts performance especially if the class is instantiated frequently.

Add the @staticmethod decorator to _validate_query_input to avoid binding to the instance and improve efficiency without changing how the method is called.

"""
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]:

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, wasting memory and computation


Defining _filter_documents_by_score as an instance method incurs overhead by creating a bound method for each class instance without using instance state. This wastes memory and CPU resources.

Add the @staticmethod decorator to _filter_documents_by_score to define it as a static method, avoiding instance binding and improving performance.

"""
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:

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` uses more memory and CPU


The _compute_average_score method accepts self but does not use any instance attributes or methods, leading to unnecessary bound method creation for each class instance. This consumes extra memory and processing time at runtime.

Add the @staticmethod decorator to _compute_average_score and remove the self parameter to avoid this overhead and clarify that the method does not rely on instance state.

"""
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.

Division by `effective_size` which can be zero or negative


The calculation of effective_size does not handle cases where self._chunk_size is less than or equal to self._chunk_overlap. This can lead to a ZeroDivisionError if they are equal, or incorrect negative chunk counts if overlap is larger, which can crash the process and cause a denial of service.

Add a check to ensure effective_size is positive before the division. Raise a ValueError if chunk_size &lt;= overlap to prevent the application from crashing with an unhandled exception.


@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)

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