diff --git a/api/utils/log_utils.py b/api/utils/log_utils.py index 7ec1cafef..7d07f7909 100644 --- a/api/utils/log_utils.py +++ b/api/utils/log_utils.py @@ -77,4 +77,11 @@ def initRootLogger(logfile_basename: str, log_format: str = "%(asctime)-15s %(le pkg_logger.setLevel(pkg_level) msg = f"{logfile_basename} log path: {log_path}, log levels: {pkg_levels}" - logger.info(msg) \ No newline at end of file + logger.info(msg) + + +def log_exception(e, *args): + logging.exception(e) + for a in args: + logging.error(str(a)) + raise e \ No newline at end of file diff --git a/rag/llm/embedding_model.py b/rag/llm/embedding_model.py index e02c81c27..4b1a48abe 100644 --- a/rag/llm/embedding_model.py +++ b/rag/llm/embedding_model.py @@ -19,7 +19,6 @@ import threading from urllib.parse import urljoin import requests -from requests.exceptions import JSONDecodeError from huggingface_hub import snapshot_download from zhipuai import ZhipuAI import os @@ -32,6 +31,7 @@ import asyncio from api import settings from api.utils.file_utils import get_home_cache_dir +from api.utils.log_utils import log_exception from rag.utils import num_tokens_from_string, truncate import google.generativeai as genai import json @@ -130,8 +130,11 @@ class OpenAIEmbed(Base): for i in range(0, len(texts), batch_size): res = self.client.embeddings.create(input=texts[i:i + batch_size], model=self.model_name) - ress.extend([d.embedding for d in res.data]) - total_tokens += self.total_token_count(res) + try: + ress.extend([d.embedding for d in res.data]) + total_tokens += self.total_token_count(res) + except Exception as _e: + log_exception(_e, res) return np.array(ress), total_tokens def encode_queries(self, text): @@ -153,7 +156,10 @@ class LocalAIEmbed(Base): ress = [] for i in range(0, len(texts), batch_size): res = self.client.embeddings.create(input=texts[i:i + batch_size], model=self.model_name) - ress.extend([d.embedding for d in res.data]) + try: + ress.extend([d.embedding for d in res.data]) + except Exception as _e: + log_exception(_e, res) # local embedding for LmStudio donot count tokens return np.array(ress), 1024 @@ -188,40 +194,39 @@ class QWenEmbed(Base): def encode(self, texts: list): import dashscope batch_size = 4 - try: - res = [] - token_count = 0 - texts = [truncate(t, 2048) for t in texts] - for i in range(0, len(texts), batch_size): - resp = dashscope.TextEmbedding.call( - model=self.model_name, - input=texts[i:i + batch_size], - api_key=self.key, - text_type="document" - ) + res = [] + token_count = 0 + texts = [truncate(t, 2048) for t in texts] + for i in range(0, len(texts), batch_size): + resp = dashscope.TextEmbedding.call( + model=self.model_name, + input=texts[i:i + batch_size], + api_key=self.key, + text_type="document" + ) + try: embds = [[] for _ in range(len(resp["output"]["embeddings"]))] for e in resp["output"]["embeddings"]: embds[e["text_index"]] = e["embedding"] res.extend(embds) token_count += self.total_token_count(resp) - return np.array(res), token_count - except Exception as e: - raise Exception("Account abnormal. Please ensure it's on good standing to use QWen's "+self.model_name) - return np.array([]), 0 + except Exception as _e: + log_exception(_e, resp) + raise + return np.array(res), token_count def encode_queries(self, text): + resp = dashscope.TextEmbedding.call( + model=self.model_name, + input=text[:2048], + api_key=self.key, + text_type="query" + ) try: - resp = dashscope.TextEmbedding.call( - model=self.model_name, - input=text[:2048], - api_key=self.key, - text_type="query" - ) return np.array(resp["output"]["embeddings"][0] ["embedding"]), self.total_token_count(resp) - except Exception: - raise Exception("Account abnormal. Please ensure it's on good standing to use QWen's "+self.model_name) - return np.array([]), 0 + except Exception as _e: + log_exception(_e, resp) class ZhipuEmbed(Base): @@ -243,14 +248,20 @@ class ZhipuEmbed(Base): for txt in texts: res = self.client.embeddings.create(input=txt, model=self.model_name) - arr.append(res.data[0].embedding) - tks_num += self.total_token_count(res) + try: + arr.append(res.data[0].embedding) + tks_num += self.total_token_count(res) + except Exception as _e: + log_exception(_e, res) return np.array(arr), tks_num def encode_queries(self, text): res = self.client.embeddings.create(input=text, model=self.model_name) - return np.array(res.data[0].embedding), self.total_token_count(res) + try: + return np.array(res.data[0].embedding), self.total_token_count(res) + except Exception as _e: + log_exception(_e, res) class OllamaEmbed(Base): @@ -266,7 +277,10 @@ class OllamaEmbed(Base): res = self.client.embeddings(prompt=txt, model=self.model_name, options={"use_mmap": True}) - arr.append(res["embedding"]) + try: + arr.append(res["embedding"]) + except Exception as _e: + log_exception(_e, res) tks_num += 128 return np.array(arr), tks_num @@ -274,7 +288,10 @@ class OllamaEmbed(Base): res = self.client.embeddings(prompt=text, model=self.model_name, options={"use_mmap": True}) - return np.array(res["embedding"]), 128 + try: + return np.array(res["embedding"]), 128 + except Exception as _e: + log_exception(_e, res) class FastEmbed(DefaultEmbedding): @@ -334,14 +351,20 @@ class XinferenceEmbed(Base): total_tokens = 0 for i in range(0, len(texts), batch_size): res = self.client.embeddings.create(input=texts[i:i + batch_size], model=self.model_name) - ress.extend([d.embedding for d in res.data]) - total_tokens += self.total_token_count(res) + try: + ress.extend([d.embedding for d in res.data]) + total_tokens += self.total_token_count(res) + except Exception as _e: + log_exception(_e, res) return np.array(ress), total_tokens def encode_queries(self, text): res = self.client.embeddings.create(input=[text], model=self.model_name) - return np.array(res.data[0].embedding), self.total_token_count(res) + try: + return np.array(res.data[0].embedding), self.total_token_count(res) + except Exception as _e: + log_exception(_e, res) class YoudaoEmbed(Base): @@ -401,11 +424,10 @@ class JinaEmbed(Base): response = requests.post(self.base_url, headers=self.headers, json=data) try: res = response.json() - except JSONDecodeError as e: - logging.error(f"JSON decode error: {e}\nResponse content: {response.text[:2000]}") - raise - ress.extend([d["embedding"] for d in res["data"]]) - token_count += self.total_token_count(res) + ress.extend([d["embedding"] for d in res["data"]]) + token_count += self.total_token_count(res) + except Exception as _e: + log_exception(_e, response) return np.array(ress), token_count def encode_queries(self, text): @@ -468,14 +490,20 @@ class MistralEmbed(Base): for i in range(0, len(texts), batch_size): res = self.client.embeddings(input=texts[i:i + batch_size], model=self.model_name) - ress.extend([d.embedding for d in res.data]) - token_count += self.total_token_count(res) + try: + ress.extend([d.embedding for d in res.data]) + token_count += self.total_token_count(res) + except Exception as _e: + log_exception(_e, res) return np.array(ress), token_count def encode_queries(self, text): res = self.client.embeddings(input=[truncate(text, 8196)], model=self.model_name) - return np.array(res.data[0].embedding), self.total_token_count(res) + try: + return np.array(res.data[0].embedding), self.total_token_count(res) + except Exception as _e: + log_exception(_e, res) class BedrockEmbed(Base): @@ -505,9 +533,12 @@ class BedrockEmbed(Base): body = {"texts": [text], "input_type": 'search_document'} response = self.client.invoke_model(modelId=self.model_name, body=json.dumps(body)) - model_response = json.loads(response["body"].read()) - embeddings.extend([model_response["embedding"]]) - token_count += num_tokens_from_string(text) + try: + model_response = json.loads(response["body"].read()) + embeddings.extend([model_response["embedding"]]) + token_count += num_tokens_from_string(text) + except Exception as _e: + log_exception(_e, response) return np.array(embeddings), token_count @@ -520,8 +551,11 @@ class BedrockEmbed(Base): body = {"texts": [truncate(text, 8196)], "input_type": 'search_query'} response = self.client.invoke_model(modelId=self.model_name, body=json.dumps(body)) - model_response = json.loads(response["body"].read()) - embeddings.extend(model_response["embedding"]) + try: + model_response = json.loads(response["body"].read()) + embeddings.extend(model_response["embedding"]) + except Exception as _e: + log_exception(_e, response) return np.array(embeddings), token_count @@ -544,7 +578,10 @@ class GeminiEmbed(Base): content=texts[i: i + batch_size], task_type="retrieval_document", title="Embedding of single string") - ress.extend(result['embedding']) + try: + ress.extend(result['embedding']) + except Exception as _e: + log_exception(_e, result) return np.array(ress),token_count def encode_queries(self, text): @@ -555,7 +592,10 @@ class GeminiEmbed(Base): task_type="retrieval_document", title="Embedding of single string") token_count = num_tokens_from_string(text) - return np.array(result['embedding']), token_count + try: + return np.array(result['embedding']), token_count + except Exception as _e: + log_exception(_e, result) class NvidiaEmbed(Base): @@ -593,9 +633,8 @@ class NvidiaEmbed(Base): response = requests.post(self.base_url, headers=self.headers, json=payload) try: res = response.json() - except JSONDecodeError as e: - logging.error(f"JSON decode error: {e}\nResponse content: {response.text[:2000]}") - raise + except Exception as _e: + log_exception(_e, response) ress.extend([d["embedding"] for d in res["data"]]) token_count += self.total_token_count(res) return np.array(ress), token_count @@ -641,8 +680,11 @@ class CoHereEmbed(Base): input_type="search_document", embedding_types=["float"], ) - ress.extend([d for d in res.embeddings.float]) - token_count += res.meta.billed_units.input_tokens + try: + ress.extend([d for d in res.embeddings.float]) + token_count += res.meta.billed_units.input_tokens + except Exception as _e: + log_exception(_e, res) return np.array(ress), token_count def encode_queries(self, text): @@ -652,9 +694,10 @@ class CoHereEmbed(Base): input_type="search_query", embedding_types=["float"], ) - return np.array(res.embeddings.float[0]), int( - res.meta.billed_units.input_tokens - ) + try: + return np.array(res.embeddings.float[0]), int(res.meta.billed_units.input_tokens) + except Exception as _e: + log_exception(_e, res) class TogetherAIEmbed(OpenAIEmbed): @@ -706,13 +749,11 @@ class SILICONFLOWEmbed(Base): response = requests.post(self.base_url, json=payload, headers=self.headers) try: res = response.json() - except JSONDecodeError as e: - logging.error(f"JSON decode error: {e}\nResponse content: {response.text[:2000]}") - raise - if "data" not in res or not isinstance(res["data"], list) or len(res["data"]) != len(texts_batch): - raise ValueError(f"SILICONFLOWEmbed.encode got invalid response from {self.base_url}") - ress.extend([d["embedding"] for d in res["data"]]) - token_count += self.total_token_count(res) + ress.extend([d["embedding"] for d in res["data"]]) + token_count += self.total_token_count(res) + except Exception as _e: + log_exception(_e, response) + return np.array(ress), token_count def encode_queries(self, text): @@ -724,12 +765,9 @@ class SILICONFLOWEmbed(Base): response = requests.post(self.base_url, json=payload, headers=self.headers).json() try: res = response.json() - except JSONDecodeError as e: - logging.error(f"JSON decode error: {e}\nResponse content: {response.text[:2000]}") - raise - if "data" not in res or not isinstance(res["data"], list) or len(res["data"])!= 1: - raise ValueError(f"SILICONFLOWEmbed.encode_queries got invalid response from {self.base_url}") - return np.array(res["data"][0]["embedding"]), self.total_token_count(res) + return np.array(res["data"][0]["embedding"]), self.total_token_count(res) + except Exception as _e: + log_exception(_e, response) class ReplicateEmbed(Base): @@ -765,17 +803,23 @@ class BaiduYiyanEmbed(Base): def encode(self, texts: list, batch_size=16): res = self.client.do(model=self.model_name, texts=texts).body - return ( - np.array([r["embedding"] for r in res["data"]]), - self.total_token_count(res), - ) + try: + return ( + np.array([r["embedding"] for r in res["data"]]), + self.total_token_count(res), + ) + except Exception as _e: + log_exception(_e, res) def encode_queries(self, text): res = self.client.do(model=self.model_name, texts=[text]).body - return ( - np.array([r["embedding"] for r in res["data"]]), - self.total_token_count(res), - ) + try: + return ( + np.array([r["embedding"] for r in res["data"]]), + self.total_token_count(res), + ) + except Exception as _e: + log_exception(_e, res) class VoyageEmbed(Base): @@ -793,15 +837,21 @@ class VoyageEmbed(Base): res = self.client.embed( texts=texts[i : i + batch_size], model=self.model_name, input_type="document" ) - ress.extend(res.embeddings) - token_count += res.total_tokens + try: + ress.extend(res.embeddings) + token_count += res.total_tokens + except Exception as _e: + log_exception(_e, res) return np.array(ress), token_count def encode_queries(self, text): res = self.client.embed( texts=text, model=self.model_name, input_type="query" ) - return np.array(res.embeddings)[0], res.total_tokens + try: + return np.array(res.embeddings)[0], res.total_tokens + except Exception as _e: + log_exception(_e, res) class HuggingFaceEmbed(Base): @@ -821,11 +871,14 @@ class HuggingFaceEmbed(Base): headers={'Content-Type': 'application/json'} ) if response.status_code == 200: - embedding = response.json() - embeddings.append(embedding[0]) + try: + embedding = response.json() + embeddings.append(embedding[0]) + return np.array(embeddings), sum([num_tokens_from_string(text) for text in texts]) + except Exception as _e: + log_exception(_e, response) else: raise Exception(f"Error: {response.status_code} - {response.text}") - return np.array(embeddings), sum([num_tokens_from_string(text) for text in texts]) def encode_queries(self, text): response = requests.post( @@ -834,8 +887,11 @@ class HuggingFaceEmbed(Base): headers={'Content-Type': 'application/json'} ) if response.status_code == 200: - embedding = response.json() - return np.array(embedding[0]), num_tokens_from_string(text) + try: + embedding = response.json() + return np.array(embedding[0]), num_tokens_from_string(text) + except Exception as _e: + log_exception(_e, response) else: raise Exception(f"Error: {response.status_code} - {response.text}") @@ -848,6 +904,7 @@ class VolcEngineEmbed(OpenAIEmbed): model_name = json.loads(key).get('ep_id', '') + json.loads(key).get('endpoint_id', '') super().__init__(ark_api_key,model_name,base_url) + class GPUStackEmbed(OpenAIEmbed): def __init__(self, key, model_name, base_url): if not base_url: diff --git a/rag/llm/rerank_model.py b/rag/llm/rerank_model.py index 49f46349a..45ef4bb0f 100644 --- a/rag/llm/rerank_model.py +++ b/rag/llm/rerank_model.py @@ -28,6 +28,7 @@ from yarl import URL from api import settings from api.utils.file_utils import get_home_cache_dir +from api.utils.log_utils import log_exception from rag.utils import num_tokens_from_string, truncate import json @@ -170,8 +171,11 @@ class JinaRerank(Base): } res = requests.post(self.base_url, headers=self.headers, json=data).json() rank = np.zeros(len(texts), dtype=float) - for d in res["results"]: - rank[d["index"]] = d["relevance_score"] + try: + for d in res["results"]: + rank[d["index"]] = d["relevance_score"] + except Exception as _e: + log_exception(_e, res) return rank, self.total_token_count(res) @@ -238,8 +242,11 @@ class XInferenceRerank(Base): } res = requests.post(self.base_url, headers=self.headers, json=data).json() rank = np.zeros(len(texts), dtype=float) - for d in res["results"]: - rank[d["index"]] = d["relevance_score"] + try: + for d in res["results"]: + rank[d["index"]] = d["relevance_score"] + except Exception as _e: + log_exception(_e, res) return rank, token_count @@ -269,10 +276,11 @@ class LocalAIRerank(Base): token_count += num_tokens_from_string(t) res = requests.post(self.base_url, headers=self.headers, json=data).json() rank = np.zeros(len(texts), dtype=float) - if 'results' not in res: - raise ValueError("response not contains results\n" + str(res)) - for d in res["results"]: - rank[d["index"]] = d["relevance_score"] + try: + for d in res["results"]: + rank[d["index"]] = d["relevance_score"] + except Exception as _e: + log_exception(_e, res) # Normalize the rank values to the range 0 to 1 min_rank = np.min(rank) @@ -322,8 +330,11 @@ class NvidiaRerank(Base): } res = requests.post(self.base_url, headers=self.headers, json=data).json() rank = np.zeros(len(texts), dtype=float) - for d in res["rankings"]: - rank[d["index"]] = d["logit"] + try: + for d in res["rankings"]: + rank[d["index"]] = d["logit"] + except Exception as _e: + log_exception(_e, res) return rank, token_count @@ -361,10 +372,11 @@ class OpenAI_APIRerank(Base): token_count += num_tokens_from_string(t) res = requests.post(self.base_url, headers=self.headers, json=data).json() rank = np.zeros(len(texts), dtype=float) - if 'results' not in res: - raise ValueError("response not contains results\n" + str(res)) - for d in res["results"]: - rank[d["index"]] = d["relevance_score"] + try: + for d in res["results"]: + rank[d["index"]] = d["relevance_score"] + except Exception as _e: + log_exception(_e, res) # Normalize the rank values to the range 0 to 1 min_rank = np.min(rank) @@ -398,8 +410,11 @@ class CoHereRerank(Base): return_documents=False, ) rank = np.zeros(len(texts), dtype=float) - for d in res.results: - rank[d.index] = d.relevance_score + try: + for d in res.results: + rank[d.index] = d.relevance_score + except Exception as _e: + log_exception(_e, res) return rank, token_count @@ -439,11 +454,11 @@ class SILICONFLOWRerank(Base): self.base_url, json=payload, headers=self.headers ).json() rank = np.zeros(len(texts), dtype=float) - if "results" not in response: - return rank, 0 - - for d in response["results"]: - rank[d["index"]] = d["relevance_score"] + try: + for d in response["results"]: + rank[d["index"]] = d["relevance_score"] + except Exception as _e: + log_exception(_e, response) return ( rank, response["meta"]["tokens"]["input_tokens"] + response["meta"]["tokens"]["output_tokens"], @@ -468,8 +483,11 @@ class BaiduYiyanRerank(Base): top_n=len(texts), ).body rank = np.zeros(len(texts), dtype=float) - for d in res["results"]: - rank[d["index"]] = d["relevance_score"] + try: + for d in res["results"]: + rank[d["index"]] = d["relevance_score"] + except Exception as _e: + log_exception(_e, res) return rank, self.total_token_count(res) @@ -487,8 +505,11 @@ class VoyageRerank(Base): res = self.client.rerank( query=query, documents=texts, model=self.model_name, top_k=len(texts) ) - for r in res.results: - rank[r.index] = r.relevance_score + try: + for r in res.results: + rank[r.index] = r.relevance_score + except Exception as _e: + log_exception(_e, res) return rank, res.total_tokens @@ -511,8 +532,11 @@ class QWenRerank(Base): ) rank = np.zeros(len(texts), dtype=float) if resp.status_code == HTTPStatus.OK: - for r in resp.output.results: - rank[r.index] = r.relevance_score + try: + for r in resp.output.results: + rank[r.index] = r.relevance_score + except Exception as _e: + log_exception(_e, resp) return rank, resp.usage.total_tokens else: raise ValueError(f"Error calling QWenRerank model {self.model_name}: {resp.status_code} - {resp.text}") @@ -529,6 +553,7 @@ class HuggingfaceRerank(DefaultRerank): res = requests.post(f"http://{url}/rerank", headers={"Content-Type": "application/json"}, json={"query": query, "texts": texts[i: i + batch_size], "raw_scores": False, "truncate": True}) + for o in res.json(): scores[o["index"] + i] = o["score"] except Exception as e: @@ -582,15 +607,15 @@ class GPUStackRerank(Base): response_json = response.json() rank = np.zeros(len(texts), dtype=float) - if "results" not in response_json: - return rank, 0 token_count = 0 for t in texts: token_count += num_tokens_from_string(t) - - for result in response_json["results"]: - rank[result["index"]] = result["relevance_score"] + try: + for result in response_json["results"]: + rank[result["index"]] = result["relevance_score"] + except Exception as _e: + log_exception(_e, response) return ( rank,