from __future__ import annotations
# Standard library imports
from typing import Any, Dict, Iterable, List, Optional, Tuple
from uuid import uuid4
# LangChain imports
from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings
from langchain_core.vectorstores import VectorStore
[docs]class VLite(VectorStore):
"""VLite是一个简单快速的矢量数据库，用于语义搜索。"""
[docs] def __init__(
self,
embedding_function: Embeddings,
collection: Optional[str] = None,
**kwargs: Any,
):
super().__init__()
self.embedding_function = embedding_function
self.collection = collection or f"vlite_{uuid4().hex}"
# Third-party imports
try:
from vlite import VLite
except ImportError:
raise ImportError(
"Could not import vlite python package. "
"Please install it with `pip install vlite`."
)
self.vlite = VLite(collection=self.collection, **kwargs)
[docs] def add_texts(
self,
texts: Iterable[str],
metadatas: Optional[List[dict]] = None,
**kwargs: Any,
) -> List[str]:
"""运行更多文本通过嵌入并添加到向量存储。
参数：
texts：要添加到向量存储的字符串的可迭代对象。
metadatas：与文本相关联的元数据的可选列表。
kwargs：向量存储特定参数
返回：
将文本添加到向量存储中的ID列表。
"""
texts = list(texts)
ids = kwargs.pop("ids", [str(uuid4()) for _ in texts])
embeddings = self.embedding_function.embed_documents(texts)
if not metadatas:
metadatas = [{} for _ in texts]
data_points = [
{"text": text, "metadata": metadata, "id": id, "embedding": embedding}
for text, metadata, id, embedding in zip(texts, metadatas, ids, embeddings)
]
results = self.vlite.add(data_points)
return [result[0] for result in results]
[docs] def add_documents(
self,
documents: List[Document],
**kwargs: Any,
) -> List[str]:
"""将一组文档添加到向量存储中。
参数:
documents: 要添加到向量存储中的文档列表。
kwargs: 向量存储特定参数，例如用于直接与vlite处理的"file_path"。
返回:
将文档添加到向量存储中后的id列表。
"""
ids = kwargs.pop("ids", [str(uuid4()) for _ in documents])
texts = []
metadatas = []
for doc, id in zip(documents, ids):
if "file_path" in kwargs:
# Third-party imports
try:
from vlite.utils import process_file
except ImportError:
raise ImportError(
"Could not import vlite python package. "
"Please install it with `pip install vlite`."
)
processed_data = process_file(kwargs["file_path"])
texts.extend(processed_data)
metadatas.extend([doc.metadata] * len(processed_data))
ids.extend([f"{id}_{i}" for i in range(len(processed_data))])
else:
texts.append(doc.page_content)
metadatas.append(doc.metadata)
return self.add_texts(texts, metadatas, ids=ids)
[docs] def similarity_search(
self,
query: str,
k: int = 4,
**kwargs: Any,
) -> List[Document]:
"""返回与查询最相似的文档。
参数：
query：要查找与之相似的文档的文本。
k：要返回的文档数量。默认为4。
返回：
与查询最相似的文档列表。
"""
docs_and_scores = self.similarity_search_with_score(query, k=k)
return [doc for doc, _ in docs_and_scores]
[docs] def similarity_search_with_score(
self,
query: str,
k: int = 4,
filter: Optional[Dict[str, str]] = None,
**kwargs: Any,
) -> List[Tuple[Document, float]]:
"""返回与查询最相似的文档。
参数：
query：要查找相似文档的文本。
k：要返回的文档数量。默认为4。
filter：按元数据过滤。默认为None。
返回：
元组列表（doc，score），其中score是相似度分数。
"""
metadata = filter or {}
embedding = self.embedding_function.embed_query(query)
results = self.vlite.retrieve(
text=query,
top_k=k,
metadata=metadata,
return_scores=True,
embedding=embedding,
)
documents_with_scores = [
(Document(page_content=text, metadata=metadata), score)
for text, score, metadata in results
]
return documents_with_scores
[docs] def update_document(self, document_id: str, document: Document) -> None:
"""更新向量存储中的现有文档。"""
self.vlite.update(
document_id, text=document.page_content, metadata=document.metadata
)
[docs] def get(self, ids: List[str]) -> List[Document]:
"""通过它们的ID获取文档。"""
results = self.vlite.get(ids)
documents = [
Document(page_content=text, metadata=metadata) for text, metadata in results
]
return documents
[docs] def delete(self, ids: Optional[List[str]] = None, **kwargs: Any) -> Optional[bool]:
"""根据ID删除。"""
if ids is not None:
self.vlite.delete(ids, **kwargs)
return True
return None
[docs] @classmethod
def from_existing_index(
cls,
embedding: Embeddings,
collection: str,
**kwargs: Any,
) -> VLite:
"""加载现有的VLite索引。
参数：
embedding：嵌入函数
collection：要加载的集合的名称。
返回：
VLite向量存储。
"""
vlite = cls(embedding_function=embedding, collection=collection, **kwargs)
return vlite
[docs] @classmethod
def from_texts(
cls,
texts: List[str],
embedding: Embeddings,
metadatas: Optional[List[dict]] = None,
collection: Optional[str] = None,
**kwargs: Any,
) -> VLite:
"""从原始文档构建VLite包装器。
这是一个用户友好的接口，可以：
1. 嵌入文档。
2. 将文档添加到向量存储。
这旨在是一个快速入门的方式。
示例：
.. code-block:: python
from langchain import VLite
from langchain.embeddings import OpenAIEmbeddings
embeddings = OpenAIEmbeddings()
vlite = VLite.from_texts(texts, embeddings)
"""
vlite = cls(embedding_function=embedding, collection=collection, **kwargs)
vlite.add_texts(texts, metadatas, **kwargs)
return vlite
[docs] @classmethod
def from_documents(
cls,
documents: List[Document],
embedding: Embeddings,
collection: Optional[str] = None,
**kwargs: Any,
) -> VLite:
"""从文档列表构建VLite包装器。
这是一个用户友好的接口，它：
1. 嵌入文档。
2. 将文档添加到向量存储。
这旨在是一个快速入门的方式。
示例：
.. code-block:: python
from langchain import VLite
from langchain.embeddings import OpenAIEmbeddings
embeddings = OpenAIEmbeddings()
vlite = VLite.from_documents(documents, embeddings)
"""
vlite = cls(embedding_function=embedding, collection=collection, **kwargs)
vlite.add_documents(documents, **kwargs)
return vlite