Underground-Digital/Workflow-Engine

1import base642import logging3from typing import Optional, cast4 5import numpy as np6from sqlalchemy.exc import IntegrityError7 8from configs import dify_config9from core.entities.embedding_type import EmbeddingInputType10from core.model_manager import ModelInstance11from core.model_runtime.entities.model_entities import ModelPropertyKey12from core.model_runtime.model_providers.__base.text_embedding_model import TextEmbeddingModel13from core.rag.embedding.embedding_base import Embeddings14from extensions.ext_database import db15from extensions.ext_redis import redis_client16from libs import helper17from models.dataset import Embedding18 19logger = logging.getLogger(__name__)20 21 22class CacheEmbedding(Embeddings):23    def __init__(self, model_instance: ModelInstance, user: Optional[str] = None) -> None:24        self._model_instance = model_instance25        self._user = user26 27    def embed_documents(self, texts: list[str]) -> list[list[float]]:28        """Embed search docs in batches of 10."""29        # use doc embedding cache or store if not exists30        text_embeddings = [None for _ in range(len(texts))]31        embedding_queue_indices = []32        for i, text in enumerate(texts):33            hash = helper.generate_text_hash(text)34            embedding = (35                db.session.query(Embedding)36                .filter_by(37                    model_name=self._model_instance.model, hash=hash, provider_name=self._model_instance.provider38                )39                .first()40            )41            if embedding:42                text_embeddings[i] = embedding.get_embedding()43            else:44                embedding_queue_indices.append(i)45        if embedding_queue_indices:46            embedding_queue_texts = [texts[i] for i in embedding_queue_indices]47            embedding_queue_embeddings = []48            try:49                model_type_instance = cast(TextEmbeddingModel, self._model_instance.model_type_instance)50                model_schema = model_type_instance.get_model_schema(51                    self._model_instance.model, self._model_instance.credentials52                )53                max_chunks = (54                    model_schema.model_properties[ModelPropertyKey.MAX_CHUNKS]55                    if model_schema and ModelPropertyKey.MAX_CHUNKS in model_schema.model_properties56                    else 157                )58                for i in range(0, len(embedding_queue_texts), max_chunks):59                    batch_texts = embedding_queue_texts[i : i + max_chunks]60 61                    embedding_result = self._model_instance.invoke_text_embedding(62                        texts=batch_texts, user=self._user, input_type=EmbeddingInputType.DOCUMENT63                    )64 65                    for vector in embedding_result.embeddings:66                        try:67                            normalized_embedding = (vector / np.linalg.norm(vector)).tolist()68                            embedding_queue_embeddings.append(normalized_embedding)69                        except IntegrityError:70                            db.session.rollback()71                        except Exception as e:72                            logging.exception("Failed transform embedding: %s", e)73                cache_embeddings = []74                try:75                    for i, embedding in zip(embedding_queue_indices, embedding_queue_embeddings):76                        text_embeddings[i] = embedding77                        hash = helper.generate_text_hash(texts[i])78                        if hash not in cache_embeddings:79                            embedding_cache = Embedding(80                                model_name=self._model_instance.model,81                                hash=hash,82                                provider_name=self._model_instance.provider,83                            )84                            embedding_cache.set_embedding(embedding)85                            db.session.add(embedding_cache)86                            cache_embeddings.append(hash)87                    db.session.commit()88                except IntegrityError:89                    db.session.rollback()90            except Exception as ex:91                db.session.rollback()92                logger.error("Failed to embed documents: %s", ex)93                raise ex94 95        return text_embeddings96 97    def embed_query(self, text: str) -> list[float]:98        """Embed query text."""99        # use doc embedding cache or store if not exists100        hash = helper.generate_text_hash(text)101        embedding_cache_key = f"{self._model_instance.provider}_{self._model_instance.model}_{hash}"102        embedding = redis_client.get(embedding_cache_key)103        if embedding:104            redis_client.expire(embedding_cache_key, 600)105            return list(np.frombuffer(base64.b64decode(embedding), dtype="float"))106        try:107            embedding_result = self._model_instance.invoke_text_embedding(108                texts=[text], user=self._user, input_type=EmbeddingInputType.QUERY109            )110 111            embedding_results = embedding_result.embeddings[0]112            embedding_results = (embedding_results / np.linalg.norm(embedding_results)).tolist()113        except Exception as ex:114            if dify_config.DEBUG:115                logging.exception(f"Failed to embed query text: {ex}")116            raise ex117 118        try:119            # encode embedding to base64120            embedding_vector = np.array(embedding_results)121            vector_bytes = embedding_vector.tobytes()122            # Transform to Base64123            encoded_vector = base64.b64encode(vector_bytes)124            # Transform to string125            encoded_str = encoded_vector.decode("utf-8")126            redis_client.setex(embedding_cache_key, 600, encoded_str)127        except Exception as ex:128            if dify_config.DEBUG:129                logging.exception("Failed to add embedding to redis %s", ex)130            raise ex131 132        return embedding_results133