retrieval.py

Overview

The retrieval.py file defines a Retrieval tool component designed for performing relevant content searches over datasets (knowledge bases). It provides a structured and configurable way to query indexed datasets with similarity-based retrieval and optional reranking and knowledge graph (KG) based augmentation.

This file is part of the InfiniFlow project and integrates with various system components such as knowledge base services, embedding models, and language model bundles to execute semantic search queries. The Retrieval tool supports features like:

Configurable similarity thresholds and search parameters.
Support for multiple knowledge bases and embedding models.
Optional reranking of search results using specialized models.
Cross-language query transformation.
Integration with knowledge graph retrieval.
Producing formalized content outputs suitable for downstream consumption.

Classes

`RetrievalParam`

Defines the parameters for the Retrieval tool component, including metadata, configuration options, and validation logic.

Attributes

meta: ToolMeta
Metadata about the tool including its name, description, and input parameters.
function_name: str
The function name identifier for the tool, fixed as "search_my_dateset".
description: str
A human-readable description of the tool's purpose.
similarity_threshold: float
Threshold for similarity score to consider a dataset chunk relevant. Default: 0.2.
keywords_similarity_weight: float
Weighting factor for keyword similarity during ranking. Default: 0.5.
top_n: int
Number of top results to return. Default: 8.
top_k: int
Upper bound on the number of candidates to consider during retrieval. Default: 1024.
kb_ids: list[str]
List of knowledge base IDs or variable references to search.
kb_vars: list
Variables related to knowledge bases (usage context-specific).
rerank_id: str
Optional rerank model ID for reranking retrieved results.
empty_response: str
Response to return when no relevant content is found.
use_kg: bool
Flag to enable augmentation via knowledge graph retrieval.
cross_languages: list
List of languages for cross-language query transformation.

Methods

init()
Initializes default parameter values and metadata.
check()
Validates parameter values such as threshold ranges and positive integers.
Raises exceptions if checks fail.
get_input_form() -> dict[str, dict]
Returns the expected input form schema, primarily the query string input.

Usage Example

params = RetrievalParam()
params.similarity_threshold = 0.3
params.check()
input_form = params.get_input_form()
print(input_form)
# Output: {'query': {'name': 'Query', 'type': 'line'}}

`Retrieval`

The core retrieval tool implementing search logic over knowledge bases by leveraging embedding models, reranking, and knowledge graph retrieval.

This class inherits from ToolBase and ABC (abstract base class) to integrate with the InfiniFlow agent framework.

Class Attributes

component_name: str
Static name identifier for the component, "Retrieval".

Methods

`_invoke(self, **kwargs)`

Primary method invoked to perform retrieval based on the input query.

Parameters:
- kwargs (dict): Expects a key "query" containing the search keywords string.
Returns:
- str — The formatted content string resulting from retrieval, or the empty response if no content is found.
Raises:
- Exception if no datasets are selected or specified datasets do not exist.
Description:
1. Validates presence of the query parameter; if missing, sets output as empty response.
2. Resolves knowledge base IDs from _param.kb_ids, supporting direct IDs or variable references from the canvas.
3. Fetches knowledge base objects corresponding to IDs; throws if none found.
4. Ensures all knowledge bases use the same embedding model.
5. Initializes embedding model bundle (LLMBundle) based on embedding model name.
6. Optionally initializes rerank model bundle if specified.
7. Formats the query string and applies cross-language transformation if enabled.
8. Executes the retrieval using settings.retrievaler.retrieval() with configured parameters.
9. If use_kg is enabled, performs additional knowledge graph retrieval augmentations.
10. Cleans up retrieved chunk metadata by removing heavy or unnecessary fields (vectors, long tokens).
11. If no chunks found, sets output to empty response. Otherwise:
  - Adds references to the canvas.
  - Formats content using kb_prompt.
  - Sets formatted content as output.
12. Returns the formatted content string.
Implementation Details:
- Uses a timeout decorator to limit execution time (default 12 seconds, configurable via environment variable).
- Uses regex to clean certain user prefixes from the query.
- Retrieves ranking features by labeling the question with label_question().
- Integrates tightly with InfiniFlow's KnowledgebaseService, LLMBundle, and retrievaler settings.
Usage Example:

retrieval_tool = Retrieval()
retrieval_tool._param.kb_ids = ["my_knowledgebase"]
retrieval_tool._param.top_n = 5
result = retrieval_tool._invoke(query="Explain semantic search")
print(result)

`thoughts(self) -> str`

Returns an introspective string describing the retrieval intent.

Returns:
- str — A multi-line string indicating the keywords being searched and the high-level goal.
Usage Example:

print(retrieval_tool.thoughts())
# Output:
# Keywords: Explain semantic search
# Looking for the most relevant articles.

Important Implementation Details and Algorithms

Embedding Model Consistency:
The retrieval ensures all selected knowledge bases use the same embedding model to maintain vector space consistency.
Reranking:
Optionally applies a rerank model (LLMBundle with LLMType.RERANK) to reorder retrieved results based on more complex semantic understanding.
Cross-Language Querying:
Supports multi-language queries and translation-like transformations via cross_languages.
Knowledge Graph Augmentation:
When enabled, performs additional retrieval from a knowledge graph retrieval component (settings.kg_retrievaler) to enhance results.
Result Cleaning:
Removes heavy vector and token-related fields from chunks before output to optimize response size.
Timeout Enforcement:
Uses a decorator to abort retrieval if it exceeds configured timeout, enhancing system robustness.

Interaction with Other System Components

KnowledgebaseService
Used to fetch knowledge base metadata by name or ID.
LLMBundle
Manages language model bundles for embedding and reranking.
settings.retrievaler and settings.kg_retrievaler
External configured retrieval services performing the actual search logic over embeddings and knowledge graphs.
_canvas (Context from ToolBase)
Provides access to tenant ID, variables, and reference management for the current workflow execution.
kb_prompt and label_question
Utility functions to format retrieved content and generate ranking features.
cross_languages
Handles multilingual query transformations.

Mermaid Class Diagram

classDiagram
    class RetrievalParam {
        +meta: ToolMeta
        +function_name: str
        +description: str
        +similarity_threshold: float
        +keywords_similarity_weight: float
        +top_n: int
        +top_k: int
        +kb_ids: list~str~
        +kb_vars: list
        +rerank_id: str
        +empty_response: str
        +use_kg: bool
        +cross_languages: list
        +__init__()
        +check()
        +get_input_form() dict~str, dict~
    }

    class Retrieval {
        +component_name: str
        +_invoke(kwargs)
        +thoughts() str
    }

    RetrievalParam <|-- Retrieval
    Retrieval ..|> ToolBase
    Retrieval ..|> ABC

Summary

retrieval.py provides a flexible, configurable retrieval tool component for semantic search over knowledge bases within the InfiniFlow platform. It leverages embeddings, reranking, and optionally knowledge graph data to produce relevant, formalized content based on user queries. The file encapsulates parameter definitions, input validation, core retrieval logic, and integration with external services and settings. The modular design allows it to be easily extended or configured for different datasets and retrieval strategies.