refactor: move the embedding to the rag module and abstract the rerank runner for extension (#9423)

api/core/rag/data_post_processor/data_post_processor.py

@@ -1,14 +1,14 @@
 from typing import Optional
 
-from core.model_manager import ModelManager
+from core.model_manager import ModelInstance, ModelManager
 from core.model_runtime.entities.model_entities import ModelType
 from core.model_runtime.errors.invoke import InvokeAuthorizationError
 from core.rag.data_post_processor.reorder import ReorderRunner
 from core.rag.models.document import Document
-from core.rag.rerank.constants.rerank_mode import RerankMode
 from core.rag.rerank.entity.weight import KeywordSetting, VectorSetting, Weights
-from core.rag.rerank.rerank_model import RerankModelRunner
-from core.rag.rerank.weight_rerank import WeightRerankRunner
+from core.rag.rerank.rerank_base import BaseRerankRunner
+from core.rag.rerank.rerank_factory import RerankRunnerFactory
+from core.rag.rerank.rerank_type import RerankMode
 
 
 class DataPostProcessor:
@@ -47,11 +47,12 @@ class DataPostProcessor:
         tenant_id: str,
         reranking_model: Optional[dict] = None,
         weights: Optional[dict] = None,
-    ) -> Optional[RerankModelRunner | WeightRerankRunner]:
+    ) -> Optional[BaseRerankRunner]:
         if reranking_mode == RerankMode.WEIGHTED_SCORE.value and weights:
-            return WeightRerankRunner(
-                tenant_id,
-                Weights(
+            runner = RerankRunnerFactory.create_rerank_runner(
+                runner_type=reranking_mode,
+                tenant_id=tenant_id,
+                weights=Weights(
                     vector_setting=VectorSetting(
                         vector_weight=weights["vector_setting"]["vector_weight"],
                         embedding_provider_name=weights["vector_setting"]["embedding_provider_name"],
@@ -62,23 +63,33 @@ class DataPostProcessor:
                     ),
                 ),
             )
+            return runner
         elif reranking_mode == RerankMode.RERANKING_MODEL.value:
-            if reranking_model:
-                try:
-                    model_manager = ModelManager()
-                    rerank_model_instance = model_manager.get_model_instance(
-                        tenant_id=tenant_id,
-                        provider=reranking_model["reranking_provider_name"],
-                        model_type=ModelType.RERANK,
-                        model=reranking_model["reranking_model_name"],
-                    )
-                except InvokeAuthorizationError:
-                    return None
-                return RerankModelRunner(rerank_model_instance)
-            return None
+            rerank_model_instance = self._get_rerank_model_instance(tenant_id, reranking_model)
+            if rerank_model_instance is None:
+                return None
+            runner = RerankRunnerFactory.create_rerank_runner(
+                runner_type=reranking_mode, rerank_model_instance=rerank_model_instance
+            )
+            return runner
         return None
 
     def _get_reorder_runner(self, reorder_enabled) -> Optional[ReorderRunner]:
         if reorder_enabled:
             return ReorderRunner()
         return None
+
+    def _get_rerank_model_instance(self, tenant_id: str, reranking_model: Optional[dict]) -> ModelInstance | None:
+        if reranking_model:
+            try:
+                model_manager = ModelManager()
+                rerank_model_instance = model_manager.get_model_instance(
+                    tenant_id=tenant_id,
+                    provider=reranking_model["reranking_provider_name"],
+                    model_type=ModelType.RERANK,
+                    model=reranking_model["reranking_model_name"],
+                )
+                return rerank_model_instance
+            except InvokeAuthorizationError:
+                return None
+        return None

api/core/rag/datasource/retrieval_service.py

@@ -6,7 +6,7 @@ from flask import Flask, current_app
 from core.rag.data_post_processor.data_post_processor import DataPostProcessor
 from core.rag.datasource.keyword.keyword_factory import Keyword
 from core.rag.datasource.vdb.vector_factory import Vector
-from core.rag.rerank.constants.rerank_mode import RerankMode
+from core.rag.rerank.rerank_type import RerankMode
 from core.rag.retrieval.retrieval_methods import RetrievalMethod
 from extensions.ext_database import db
 from models.dataset import Dataset

api/core/rag/datasource/vdb/analyticdb/analyticdb_vector.py

@@ -9,10 +9,10 @@ _import_err_msg = (
 )
 
 from configs import dify_config
-from core.rag.datasource.entity.embedding import Embeddings
 from core.rag.datasource.vdb.vector_base import BaseVector
 from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
 from core.rag.datasource.vdb.vector_type import VectorType
+from core.rag.embedding.embedding_base import Embeddings
 from core.rag.models.document import Document
 from extensions.ext_redis import redis_client
 from models.dataset import Dataset

api/core/rag/datasource/vdb/baidu/baidu_vector.py

@@ -12,10 +12,10 @@ from pymochow.model.schema import Field, HNSWParams, Schema, VectorIndex
 from pymochow.model.table import AnnSearch, HNSWSearchParams, Partition, Row
 
 from configs import dify_config
-from core.rag.datasource.entity.embedding import Embeddings
 from core.rag.datasource.vdb.vector_base import BaseVector
 from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
 from core.rag.datasource.vdb.vector_type import VectorType
+from core.rag.embedding.embedding_base import Embeddings
 from core.rag.models.document import Document
 from extensions.ext_redis import redis_client
 from models.dataset import Dataset

api/core/rag/datasource/vdb/chroma/chroma_vector.py

@@ -6,10 +6,10 @@ from chromadb import QueryResult, Settings
 from pydantic import BaseModel
 
 from configs import dify_config
-from core.rag.datasource.entity.embedding import Embeddings
 from core.rag.datasource.vdb.vector_base import BaseVector
 from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
 from core.rag.datasource.vdb.vector_type import VectorType
+from core.rag.embedding.embedding_base import Embeddings
 from core.rag.models.document import Document
 from extensions.ext_redis import redis_client
 from models.dataset import Dataset

api/core/rag/datasource/vdb/elasticsearch/elasticsearch_vector.py

@@ -9,11 +9,11 @@ from elasticsearch import Elasticsearch
 from flask import current_app
 from pydantic import BaseModel, model_validator
 
-from core.rag.datasource.entity.embedding import Embeddings
 from core.rag.datasource.vdb.field import Field
 from core.rag.datasource.vdb.vector_base import BaseVector
 from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
 from core.rag.datasource.vdb.vector_type import VectorType
+from core.rag.embedding.embedding_base import Embeddings
 from core.rag.models.document import Document
 from extensions.ext_redis import redis_client
 from models.dataset import Dataset

api/core/rag/datasource/vdb/milvus/milvus_vector.py

@@ -7,11 +7,11 @@ from pymilvus import MilvusClient, MilvusException
 from pymilvus.milvus_client import IndexParams
 
 from configs import dify_config
-from core.rag.datasource.entity.embedding import Embeddings
 from core.rag.datasource.vdb.field import Field
 from core.rag.datasource.vdb.vector_base import BaseVector
 from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
 from core.rag.datasource.vdb.vector_type import VectorType
+from core.rag.embedding.embedding_base import Embeddings
 from core.rag.models.document import Document
 from extensions.ext_redis import redis_client
 from models.dataset import Dataset

api/core/rag/datasource/vdb/myscale/myscale_vector.py

@@ -8,10 +8,10 @@ from clickhouse_connect import get_client
 from pydantic import BaseModel
 
 from configs import dify_config
-from core.rag.datasource.entity.embedding import Embeddings
 from core.rag.datasource.vdb.vector_base import BaseVector
 from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
 from core.rag.datasource.vdb.vector_type import VectorType
+from core.rag.embedding.embedding_base import Embeddings
 from core.rag.models.document import Document
 from models.dataset import Dataset
 

api/core/rag/datasource/vdb/opensearch/opensearch_vector.py

@@ -9,11 +9,11 @@ from opensearchpy.helpers import BulkIndexError
 from pydantic import BaseModel, model_validator
 
 from configs import dify_config
-from core.rag.datasource.entity.embedding import Embeddings
 from core.rag.datasource.vdb.field import Field
 from core.rag.datasource.vdb.vector_base import BaseVector
 from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
 from core.rag.datasource.vdb.vector_type import VectorType
+from core.rag.embedding.embedding_base import Embeddings
 from core.rag.models.document import Document
 from extensions.ext_redis import redis_client
 from models.dataset import Dataset

api/core/rag/datasource/vdb/oracle/oraclevector.py

@@ -13,10 +13,10 @@ from nltk.corpus import stopwords
 from pydantic import BaseModel, model_validator
 
 from configs import dify_config
-from core.rag.datasource.entity.embedding import Embeddings
 from core.rag.datasource.vdb.vector_base import BaseVector
 from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
 from core.rag.datasource.vdb.vector_type import VectorType
+from core.rag.embedding.embedding_base import Embeddings
 from core.rag.models.document import Document
 from extensions.ext_redis import redis_client
 from models.dataset import Dataset

api/core/rag/datasource/vdb/pgvecto_rs/pgvecto_rs.py

@@ -12,11 +12,11 @@ from sqlalchemy.dialects import postgresql
 from sqlalchemy.orm import Mapped, Session, mapped_column
 
 from configs import dify_config
-from core.rag.datasource.entity.embedding import Embeddings
 from core.rag.datasource.vdb.pgvecto_rs.collection import CollectionORM
 from core.rag.datasource.vdb.vector_base import BaseVector
 from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
 from core.rag.datasource.vdb.vector_type import VectorType
+from core.rag.embedding.embedding_base import Embeddings
 from core.rag.models.document import Document
 from extensions.ext_redis import redis_client
 from models.dataset import Dataset

api/core/rag/datasource/vdb/pgvector/pgvector.py

@@ -8,10 +8,10 @@ import psycopg2.pool
 from pydantic import BaseModel, model_validator
 
 from configs import dify_config
-from core.rag.datasource.entity.embedding import Embeddings
 from core.rag.datasource.vdb.vector_base import BaseVector
 from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
 from core.rag.datasource.vdb.vector_type import VectorType
+from core.rag.embedding.embedding_base import Embeddings
 from core.rag.models.document import Document
 from extensions.ext_redis import redis_client
 from models.dataset import Dataset

api/core/rag/datasource/vdb/qdrant/qdrant_vector.py

@@ -20,11 +20,11 @@ from qdrant_client.http.models import (
 from qdrant_client.local.qdrant_local import QdrantLocal
 
 from configs import dify_config
-from core.rag.datasource.entity.embedding import Embeddings
 from core.rag.datasource.vdb.field import Field
 from core.rag.datasource.vdb.vector_base import BaseVector
 from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
 from core.rag.datasource.vdb.vector_type import VectorType
+from core.rag.embedding.embedding_base import Embeddings
 from core.rag.models.document import Document
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client

api/core/rag/datasource/vdb/relyt/relyt_vector.py

@@ -8,9 +8,9 @@ from sqlalchemy import text as sql_text
 from sqlalchemy.dialects.postgresql import JSON, TEXT
 from sqlalchemy.orm import Session
 
-from core.rag.datasource.entity.embedding import Embeddings
 from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
 from core.rag.datasource.vdb.vector_type import VectorType
+from core.rag.embedding.embedding_base import Embeddings
 from models.dataset import Dataset
 
 try:

api/core/rag/datasource/vdb/tencent/tencent_vector.py

@@ -8,10 +8,10 @@ from tcvectordb.model import index as vdb_index
 from tcvectordb.model.document import Filter
 
 from configs import dify_config
-from core.rag.datasource.entity.embedding import Embeddings
 from core.rag.datasource.vdb.vector_base import BaseVector
 from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
 from core.rag.datasource.vdb.vector_type import VectorType
+from core.rag.embedding.embedding_base import Embeddings
 from core.rag.models.document import Document
 from extensions.ext_redis import redis_client
 from models.dataset import Dataset

api/core/rag/datasource/vdb/tidb_vector/tidb_vector.py

@@ -9,10 +9,10 @@ from sqlalchemy import text as sql_text
 from sqlalchemy.orm import Session, declarative_base
 
 from configs import dify_config
-from core.rag.datasource.entity.embedding import Embeddings
 from core.rag.datasource.vdb.vector_base import BaseVector
 from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
 from core.rag.datasource.vdb.vector_type import VectorType
+from core.rag.embedding.embedding_base import Embeddings
 from core.rag.models.document import Document
 from extensions.ext_redis import redis_client
 from models.dataset import Dataset

api/core/rag/datasource/vdb/vector_factory.py

@@ -2,12 +2,12 @@ from abc import ABC, abstractmethod
 from typing import Any, Optional
 
 from configs import dify_config
-from core.embedding.cached_embedding import CacheEmbedding
 from core.model_manager import ModelManager
 from core.model_runtime.entities.model_entities import ModelType
-from core.rag.datasource.entity.embedding import Embeddings
 from core.rag.datasource.vdb.vector_base import BaseVector
 from core.rag.datasource.vdb.vector_type import VectorType
+from core.rag.embedding.cached_embedding import CacheEmbedding
+from core.rag.embedding.embedding_base import Embeddings
 from core.rag.models.document import Document
 from extensions.ext_redis import redis_client
 from models.dataset import Dataset

api/core/rag/datasource/vdb/vikingdb/vikingdb_vector.py

@@ -14,11 +14,11 @@ from volcengine.viking_db import (
 )
 
 from configs import dify_config
-from core.rag.datasource.entity.embedding import Embeddings
 from core.rag.datasource.vdb.field import Field as vdb_Field
 from core.rag.datasource.vdb.vector_base import BaseVector
 from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
 from core.rag.datasource.vdb.vector_type import VectorType
+from core.rag.embedding.embedding_base import Embeddings
 from core.rag.models.document import Document
 from extensions.ext_redis import redis_client
 from models.dataset import Dataset

api/core/rag/datasource/vdb/weaviate/weaviate_vector.py

@@ -7,11 +7,11 @@ import weaviate
 from pydantic import BaseModel, model_validator
 
 from configs import dify_config
-from core.rag.datasource.entity.embedding import Embeddings
 from core.rag.datasource.vdb.field import Field
 from core.rag.datasource.vdb.vector_base import BaseVector
 from core.rag.datasource.vdb.vector_factory import AbstractVectorFactory
 from core.rag.datasource.vdb.vector_type import VectorType
+from core.rag.embedding.embedding_base import Embeddings
 from core.rag.models.document import Document
 from extensions.ext_redis import redis_client
 from models.dataset import Dataset

api/core/rag/embedding/cached_embedding.py

@@ -0,0 +1,132 @@
+import base64
+import logging
+from typing import Optional, cast
+
+import numpy as np
+from sqlalchemy.exc import IntegrityError
+
+from configs import dify_config
+from core.entities.embedding_type import EmbeddingInputType
+from core.model_manager import ModelInstance
+from core.model_runtime.entities.model_entities import ModelPropertyKey
+from core.model_runtime.model_providers.__base.text_embedding_model import TextEmbeddingModel
+from core.rag.embedding.embedding_base import Embeddings
+from extensions.ext_database import db
+from extensions.ext_redis import redis_client
+from libs import helper
+from models.dataset import Embedding
+
+logger = logging.getLogger(__name__)
+
+
+class CacheEmbedding(Embeddings):
+    def __init__(self, model_instance: ModelInstance, user: Optional[str] = None) -> None:
+        self._model_instance = model_instance
+        self._user = user
+
+    def embed_documents(self, texts: list[str]) -> list[list[float]]:
+        """Embed search docs in batches of 10."""
+        # use doc embedding cache or store if not exists
+        text_embeddings = [None for _ in range(len(texts))]
+        embedding_queue_indices = []
+        for i, text in enumerate(texts):
+            hash = helper.generate_text_hash(text)
+            embedding = (
+                db.session.query(Embedding)
+                .filter_by(
+                    model_name=self._model_instance.model, hash=hash, provider_name=self._model_instance.provider
+                )
+                .first()
+            )
+            if embedding:
+                text_embeddings[i] = embedding.get_embedding()
+            else:
+                embedding_queue_indices.append(i)
+        if embedding_queue_indices:
+            embedding_queue_texts = [texts[i] for i in embedding_queue_indices]
+            embedding_queue_embeddings = []
+            try:
+                model_type_instance = cast(TextEmbeddingModel, self._model_instance.model_type_instance)
+                model_schema = model_type_instance.get_model_schema(
+                    self._model_instance.model, self._model_instance.credentials
+                )
+                max_chunks = (
+                    model_schema.model_properties[ModelPropertyKey.MAX_CHUNKS]
+                    if model_schema and ModelPropertyKey.MAX_CHUNKS in model_schema.model_properties
+                    else 1
+                )
+                for i in range(0, len(embedding_queue_texts), max_chunks):
+                    batch_texts = embedding_queue_texts[i : i + max_chunks]
+
+                    embedding_result = self._model_instance.invoke_text_embedding(
+                        texts=batch_texts, user=self._user, input_type=EmbeddingInputType.DOCUMENT
+                    )
+
+                    for vector in embedding_result.embeddings:
+                        try:
+                            normalized_embedding = (vector / np.linalg.norm(vector)).tolist()
+                            embedding_queue_embeddings.append(normalized_embedding)
+                        except IntegrityError:
+                            db.session.rollback()
+                        except Exception as e:
+                            logging.exception("Failed transform embedding: %s", e)
+                cache_embeddings = []
+                try:
+                    for i, embedding in zip(embedding_queue_indices, embedding_queue_embeddings):
+                        text_embeddings[i] = embedding
+                        hash = helper.generate_text_hash(texts[i])
+                        if hash not in cache_embeddings:
+                            embedding_cache = Embedding(
+                                model_name=self._model_instance.model,
+                                hash=hash,
+                                provider_name=self._model_instance.provider,
+                            )
+                            embedding_cache.set_embedding(embedding)
+                            db.session.add(embedding_cache)
+                            cache_embeddings.append(hash)
+                    db.session.commit()
+                except IntegrityError:
+                    db.session.rollback()
+            except Exception as ex:
+                db.session.rollback()
+                logger.error("Failed to embed documents: %s", ex)
+                raise ex
+
+        return text_embeddings
+
+    def embed_query(self, text: str) -> list[float]:
+        """Embed query text."""
+        # use doc embedding cache or store if not exists
+        hash = helper.generate_text_hash(text)
+        embedding_cache_key = f"{self._model_instance.provider}_{self._model_instance.model}_{hash}"
+        embedding = redis_client.get(embedding_cache_key)
+        if embedding:
+            redis_client.expire(embedding_cache_key, 600)
+            return list(np.frombuffer(base64.b64decode(embedding), dtype="float"))
+        try:
+            embedding_result = self._model_instance.invoke_text_embedding(
+                texts=[text], user=self._user, input_type=EmbeddingInputType.QUERY
+            )
+
+            embedding_results = embedding_result.embeddings[0]
+            embedding_results = (embedding_results / np.linalg.norm(embedding_results)).tolist()
+        except Exception as ex:
+            if dify_config.DEBUG:
+                logging.exception(f"Failed to embed query text: {ex}")
+            raise ex
+
+        try:
+            # encode embedding to base64
+            embedding_vector = np.array(embedding_results)
+            vector_bytes = embedding_vector.tobytes()
+            # Transform to Base64
+            encoded_vector = base64.b64encode(vector_bytes)
+            # Transform to string
+            encoded_str = encoded_vector.decode("utf-8")
+            redis_client.setex(embedding_cache_key, 600, encoded_str)
+        except Exception as ex:
+            if dify_config.DEBUG:
+                logging.exception("Failed to add embedding to redis %s", ex)
+            raise ex
+
+        return embedding_results

api/core/rag/embedding/embedding_base.py

@@ -7,10 +7,12 @@ class Embeddings(ABC):
     @abstractmethod
     def embed_documents(self, texts: list[str]) -> list[list[float]]:
         """Embed search docs."""
+        raise NotImplementedError
 
     @abstractmethod
     def embed_query(self, text: str) -> list[float]:
         """Embed query text."""
+        raise NotImplementedError
 
     async def aembed_documents(self, texts: list[str]) -> list[list[float]]:
         """Asynchronous Embed search docs."""

api/core/rag/rerank/rerank_base.py

@@ -0,0 +1,26 @@
+from abc import ABC, abstractmethod
+from typing import Optional
+
+from core.rag.models.document import Document
+
+
+class BaseRerankRunner(ABC):
+    @abstractmethod
+    def run(
+        self,
+        query: str,
+        documents: list[Document],
+        score_threshold: Optional[float] = None,
+        top_n: Optional[int] = None,
+        user: Optional[str] = None,
+    ) -> list[Document]:
+        """
+        Run rerank model
+        :param query: search query
+        :param documents: documents for reranking
+        :param score_threshold: score threshold
+        :param top_n: top n
+        :param user: unique user id if needed
+        :return:
+        """
+        raise NotImplementedError

api/core/rag/rerank/rerank_factory.py

@@ -0,0 +1,16 @@
+from core.rag.rerank.rerank_base import BaseRerankRunner
+from core.rag.rerank.rerank_model import RerankModelRunner
+from core.rag.rerank.rerank_type import RerankMode
+from core.rag.rerank.weight_rerank import WeightRerankRunner
+
+
+class RerankRunnerFactory:
+    @staticmethod
+    def create_rerank_runner(runner_type: str, *args, **kwargs) -> BaseRerankRunner:
+        match runner_type:
+            case RerankMode.RERANKING_MODEL.value:
+                return RerankModelRunner(*args, **kwargs)
+            case RerankMode.WEIGHTED_SCORE.value:
+                return WeightRerankRunner(*args, **kwargs)
+            case _:
+                raise ValueError(f"Unknown runner type: {runner_type}")

api/core/rag/rerank/rerank_model.py

@@ -2,9 +2,10 @@ from typing import Optional
 
 from core.model_manager import ModelInstance
 from core.rag.models.document import Document
+from core.rag.rerank.rerank_base import BaseRerankRunner
 
 
-class RerankModelRunner:
+class RerankModelRunner(BaseRerankRunner):
     def __init__(self, rerank_model_instance: ModelInstance) -> None:
         self.rerank_model_instance = rerank_model_instance
 

api/core/rag/rerank/weight_rerank.py

@@ -4,15 +4,16 @@ from typing import Optional
 
 import numpy as np
 
-from core.embedding.cached_embedding import CacheEmbedding
 from core.model_manager import ModelManager
 from core.model_runtime.entities.model_entities import ModelType
 from core.rag.datasource.keyword.jieba.jieba_keyword_table_handler import JiebaKeywordTableHandler
+from core.rag.embedding.cached_embedding import CacheEmbedding
 from core.rag.models.document import Document
 from core.rag.rerank.entity.weight import VectorSetting, Weights
+from core.rag.rerank.rerank_base import BaseRerankRunner
 
 
-class WeightRerankRunner:
+class WeightRerankRunner(BaseRerankRunner):
     def __init__(self, tenant_id: str, weights: Weights) -> None:
         self.tenant_id = tenant_id
         self.weights = weights