from __future__ import annotations
import logging
import uuid
from typing import Any, Iterable, List, Optional, Type
import numpy as np
from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings
from langchain_core.vectorstores import VectorStore
logger = logging.getLogger(__name__)
[docs]class AtlasDB(VectorStore):
"""Atlas向量存储。
Atlas是Nomic的神经数据库和根茎仪器。
要使用，您应该安装``nomic`` python包。
示例：
.. code-block:: python
from langchain_community.vectorstores import AtlasDB
from langchain_community.embeddings.openai import OpenAIEmbeddings
embeddings = OpenAIEmbeddings()
vectorstore = AtlasDB("my_project", embeddings.embed_query)
"""
_ATLAS_DEFAULT_ID_FIELD = "atlas_id"
[docs] def __init__(
self,
name: str,
embedding_function: Optional[Embeddings] = None,
api_key: Optional[str] = None,
description: str = "A description for your project",
is_public: bool = True,
reset_project_if_exists: bool = False,
) -> None:
"""初始化Atlas客户端
参数:
name (str): 你的项目名称。如果项目已经存在，将被加载。
embedding_function (Optional[Embeddings]): 用于嵌入数据的可选函数。如果为None，则数据将使用Nomic的嵌入模型进行嵌入。
api_key (str): 你的Nomic API密钥
description (str): 项目的描述。
is_public (bool): 你的项目是否可以公开访问。默认为True。
reset_project_if_exists (bool): 如果项目已经存在，是否重置该项目。默认为False。
通常在开发和测试过程中很有用。
"""
try:
import nomic
from nomic import AtlasProject
except ImportError:
raise ImportError(
"Could not import nomic python package. "
"Please install it with `pip install nomic`."
)
if api_key is None:
raise ValueError("No API key provided. Sign up at atlas.nomic.ai!")
nomic.login(api_key)
self._embedding_function = embedding_function
modality = "text"
if self._embedding_function is not None:
modality = "embedding"
# Check if the project exists, create it if not
self.project = AtlasProject(
name=name,
description=description,
modality=modality,
is_public=is_public,
reset_project_if_exists=reset_project_if_exists,
unique_id_field=AtlasDB._ATLAS_DEFAULT_ID_FIELD,
)
self.project._latest_project_state()
@property
def embeddings(self) -> Optional[Embeddings]:
return self._embedding_function
[docs] def add_texts(
self,
texts: Iterable[str],
metadatas: Optional[List[dict]] = None,
ids: Optional[List[str]] = None,
refresh: bool = True,
**kwargs: Any,
) -> List[str]:
"""运行更多的文本通过嵌入并添加到向量存储。
参数：
texts (Iterable[str]): 要添加到向量存储的文本。
metadatas (Optional[List[dict]], optional): 元数据的可选列表。
ids (Optional[List[str]]): 一个可选的ID列表。
refresh(bool): 是否刷新索引以更新数据。默认为True。
返回：
List[str]: 添加的文本的ID列表。
"""
if (
metadatas is not None
and len(metadatas) > 0
and "text" in metadatas[0].keys()
):
raise ValueError("Cannot accept key text in metadata!")
texts = list(texts)
if ids is None:
ids = [str(uuid.uuid4()) for _ in texts]
# Embedding upload case
if self._embedding_function is not None:
_embeddings = self._embedding_function.embed_documents(texts)
embeddings = np.stack(_embeddings)
if metadatas is None:
data = [
{AtlasDB._ATLAS_DEFAULT_ID_FIELD: ids[i], "text": texts[i]}
for i, _ in enumerate(texts)
]
else:
for i in range(len(metadatas)):
metadatas[i][AtlasDB._ATLAS_DEFAULT_ID_FIELD] = ids[i]
metadatas[i]["text"] = texts[i]
data = metadatas
self.project._validate_map_data_inputs(
[], id_field=AtlasDB._ATLAS_DEFAULT_ID_FIELD, data=data
)
with self.project.wait_for_project_lock():
self.project.add_embeddings(embeddings=embeddings, data=data)
# Text upload case
else:
if metadatas is None:
data = [
{"text": text, AtlasDB._ATLAS_DEFAULT_ID_FIELD: ids[i]}
for i, text in enumerate(texts)
]
else:
for i, text in enumerate(texts):
metadatas[i]["text"] = texts
metadatas[i][AtlasDB._ATLAS_DEFAULT_ID_FIELD] = ids[i]
data = metadatas
self.project._validate_map_data_inputs(
[], id_field=AtlasDB._ATLAS_DEFAULT_ID_FIELD, data=data
)
with self.project.wait_for_project_lock():
self.project.add_text(data)
if refresh:
if len(self.project.indices) > 0:
with self.project.wait_for_project_lock():
self.project.rebuild_maps()
return ids
[docs] def create_index(self, **kwargs: Any) -> Any:
"""在您的项目中创建一个索引。
详细信息请参见
https://docs.nomic.ai/atlas_api.html#nomic.project.AtlasProject.create_index
"""
with self.project.wait_for_project_lock():
return self.project.create_index(**kwargs)
[docs] def similarity_search(
self,
query: str,
k: int = 4,
**kwargs: Any,
) -> List[Document]:
"""使用AtlasDB运行相似性搜索
参数：
query (str): 要搜索的查询文本。
k (int): 要返回的结果数量。默认为4。
返回：
List[Document]: 与查询文本最相似的文档列表。
"""
if self._embedding_function is None:
raise NotImplementedError(
"AtlasDB requires an embedding_function for text similarity search!"
)
_embedding = self._embedding_function.embed_documents([query])[0]
embedding = np.array(_embedding).reshape(1, -1)
with self.project.wait_for_project_lock():
neighbors, _ = self.project.projections[0].vector_search(
queries=embedding, k=k
)
data = self.project.get_data(ids=neighbors[0])
docs = [
Document(page_content=data[i]["text"], metadata=data[i])
for i, neighbor in enumerate(neighbors)
]
return docs
[docs] @classmethod
def from_texts(
cls: Type[AtlasDB],
texts: List[str],
embedding: Optional[Embeddings] = None,
metadatas: Optional[List[dict]] = None,
ids: Optional[List[str]] = None,
name: Optional[str] = None,
api_key: Optional[str] = None,
description: str = "A description for your project",
is_public: bool = True,
reset_project_if_exists: bool = False,
index_kwargs: Optional[dict] = None,
**kwargs: Any,
) -> AtlasDB:
"""从原始文档创建一个AtlasDB向量存储。
参数：
texts（List[str]）：要导入的文本列表。
name（str）：要创建的项目名称。
api_key（str）：您的nomic API密钥。
embedding（Optional[Embeddings]）：嵌入函数。默认为None。
metadatas（Optional[List[dict]]）：元数据列表。默认为None。
ids（Optional[List[str]）：文档ID的可选列表。如果为None，
将自动创建ids。
description（str）：项目的描述。
is_public（bool）：您的项目是否可以公开访问。
默认为True。
reset_project_if_exists（bool）：如果项目已经存在，是否重置该项目。
默认为False。
在开发和测试过程中通常很有用。
index_kwargs（Optional[dict]）：用于索引创建的kwargs字典。
请参阅https://docs.nomic.ai/atlas_api.html
返回：
AtlasDB：Nomic的神经数据库和最精致的根茎仪器
"""
if name is None or api_key is None:
raise ValueError("`name` and `api_key` cannot be None.")
# Inject relevant kwargs
all_index_kwargs = {"name": name + "_index", "indexed_field": "text"}
if index_kwargs is not None:
for k, v in index_kwargs.items():
all_index_kwargs[k] = v
# Build project
atlasDB = cls(
name,
embedding_function=embedding,
api_key=api_key,
description="A description for your project",
is_public=is_public,
reset_project_if_exists=reset_project_if_exists,
)
with atlasDB.project.wait_for_project_lock():
atlasDB.add_texts(texts=texts, metadatas=metadatas, ids=ids)
atlasDB.create_index(**all_index_kwargs)
return atlasDB
[docs] @classmethod
def from_documents(
cls: Type[AtlasDB],
documents: List[Document],
embedding: Optional[Embeddings] = None,
ids: Optional[List[str]] = None,
name: Optional[str] = None,
api_key: Optional[str] = None,
persist_directory: Optional[str] = None,
description: str = "A description for your project",
is_public: bool = True,
reset_project_if_exists: bool = False,
index_kwargs: Optional[dict] = None,
**kwargs: Any,
) -> AtlasDB:
"""从文档列表中创建一个AtlasDB向量存储。
参数:
name (str): 要创建的集合的名称。
api_key (str): 您的nomic API密钥。
documents (List[Document]): 要添加到向量存储的文档列表。
embedding (Optional[Embeddings]): 嵌入函数。默认为None。
ids (Optional[List[str]]): 文档ID的可选列表。如果为None，
将自动创建ID。
description (str): 项目的描述。
is_public (bool): 项目是否可以公开访问。
默认为True。
reset_project_if_exists (bool): 如果项目已经存在，是否重置此项目。
默认为False。
通常在开发和测试过程中很有用。
index_kwargs (Optional[dict]): 用于索引创建的kwargs字典。
请参阅https://docs.nomic.ai/atlas_api.html
返回:
AtlasDB: Nomic的神经数据库和最精致的根茎仪器
"""
if name is None or api_key is None:
raise ValueError("`name` and `api_key` cannot be None.")
texts = [doc.page_content for doc in documents]
metadatas = [doc.metadata for doc in documents]
return cls.from_texts(
name=name,
api_key=api_key,
texts=texts,
embedding=embedding,
metadatas=metadatas,
ids=ids,
description=description,
is_public=is_public,
reset_project_if_exists=reset_project_if_exists,
index_kwargs=index_kwargs,
)