codekingpro/portable-devtools

1from __future__ import annotations2 3import logging4import uuid5from typing import Any, Iterable, List, Optional, Type6 7import numpy as np8from langchain_core.documents import Document9from langchain_core.embeddings import Embeddings10from langchain_core.vectorstores import VectorStore11 12logger = logging.getLogger(__name__)13 14 15class AtlasDB(VectorStore):16    """`Atlas` vector store.17 18     Atlas is the `Nomic's` neural database and `rhizomatic` instrument.19 20    To use, you should have the ``nomic`` python package installed.21 22    Example:23        .. code-block:: python24 25                from langchain_community.vectorstores import AtlasDB26                from langchain_community.embeddings.openai import OpenAIEmbeddings27 28                embeddings = OpenAIEmbeddings()29                vectorstore = AtlasDB("my_project", embeddings.embed_query)30    """31 32    _ATLAS_DEFAULT_ID_FIELD: str = "atlas_id"33 34    def __init__(35        self,36        name: str,37        embedding_function: Optional[Embeddings] = None,38        api_key: Optional[str] = None,39        description: str = "A description for your project",40        is_public: bool = True,41        reset_project_if_exists: bool = False,42    ) -> None:43        """44        Initialize the Atlas Client45 46        Args:47            name (str): The name of your project. If the project already exists,48                it will be loaded.49            embedding_function (Optional[Embeddings]): An optional function used for50                embedding your data. If None, data will be embedded with51                Nomic's embed model.52            api_key (str): Your nomic API key53            description (str): A description for your project.54            is_public (bool): Whether your project is publicly accessible.55                True by default.56            reset_project_if_exists (bool): Whether to reset this project if it57                already exists. Default False.58                Generally useful during development and testing.59        """60        try:61            import nomic62            from nomic import AtlasProject63        except ImportError:64            raise ImportError(65                "Could not import nomic python package. "66                "Please install it with `pip install nomic`."67            )68 69        if api_key is None:70            raise ValueError("No API key provided. Sign up at atlas.nomic.ai!")71        nomic.login(api_key)72 73        self._embedding_function = embedding_function74        modality = "text"75        if self._embedding_function is not None:76            modality = "embedding"77 78        # Check if the project exists, create it if not79        self.project = AtlasProject(80            name=name,81            description=description,82            modality=modality,83            is_public=is_public,84            reset_project_if_exists=reset_project_if_exists,85            unique_id_field=AtlasDB._ATLAS_DEFAULT_ID_FIELD,86        )87        self.project._latest_project_state()88 89    @property90    def embeddings(self) -> Optional[Embeddings]:91        return self._embedding_function92 93    def add_texts(94        self,95        texts: Iterable[str],96        metadatas: Optional[List[dict]] = None,97        ids: Optional[List[str]] = None,98        refresh: bool = True,99        **kwargs: Any,100    ) -> List[str]:101        """Run more texts through the embeddings and add to the vectorstore.102 103        Args:104            texts (Iterable[str]): Texts to add to the vectorstore.105            metadatas (Optional[List[dict]], optional): Optional list of metadatas.106            ids (Optional[List[str]]): An optional list of ids.107            refresh(bool): Whether or not to refresh indices with the updated data.108                Default True.109        Returns:110            List[str]: List of IDs of the added texts.111        """112 113        if (114            metadatas is not None115            and len(metadatas) > 0116            and "text" in metadatas[0].keys()117        ):118            raise ValueError("Cannot accept key text in metadata!")119 120        texts = list(texts)121        if ids is None:122            ids = [str(uuid.uuid4()) for _ in texts]123 124        # Embedding upload case125        if self._embedding_function is not None:126            _embeddings = self._embedding_function.embed_documents(texts)127            embeddings = np.stack(_embeddings)128            if metadatas is None:129                data = [130                    {AtlasDB._ATLAS_DEFAULT_ID_FIELD: ids[i], "text": texts[i]}131                    for i, _ in enumerate(texts)132                ]133            else:134                for i in range(len(metadatas)):135                    metadatas[i][AtlasDB._ATLAS_DEFAULT_ID_FIELD] = ids[i]136                    metadatas[i]["text"] = texts[i]137                data = metadatas138 139            self.project._validate_map_data_inputs(140                [], id_field=AtlasDB._ATLAS_DEFAULT_ID_FIELD, data=data141            )142            with self.project.wait_for_project_lock():143                self.project.add_embeddings(embeddings=embeddings, data=data)144        # Text upload case145        else:146            if metadatas is None:147                data = [148                    {"text": text, AtlasDB._ATLAS_DEFAULT_ID_FIELD: ids[i]}149                    for i, text in enumerate(texts)150                ]151            else:152                for i, text in enumerate(texts):153                    metadatas[i]["text"] = texts154                    metadatas[i][AtlasDB._ATLAS_DEFAULT_ID_FIELD] = ids[i]155                data = metadatas156 157            self.project._validate_map_data_inputs(158                [], id_field=AtlasDB._ATLAS_DEFAULT_ID_FIELD, data=data159            )160 161            with self.project.wait_for_project_lock():162                self.project.add_text(data)163 164        if refresh:165            if len(self.project.indices) > 0:166                with self.project.wait_for_project_lock():167                    self.project.rebuild_maps()168 169        return ids170 171    def create_index(self, **kwargs: Any) -> Any:172        """Creates an index in your project.173 174        See175        https://docs.nomic.ai/atlas_api.html#nomic.project.AtlasProject.create_index176        for full detail.177        """178        with self.project.wait_for_project_lock():179            return self.project.create_index(**kwargs)180 181    def similarity_search(182        self,183        query: str,184        k: int = 4,185        **kwargs: Any,186    ) -> List[Document]:187        """Run similarity search with AtlasDB188 189        Args:190            query (str): Query text to search for.191            k (int): Number of results to return. Defaults to 4.192 193        Returns:194            List[Document]: List of documents most similar to the query text.195        """196        if self._embedding_function is None:197            raise NotImplementedError(198                "AtlasDB requires an embedding_function for text similarity search!"199            )200 201        _embedding = self._embedding_function.embed_documents([query])[0]202        embedding = np.array(_embedding).reshape(1, -1)203        with self.project.wait_for_project_lock():204            neighbors, _ = self.project.projections[0].vector_search(205                queries=embedding, k=k206            )207            data = self.project.get_data(ids=neighbors[0])208 209        docs = [210            Document(page_content=data[i]["text"], metadata=data[i])211            for i, neighbor in enumerate(neighbors)212        ]213        return docs214 215    @classmethod216    def from_texts(217        cls: Type[AtlasDB],218        texts: List[str],219        embedding: Optional[Embeddings] = None,220        metadatas: Optional[List[dict]] = None,221        ids: Optional[List[str]] = None,222        name: Optional[str] = None,223        api_key: Optional[str] = None,224        description: str = "A description for your project",225        is_public: bool = True,226        reset_project_if_exists: bool = False,227        index_kwargs: Optional[dict] = None,228        **kwargs: Any,229    ) -> AtlasDB:230        """Create an AtlasDB vectorstore from a raw documents.231 232        Args:233            texts (List[str]): The list of texts to ingest.234            name (str): Name of the project to create.235            api_key (str): Your nomic API key,236            embedding (Optional[Embeddings]): Embedding function. Defaults to None.237            metadatas (Optional[List[dict]]): List of metadatas. Defaults to None.238            ids (Optional[List[str]]): Optional list of document IDs. If None,239                ids will be auto created240            description (str): A description for your project.241            is_public (bool): Whether your project is publicly accessible.242                True by default.243            reset_project_if_exists (bool): Whether to reset this project if it244                already exists. Default False.245                Generally useful during development and testing.246            index_kwargs (Optional[dict]): Dict of kwargs for index creation.247                See https://docs.nomic.ai/atlas_api.html248 249        Returns:250            AtlasDB: Nomic's neural database and finest rhizomatic instrument251        """252        if name is None or api_key is None:253            raise ValueError("`name` and `api_key` cannot be None.")254 255        # Inject relevant kwargs256        all_index_kwargs = {"name": name + "_index", "indexed_field": "text"}257        if index_kwargs is not None:258            for k, v in index_kwargs.items():259                all_index_kwargs[k] = v260 261        # Build project262        atlasDB = cls(263            name,264            embedding_function=embedding,265            api_key=api_key,266            description="A description for your project",267            is_public=is_public,268            reset_project_if_exists=reset_project_if_exists,269        )270        with atlasDB.project.wait_for_project_lock():271            atlasDB.add_texts(texts=texts, metadatas=metadatas, ids=ids)272            atlasDB.create_index(**all_index_kwargs)273        return atlasDB274 275    @classmethod276    def from_documents(277        cls: Type[AtlasDB],278        documents: List[Document],279        embedding: Optional[Embeddings] = None,280        ids: Optional[List[str]] = None,281        name: Optional[str] = None,282        api_key: Optional[str] = None,283        persist_directory: Optional[str] = None,284        description: str = "A description for your project",285        is_public: bool = True,286        reset_project_if_exists: bool = False,287        index_kwargs: Optional[dict] = None,288        **kwargs: Any,289    ) -> AtlasDB:290        """Create an AtlasDB vectorstore from a list of documents.291 292        Args:293            name (str): Name of the collection to create.294            api_key (str): Your nomic API key,295            documents (List[Document]): List of documents to add to the vectorstore.296            embedding (Optional[Embeddings]): Embedding function. Defaults to None.297            ids (Optional[List[str]]): Optional list of document IDs. If None,298                ids will be auto created299            description (str): A description for your project.300            is_public (bool): Whether your project is publicly accessible.301                True by default.302            reset_project_if_exists (bool): Whether to reset this project if303                it already exists. Default False.304                Generally useful during development and testing.305            index_kwargs (Optional[dict]): Dict of kwargs for index creation.306                See https://docs.nomic.ai/atlas_api.html307 308        Returns:309            AtlasDB: Nomic's neural database and finest rhizomatic instrument310        """311        if name is None or api_key is None:312            raise ValueError("`name` and `api_key` cannot be None.")313        texts = [doc.page_content for doc in documents]314        metadatas = [doc.metadata for doc in documents]315        return cls.from_texts(316            name=name,317            api_key=api_key,318            texts=texts,319            embedding=embedding,320            metadatas=metadatas,321            ids=ids,322            description=description,323            is_public=is_public,324            reset_project_if_exists=reset_project_if_exists,325            index_kwargs=index_kwargs,326        )327