dify_retrieval.py

Overview

The dify_retrieval.py file implements an API endpoint for performing knowledge retrieval queries within the InfiniFlow system. It provides a RESTful POST interface /dify/retrieval that accepts a query, knowledge base identifier, and various retrieval parameters to return relevant document chunks ranked by semantic similarity.

This retrieval endpoint supports filtering based on metadata conditions, applies similarity thresholds, and optionally integrates knowledge graph (KG) based retrieval results. It leverages embedding models and ranking utilities to return relevant content snippets from a specified knowledge base, facilitating advanced question answering and information retrieval capabilities in tenant-scoped environments.

Detailed Description

API Endpoint: `/dify/retrieval`

HTTP Method: POST
Decorators:
- @apikey_required: Ensures that the request includes a valid API key for authorization.
- @validate_request("knowledge_id", "query"): Validates that the JSON payload contains required fields knowledge_id and query.
Parameters:
- The endpoint expects a JSON body with the following fields:
  - knowledge_id (str): Identifier of the knowledge base to query.
  - query (str): The question or search query string.
  - use_kg (bool, optional): Whether to augment results with knowledge graph retrieval (default False).
  - retrieval_setting (dict, optional): Retrieval parameters including:
    - score_threshold (float, optional): Minimum similarity score to consider.
    - top_k (int, optional): Number of top results to return.
  - metadata_condition (dict, optional): Metadata filters to constrain the search scope.
Returns:
- JSON response containing a list of document chunks with the following structure per record:
  - content: The retrieved text snippet, potentially weighted or annotated.
  - score: Similarity score of the chunk to the query.
  - title: Document title or keyword associated with the chunk.
  - metadata: Metadata fields associated with the document chunk.
Error Handling:
- Returns error JSON with appropriate HTTP status codes if knowledge base is not found or no chunks match.
- Logs exceptions and returns server error responses if unexpected errors occur.

Functions and Methods

`retrieval(tenant_id)`

The main view function handling the /dify/retrieval POST requests.

Parameters:

tenant_id (str): Tenant identifier implicitly passed by the route context (multi-tenant support).

Workflow:

Extracts and validates the request JSON.
Retrieves metadata for documents associated with the knowledge base.
Fetches the knowledge base object by ID.
Constructs an embedding model bundle for vector-based retrieval.
Converts metadata condition filters and applies them to filter document IDs.
Invokes the retrieval engine to get ranked document chunks based on semantic similarity to the query.
Optionally performs knowledge graph enhanced retrieval and merges results.
Constructs a response list with relevant document content, scores, titles, and metadata.
Returns the results as JSON.

Returns:

Flask jsonify response with retrieval results or error message.

Usage Example:

import requests

url = "https://api.example.com/dify/retrieval"
payload = {
    "knowledge_id": "kb123",
    "query": "What is the refund policy?",
    "use_kg": True,
    "retrieval_setting": {
        "score_threshold": 0.5,
        "top_k": 10
    },
    "metadata_condition": {
        "category": "policy"
    }
}
headers = {
    "apikey": "your_api_key_here"
}

response = requests.post(url, json=payload, headers=headers)
print(response.json())

Important Implementation Details and Algorithms

Metadata Filtering:
The system filters document candidates based on metadata conditions using meta_filter and convert_conditions. This ensures retrieval results conform to user-specified constraints.
Vector Similarity Retrieval:
Uses embedding models (LLMBundle) to convert queries and documents into vector space and ranks results by cosine similarity or other metrics.
The retrieval process considers:
- Similarity threshold (score_threshold)
- Number of top results (top_k)
- Weighting of vector similarity (defaulted to 0.3)
Knowledge Graph Integration:
If use_kg is enabled, additional retrieval occurs through a KG-based retrievaler, which may enrich or reorder the results by incorporating structured knowledge relationships.
Ranking with Feature Labeling:
The question is labeled using label_question to enhance ranking features, potentially improving relevance.
Error Handling:
Specific checks for "not_found" exceptions provide user-friendly error messages when no chunks are found or knowledge base is missing.

Interaction with Other System Components

Database Services:
- DocumentService: Fetches document metadata and content by document IDs.
- KnowledgebaseService: Retrieves knowledge base information.
- LLMBundle: Provides embedding models for vectorization.
- DialogService: Provides utilities (meta_filter, convert_conditions) for filtering documents.
Settings and Configuration:
- settings.retrievaler: The main retrieval engine instance.
- settings.kg_retrievaler: Optional knowledge graph retrieval engine.
- settings.RetCode: Standardized return codes for response consistency.
API Utilities:
- validate_request and apikey_required decorators enforce security and input validation.
- build_error_result standardizes error response formatting.
External Modules:
- rag.app.tag.label_question: Used to enhance ranking by labeling queries.

This file acts as the controller layer connecting HTTP requests to backend retrieval services and orchestrates multiple services to produce a unified retrieval result.

Class & Function Diagram

flowchart TD
    A[POST /dify/retrieval Endpoint]
    A -->|Requires API Key| B[apikey_required Decorator]
    A -->|Validates JSON Fields| C[validate_request Decorator]

    A --> D[retrieval(tenant_id) Function]
    D --> E[DocumentService.get_meta_by_kbs(kb_id)]
    D --> F[KnowledgebaseService.get_by_id(kb_id)]
    D --> G[LLMBundle Embedding Model]
    D --> H[meta_filter + convert_conditions(metadata_condition)]
    D --> I[settings.retrievaler.retrieval(query, ...)]
    D --> J[Optional: settings.kg_retrievaler.retrieval(...) if use_kg]
    D --> K[DocumentService.get_by_id(doc_id)]
    D --> L[build_error_result on errors]

    subgraph Retrieval Flow
        E --> H --> I --> J --> K
    end

Summary

The dify_retrieval.py file is a critical component in the InfiniFlow application that handles knowledge retrieval requests through a well-defined API endpoint. It integrates metadata filtering, vector similarity ranking, and optional knowledge graph augmentation to deliver relevant document chunks from tenant-specific knowledge bases. The file leverages multiple backend services and utility functions to validate inputs, process queries, and return structured JSON responses optimized for downstream consumption by client applications or other system modules.