# mypy: disable-error-code=func-returns-value
from __future__ import annotations
import json
import uuid
from typing import Any, Iterable, List, Optional, Type
from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings
from langchain_core.vectorstores import VST, VectorStore
[docs]class DuckDB(VectorStore):
"""`DuckDB` 矢量存储。
该类提供了一个矢量存储接口，用于添加文本并使用 DuckDB 执行相似性搜索。
有关 DuckDB 的更多信息，请参见：https://duckdb.org/
此集成需要 `duckdb` Python 包。
您可以使用 `pip install duckdb` 进行安装。
*安全提示*：默认的 DuckDB 配置不安全。
默认情况下，DuckDB 可以与整个文件系统中的文件交互，包括读取、写入和列出文件和目录的能力。
它还可以访问全局命名空间中存在的一些 Python 变量。
在使用此 DuckDB 矢量存储时，建议您使用安全配置初始化 DuckDB 连接。
例如，您可以在连接配置中将 `enable_external_access` 设置为 `false`，以禁用对 DuckDB 连接的外部访问。
您可以在此处查看 DuckDB 配置选项：
https://duckdb.org/docs/configuration/overview.html
请查阅 DuckDB 文档中的其他相关安全注意事项。（例如，"autoinstall_known_extensions": "false"，"autoload_known_extensions": "false"）
有关更多信息，请参阅 https://python.langchain.com/docs/security。
Args:
connection: 可选的 DuckDB 连接
embedding: 用于生成嵌入的嵌入函数或模型。
vector_key: 用于存储向量的列名。默认为 `embedding`。
id_key: 用于存储唯一标识符的列名。默认为 `id`。
text_key: 用于存储文本的列名。默认为 `text`。
table_name: 用于存储嵌入的表的名称。默认为 `embeddings`。
示例:
.. code-block:: python
import duckdb
conn = duckdb.connect(database=':memory:',
config={
# 限制某些 DuckDB 功能的示例配置
# 列出不详尽。请查阅 DuckDB 文档。
"enable_external_access": "false",
"autoinstall_known_extensions": "false",
"autoload_known_extensions": "false"
}
)
embedding_function = ... # 在此处定义或导入您的嵌入函数
vector_store = DuckDB(conn, embedding_function)
vector_store.add_texts(['text1', 'text2'])
result = vector_store.similarity_search('text1')"""
[docs] def __init__(
self,
*,
connection: Optional[Any] = None,
embedding: Embeddings,
vector_key: str = "embedding",
id_key: str = "id",
text_key: str = "text",
table_name: str = "vectorstore",
):
"""使用DuckDB连接进行初始化，并设置为向量存储。"""
try:
import duckdb
except ImportError:
raise ImportError(
"Could not import duckdb package. "
"Please install it with `pip install duckdb`."
)
self.duckdb = duckdb
self._embedding = embedding
self._vector_key = vector_key
self._id_key = id_key
self._text_key = text_key
self._table_name = table_name
if self._embedding is None:
raise ValueError("An embedding function or model must be provided.")
if connection is None:
import warnings
warnings.warn(
"No DuckDB connection provided. A new connection will be created."
"This connection is running in memory and no data will be persisted."
"To persist data, specify `connection=duckdb.connect(...)` when using "
"the API. Please review the documentation of the vectorstore for "
"security recommendations on configuring the connection."
)
self._connection = connection or self.duckdb.connect(
database=":memory:", config={"enable_external_access": "false"}
)
self._ensure_table()
self._table = self._connection.table(self._table_name)
@property
def embeddings(self) -> Optional[Embeddings]:
"""返回向量存储中使用的嵌入对象。"""
return self._embedding
[docs] def add_texts(
self,
texts: Iterable[str],
metadatas: Optional[List[dict]] = None,
**kwargs: Any,
) -> List[str]:
"""将文本转换为嵌入向量，并使用Pandas DataFrame 将其添加到数据库中
参数：
texts: 要添加到向量存储中的字符串的可迭代对象。
metadatas: 与文本相关的元数据的可选列表。
kwargs: 包括可选的与文本关联的 'ids' 在内的其他参数。
返回：
添加的文本的id列表。
"""
# Extract ids from kwargs or generate new ones if not provided
ids = kwargs.pop("ids", [str(uuid.uuid4()) for _ in texts])
# Embed texts and create documents
ids = ids or [str(uuid.uuid4()) for _ in texts]
embeddings = self._embedding.embed_documents(list(texts))
for idx, text in enumerate(texts):
embedding = embeddings[idx]
# Serialize metadata if present, else default to None
metadata = (
json.dumps(metadatas[idx])
if metadatas and idx < len(metadatas)
else None
)
self._connection.execute(
f"INSERT INTO {self._table_name} VALUES (?,?,?,?)",
[ids[idx], text, embedding, metadata],
)
return ids
[docs] def similarity_search(
self, query: str, k: int = 4, **kwargs: Any
) -> List[Document]:
"""执行给定查询字符串的相似性搜索。
参数：
query：要搜索的查询字符串。
k：要返回的相似文本数量。
返回：
与查询最相似的文档列表。
"""
embedding = self._embedding.embed_query(query) # type: ignore
list_cosine_similarity = self.duckdb.FunctionExpression(
"list_cosine_similarity",
self.duckdb.ColumnExpression(self._vector_key),
self.duckdb.ConstantExpression(embedding),
)
docs = (
self._table.select(
*[
self.duckdb.StarExpression(exclude=[]),
list_cosine_similarity.alias("similarity"),
]
)
.order("similarity desc")
.limit(k)
.select(
self.duckdb.StarExpression(exclude=["similarity", self._vector_key])
)
.fetchdf()
)
return [
Document(
page_content=docs[self._text_key][idx],
metadata=json.loads(docs["metadata"][idx])
if docs["metadata"][idx]
else {},
)
for idx in range(len(docs))
]
[docs] @classmethod
def from_texts(
cls: Type[VST],
texts: List[str],
embedding: Embeddings,
metadatas: Optional[List[dict]] = None,
**kwargs: Any,
) -> DuckDB:
"""创建一个DuckDB实例，并用文本及其嵌入填充它。
参数:
texts: 要添加到向量存储中的字符串列表。
embedding: 用于生成嵌入的嵌入函数或模型。
metadatas: 与文本相关的元数据字典的可选列表。
**kwargs: 包括以下附加关键字参数:
- connection: DuckDB连接。如果未提供，将创建一个新连接。
- vector_key: 用于存储向量的列名。默认为"vector"。
- id_key: 用于存储唯一标识符的列名。默认为"id"。
- text_key: 用于存储文本的列名。默认为"text"。
- table_name: 用于存储嵌入的表的名称。默认为"embeddings"。
返回:
包含提供的文本及其嵌入的DuckDB实例。
"""
# Extract kwargs for DuckDB instance creation
connection = kwargs.get("connection", None)
vector_key = kwargs.get("vector_key", "vector")
id_key = kwargs.get("id_key", "id")
text_key = kwargs.get("text_key", "text")
table_name = kwargs.get("table_name", "embeddings")
# Create an instance of DuckDB
instance = DuckDB(
connection=connection,
embedding=embedding,
vector_key=vector_key,
id_key=id_key,
text_key=text_key,
table_name=table_name,
)
# Add texts and their embeddings to the DuckDB vector store
instance.add_texts(texts, metadatas=metadatas, **kwargs)
return instance
def _ensure_table(self) -> None:
"""确保用于存储嵌入的表存在。"""
create_table_sql = f"""
CREATE TABLE IF NOT EXISTS {self._table_name} (
{self._id_key} VARCHAR PRIMARY KEY,
{self._text_key} VARCHAR,
{self._vector_key} FLOAT[],
metadata VARCHAR
)
"""
self._connection.execute(create_table_sql)