fix embedding doc err

This commit is contained in:
huangzhuo 2025-03-05 19:58:23 +08:00
parent 2c17bb2c36
commit d88624755a

View File

@ -1,5 +1,7 @@
import base64 import base64
import logging import logging
import random
import time
from typing import Any, Optional, cast from typing import Any, Optional, cast
import numpy as np import numpy as np
@ -19,11 +21,49 @@ from models.dataset import Embedding
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
def retry(max_retries: int, base_delay: int, max_wait_time: int = 16):
"""
A retry decorator that uses an exponential backoff algorithm.
:param max_retries: The maximum number of retries.
:param base_delay: The base delay time in seconds.
:param max_wait_time: The maximum wait time in seconds.
:return: The decorated function.
"""
def decorator(func):
def wrapper(*args, **kwargs):
retries = 0
while retries <= max_retries:
try:
return func(*args, **kwargs)
except Exception as e:
if retries == max_retries:
raise e
logger.warning(f"Attempt {retries + 1} failed: {e}")
retries += 1
delay = base_delay * 2**retries
if delay > max_wait_time:
delay = max_wait_time
delay += random.uniform(0, 1)
logger.info(f"Retrying in {delay:.2f} seconds...")
time.sleep(delay)
return wrapper
return decorator
class CacheEmbedding(Embeddings): class CacheEmbedding(Embeddings):
def __init__(self, model_instance: ModelInstance, user: Optional[str] = None) -> None: def __init__(self, model_instance: ModelInstance, user: Optional[str] = None) -> None:
self._model_instance = model_instance self._model_instance = model_instance
self._user = user self._user = user
@retry(max_retries=10, base_delay=1, max_wait_time=16)
def text_embedding(self, texts: list[str]):
return self._model_instance.invoke_text_embedding(
texts=texts, user=self._user, input_type=EmbeddingInputType.DOCUMENT
)
def embed_documents(self, texts: list[str]) -> list[list[float]]: def embed_documents(self, texts: list[str]) -> list[list[float]]:
"""Embed search docs in batches of 10.""" """Embed search docs in batches of 10."""
# use doc embedding cache or store if not exists # use doc embedding cache or store if not exists
@ -58,9 +98,7 @@ class CacheEmbedding(Embeddings):
for i in range(0, len(embedding_queue_texts), max_chunks): for i in range(0, len(embedding_queue_texts), max_chunks):
batch_texts = embedding_queue_texts[i : i + max_chunks] batch_texts = embedding_queue_texts[i : i + max_chunks]
embedding_result = self._model_instance.invoke_text_embedding( embedding_result = self.text_embedding(batch_texts)
texts=batch_texts, user=self._user, input_type=EmbeddingInputType.DOCUMENT
)
for vector in embedding_result.embeddings: for vector in embedding_result.embeddings:
try: try: