codekingpro/portable-devtools
114k
1from __future__ import annotations2 3import os4import time5import uuid6from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple, Type7 8import numpy as np9from langchain_core.documents import Document10from langchain_core.embeddings import Embeddings11from langchain_core.vectorstores import VectorStore12 13if TYPE_CHECKING:14 import vearch15 16DEFAULT_TOPN = 417 18 19class Vearch(VectorStore):20 _DEFAULT_TABLE_NAME = "langchain_vearch"21 _DEFAULT_CLUSTER_DB_NAME = "cluster_client_db"22 _DEFAULT_VERSION = 123 24 def __init__(25 self,26 embedding_function: Embeddings,27 path_or_url: Optional[str] = None,28 table_name: str = _DEFAULT_TABLE_NAME,29 db_name: str = _DEFAULT_CLUSTER_DB_NAME,30 flag: int = _DEFAULT_VERSION,31 **kwargs: Any,32 ) -> None:33 """Initialize vearch vector store34 flag 1 for cluster,0 for standalone35 """36 try:37 if flag:38 import vearch_cluster39 else:40 import vearch41 except ImportError:42 raise ImportError(43 "Could not import suitable python package. "44 "Please install it with `pip install vearch or vearch_cluster`."45 )46 47 if flag:48 if path_or_url is None:49 raise ValueError("Please input url of cluster")50 if not db_name:51 db_name = self._DEFAULT_CLUSTER_DB_NAME52 db_name += "_"53 db_name += str(uuid.uuid4()).split("-")[-1]54 self.using_db_name = db_name55 self.url = path_or_url56 self.vearch = vearch_cluster.VearchCluster(path_or_url)57 58 else:59 if path_or_url is None:60 metadata_path = os.getcwd().replace("\\", "/")61 else:62 metadata_path = path_or_url63 if not os.path.isdir(metadata_path):64 os.makedirs(metadata_path)65 log_path = os.path.join(metadata_path, "log")66 if not os.path.isdir(log_path):67 os.makedirs(log_path)68 self.vearch = vearch.Engine(metadata_path, log_path)69 self.using_metapath = metadata_path70 if not table_name:71 table_name = self._DEFAULT_TABLE_NAME72 table_name += "_"73 table_name += str(uuid.uuid4()).split("-")[-1]74 self.using_table_name = table_name75 self.embedding_func = embedding_function76 self.flag = flag77 78 @property79 def embeddings(self) -> Optional[Embeddings]:80 return self.embedding_func81 82 @classmethod83 def from_documents(84 cls: Type[Vearch],85 documents: List[Document],86 embedding: Embeddings,87 path_or_url: Optional[str] = None,88 table_name: str = _DEFAULT_TABLE_NAME,89 db_name: str = _DEFAULT_CLUSTER_DB_NAME,90 flag: int = _DEFAULT_VERSION,91 **kwargs: Any,92 ) -> Vearch:93 """Return Vearch VectorStore"""94 95 texts = [d.page_content for d in documents]96 metadatas = [d.metadata for d in documents]97 98 return cls.from_texts(99 texts=texts,100 embedding=embedding,101 metadatas=metadatas,102 path_or_url=path_or_url,103 table_name=table_name,104 db_name=db_name,105 flag=flag,106 **kwargs,107 )108 109 @classmethod110 def from_texts(111 cls: Type[Vearch],112 texts: List[str],113 embedding: Embeddings,114 metadatas: Optional[List[dict]] = None,115 path_or_url: Optional[str] = None,116 table_name: str = _DEFAULT_TABLE_NAME,117 db_name: str = _DEFAULT_CLUSTER_DB_NAME,118 flag: int = _DEFAULT_VERSION,119 **kwargs: Any,120 ) -> Vearch:121 """Return Vearch VectorStore"""122 123 vearch_db = cls(124 embedding_function=embedding,125 embedding=embedding,126 path_or_url=path_or_url,127 db_name=db_name,128 table_name=table_name,129 flag=flag,130 )131 vearch_db.add_texts(texts=texts, metadatas=metadatas)132 return vearch_db133 134 def _create_table(135 self,136 dim: int = 1024,137 field_list: List[dict] = [138 {"field": "text", "type": "str"},139 {"field": "metadata", "type": "str"},140 ],141 ) -> int:142 """143 Create VectorStore Table144 Args:145 dim:dimension of vector146 fields_list: the field you want to store147 Return:148 code,0 for success,1 for failed149 """150 151 type_dict = {"int": vearch.dataType.INT, "str": vearch.dataType.STRING}152 engine_info = {153 "index_size": 10000,154 "retrieval_type": "IVFPQ",155 "retrieval_param": {"ncentroids": 2048, "nsubvector": 32},156 }157 fields = [158 vearch.GammaFieldInfo(fi["field"], type_dict[fi["type"]])159 for fi in field_list160 ]161 vector_field = vearch.GammaVectorInfo(162 name="text_embedding",163 type=vearch.dataType.VECTOR,164 is_index=True,165 dimension=dim,166 model_id="",167 store_type="MemoryOnly",168 store_param={"cache_size": 10000},169 has_source=False,170 )171 response_code = self.vearch.create_table(172 engine_info,173 name=self.using_table_name,174 fields=fields,175 vector_field=vector_field,176 )177 return response_code178 179 def _create_space(180 self,181 dim: int = 1024,182 ) -> int:183 """184 Create VectorStore space185 Args:186 dim:dimension of vector187 Return:188 code,0 failed for ,1 for success189 """190 space_config = {191 "name": self.using_table_name,192 "partition_num": 1,193 "replica_num": 1,194 "engine": {195 "name": "gamma",196 "index_size": 1,197 "retrieval_type": "FLAT",198 "retrieval_param": {199 "metric_type": "L2",200 },201 },202 "properties": {203 "text": {204 "type": "string",205 },206 "metadata": {207 "type": "string",208 },209 "text_embedding": {210 "type": "vector",211 "index": True,212 "dimension": dim,213 "store_type": "MemoryOnly",214 },215 },216 }217 response_code = self.vearch.create_space(self.using_db_name, space_config)218 219 return response_code220 221 def add_texts(222 self,223 texts: Iterable[str],224 metadatas: Optional[List[dict]] = None,225 **kwargs: Any,226 ) -> List[str]:227 """228 Returns:229 List of ids from adding the texts into the vectorstore.230 """231 embeddings = None232 if self.embedding_func is not None:233 embeddings = self.embedding_func.embed_documents(list(texts))234 if embeddings is None:235 raise ValueError("embeddings is None")236 if self.flag:237 dbs_list = self.vearch.list_dbs()238 if self.using_db_name not in dbs_list:239 create_db_code = self.vearch.create_db(self.using_db_name)240 if not create_db_code:241 raise ValueError("create db failed!!!")242 space_list = self.vearch.list_spaces(self.using_db_name)243 if self.using_table_name not in space_list:244 create_space_code = self._create_space(len(embeddings[0]))245 if not create_space_code:246 raise ValueError("create space failed!!!")247 docid = []248 if embeddings is not None and metadatas is not None:249 for text, metadata, embed in zip(texts, metadatas, embeddings):250 profiles: dict[str, Any] = {}251 profiles["text"] = text252 profiles["metadata"] = metadata["source"]253 embed_np = np.array(embed)254 profiles["text_embedding"] = {255 "feature": (embed_np / np.linalg.norm(embed_np)).tolist()256 }257 insert_res = self.vearch.insert_one(258 self.using_db_name, self.using_table_name, profiles259 )260 if insert_res["status"] == 200:261 docid.append(insert_res["_id"])262 continue263 else:264 retry_insert = self.vearch.insert_one(265 self.using_db_name, self.using_table_name, profiles266 )267 docid.append(retry_insert["_id"])268 continue269 else:270 table_path = os.path.join(271 self.using_metapath, self.using_table_name + ".schema"272 )273 if not os.path.exists(table_path):274 dim = len(embeddings[0])275 response_code = self._create_table(dim)276 if response_code:277 raise ValueError("create table failed!!!")278 if embeddings is not None and metadatas is not None:279 doc_items = []280 for text, metadata, embed in zip(texts, metadatas, embeddings):281 profiles_v: dict[str, Any] = {}282 profiles_v["text"] = text283 profiles_v["metadata"] = metadata["source"]284 embed_np = np.array(embed)285 profiles_v["text_embedding"] = embed_np / np.linalg.norm(embed_np)286 doc_items.append(profiles_v)287 288 docid = self.vearch.add(doc_items)289 t_time = 0290 while len(docid) != len(embeddings):291 time.sleep(0.5)292 if t_time > 6:293 break294 t_time += 1295 self.vearch.dump()296 return docid297 298 def _load(self) -> None:299 """300 load vearch engine for standalone vearch301 """302 self.vearch.load()303 304 @classmethod305 def load_local(306 cls,307 embedding: Embeddings,308 path_or_url: Optional[str] = None,309 table_name: str = _DEFAULT_TABLE_NAME,310 db_name: str = _DEFAULT_CLUSTER_DB_NAME,311 flag: int = _DEFAULT_VERSION,312 **kwargs: Any,313 ) -> Vearch:314 """Load the local specified table of standalone vearch.315 Returns:316 Success or failure of loading the local specified table317 """318 if not path_or_url:319 raise ValueError("No metadata path!!!")320 if not table_name:321 raise ValueError("No table name!!!")322 table_path = os.path.join(path_or_url, table_name + ".schema")323 if not os.path.exists(table_path):324 raise ValueError("vearch vectorbase table not exist!!!")325 326 vearch_db = cls(327 embedding_function=embedding,328 path_or_url=path_or_url,329 table_name=table_name,330 db_name=db_name,331 flag=flag,332 )333 vearch_db._load()334 return vearch_db335 336 def similarity_search(337 self,338 query: str,339 k: int = DEFAULT_TOPN,340 **kwargs: Any,341 ) -> List[Document]:342 """343 Return docs most similar to query.344 345 """346 if self.embedding_func is None:347 raise ValueError("embedding_func is None!!!")348 embeddings = self.embedding_func.embed_query(query)349 docs = self.similarity_search_by_vector(embeddings, k)350 return docs351 352 def similarity_search_by_vector(353 self,354 embedding: List[float],355 k: int = DEFAULT_TOPN,356 **kwargs: Any,357 ) -> List[Document]:358 """The most k similar documents and scores of the specified query.359 Args:360 embeddings: embedding vector of the query.361 k: The k most similar documents to the text query.362 min_score: the score of similar documents to the text query363 Returns:364 The k most similar documents to the specified text query.365 0 is dissimilar, 1 is the most similar.366 """367 embed = np.array(embedding)368 if self.flag:369 query_data = {370 "query": {371 "sum": [372 {373 "field": "text_embedding",374 "feature": (embed / np.linalg.norm(embed)).tolist(),375 }376 ],377 },378 "size": k,379 "fields": ["text", "metadata"],380 }381 query_result = self.vearch.search(382 self.using_db_name, self.using_table_name, query_data383 )384 res = query_result["hits"]["hits"]385 else:386 query_data = {387 "vector": [388 {389 "field": "text_embedding",390 "feature": embed / np.linalg.norm(embed),391 }392 ],393 "fields": [],394 "is_brute_search": 1,395 "retrieval_param": {"metric_type": "InnerProduct", "nprobe": 20},396 "topn": k,397 }398 query_result = self.vearch.search(query_data)399 res = query_result[0]["result_items"]400 docs = []401 for item in res:402 content = ""403 meta_data = {}404 if self.flag:405 item = item["_source"]406 for item_key in item:407 if item_key == "text":408 content = item[item_key]409 continue410 if item_key == "metadata":411 meta_data["source"] = item[item_key]412 continue413 docs.append(Document(page_content=content, metadata=meta_data))414 return docs415 416 def similarity_search_with_score(417 self,418 query: str,419 k: int = DEFAULT_TOPN,420 **kwargs: Any,421 ) -> List[Tuple[Document, float]]:422 """The most k similar documents and scores of the specified query.423 Args:424 embeddings: embedding vector of the query.425 k: The k most similar documents to the text query.426 min_score: the score of similar documents to the text query427 Returns:428 The k most similar documents to the specified text query.429 0 is dissimilar, 1 is the most similar.430 """431 if self.embedding_func is None:432 raise ValueError("embedding_func is None!!!")433 embeddings = self.embedding_func.embed_query(query)434 embed = np.array(embeddings)435 if self.flag:436 query_data = {437 "query": {438 "sum": [439 {440 "field": "text_embedding",441 "feature": (embed / np.linalg.norm(embed)).tolist(),442 }443 ],444 },445 "size": k,446 "fields": ["text_embedding", "text", "metadata"],447 }448 query_result = self.vearch.search(449 self.using_db_name, self.using_table_name, query_data450 )451 res = query_result["hits"]["hits"]452 else:453 query_data = {454 "vector": [455 {456 "field": "text_embedding",457 "feature": embed / np.linalg.norm(embed),458 }459 ],460 "fields": [],461 "is_brute_search": 1,462 "retrieval_param": {"metric_type": "InnerProduct", "nprobe": 20},463 "topn": k,464 }465 query_result = self.vearch.search(query_data)466 res = query_result[0]["result_items"]467 results: List[Tuple[Document, float]] = []468 for item in res:469 content = ""470 meta_data = {}471 if self.flag:472 score = item["_score"]473 item = item["_source"]474 for item_key in item:475 if item_key == "text":476 content = item[item_key]477 continue478 if item_key == "metadata":479 meta_data["source"] = item[item_key]480 continue481 if self.flag != 1 and item_key == "score":482 score = item[item_key]483 continue484 tmp_res = (Document(page_content=content, metadata=meta_data), score)485 results.append(tmp_res)486 return results487 488 def _similarity_search_with_relevance_scores(489 self,490 query: str,491 k: int = 4,492 **kwargs: Any,493 ) -> List[Tuple[Document, float]]:494 return self.similarity_search_with_score(query, k, **kwargs)495 496 def delete(497 self,498 ids: Optional[List[str]] = None,499 **kwargs: Any,500 ) -> Optional[bool]:501 """Delete the documents which have the specified ids.502 503 Args:504 ids: The ids of the embedding vectors.505 **kwargs: Other keyword arguments that subclasses might use.506 Returns:507 Optional[bool]: True if deletion is successful.508 False otherwise, None if not implemented.509 """510 511 ret: Optional[bool] = None512 tmp_res = []513 if ids is None or ids.__len__() == 0:514 return ret515 for _id in ids:516 if self.flag:517 ret = self.vearch.delete(self.using_db_name, self.using_table_name, _id)518 else:519 ret = self.vearch.del_doc(_id)520 tmp_res.append(ret)521 ret = all(i == 0 for i in tmp_res)522 return ret523 524 def get(525 self,526 ids: Optional[List[str]] = None,527 **kwargs: Any,528 ) -> Dict[str, Document]:529 """Return docs according ids.530 531 Args:532 ids: The ids of the embedding vectors.533 Returns:534 Documents which satisfy the input conditions.535 """536 537 results: Dict[str, Document] = {}538 if ids is None or ids.__len__() == 0:539 return results540 if self.flag:541 query_data = {"query": {"ids": ids}}542 docs_detail = self.vearch.mget_by_ids(543 self.using_db_name, self.using_table_name, query_data544 )545 for record in docs_detail:546 if record["found"] is False:547 continue548 content = ""549 meta_info = {}550 for field in record["_source"]:551 if field == "text":552 content = record["_source"][field]553 continue554 elif field == "metadata":555 meta_info["source"] = record["_source"][field]556 continue557 results[record["_id"]] = Document(558 page_content=content, metadata=meta_info559 )560 else:561 for id in ids:562 docs_detail = self.vearch.get_doc_by_id(id)563 if docs_detail == {}:564 continue565 content = ""566 meta_info = {}567 for field in docs_detail:568 if field == "text":569 content = docs_detail[field]570 continue571 elif field == "metadata":572 meta_info["source"] = docs_detail[field]573 continue574 results[docs_detail["_id"]] = Document(575 page_content=content, metadata=meta_info576 )577 return results578