codekingpro/portable-devtools

1import os2from typing import Any, Dict, Iterable, List, Optional, Type3 4from langchain_core.documents import Document5from langchain_core.embeddings import Embeddings6from langchain_core.vectorstores import VST, VectorStore7 8FIELD_TYPES = {9    "f": "files",10    "t": "texts",11    "l": "links",12}13 14 15class NucliaDB(VectorStore):16    """NucliaDB vector store."""17 18    _config: Dict[str, Any] = {}19 20    def __init__(21        self,22        knowledge_box: str,23        local: bool,24        api_key: Optional[str] = None,25        backend: Optional[str] = None,26    ) -> None:27        """Initialize the NucliaDB client.28 29        Args:30            knowledge_box: the Knowledge Box id.31            local: Whether to use a local NucliaDB instance or Nuclia Cloud32            api_key: A contributor API key for the kb (needed when local is False)33            backend: The backend url to use when local is True, defaults to34            http://localhost:808035        """36        try:37            from nuclia.sdk import NucliaAuth38        except ImportError:39            raise ImportError(40                "nuclia python package not found. "41                "Please install it with `pip install nuclia`."42            )43        self._config["LOCAL"] = local44        zone = os.environ.get("NUCLIA_ZONE", "europe-1")45        self._kb = knowledge_box46        if local:47            if not backend:48                backend = "http://localhost:8080"49            self._config["BACKEND"] = f"{backend}/api/v1"50            self._config["TOKEN"] = None51            NucliaAuth().nucliadb(url=backend)52            NucliaAuth().kb(url=self.kb_url, interactive=False)53        else:54            self._config["BACKEND"] = f"https://{zone}.nuclia.cloud/api/v1"55            self._config["TOKEN"] = api_key56            NucliaAuth().kb(57                url=self.kb_url, token=self._config["TOKEN"], interactive=False58            )59 60    @property61    def is_local(self) -> str:62        return self._config["LOCAL"]63 64    @property65    def kb_url(self) -> str:66        return f"{self._config['BACKEND']}/kb/{self._kb}"67 68    def add_texts(69        self,70        texts: Iterable[str],71        metadatas: Optional[List[dict]] = None,72        **kwargs: Any,73    ) -> List[str]:74        """Upload texts to NucliaDB"""75        ids = []76        from nuclia.sdk import NucliaResource77 78        factory = NucliaResource()79        for i, text in enumerate(texts):80            extra: Dict[str, Any] = {"metadata": ""}81            if metadatas:82                extra = {"metadata": metadatas[i]}83            id = factory.create(84                texts={"text": {"body": text}},85                extra=extra,86                url=self.kb_url,87                api_key=self._config["TOKEN"],88            )89            ids.append(id)90        return ids91 92    def delete(self, ids: Optional[List[str]] = None, **kwargs: Any) -> Optional[bool]:93        if not ids:94            return None95        from nuclia.sdk import NucliaResource96 97        factory = NucliaResource()98        results: List[bool] = []99        for id in ids:100            try:101                factory.delete(rid=id, url=self.kb_url, api_key=self._config["TOKEN"])102                results.append(True)103            except ValueError:104                results.append(False)105        return all(results)106 107    def similarity_search(108        self, query: str, k: int = 4, **kwargs: Any109    ) -> List[Document]:110        from nuclia.sdk import NucliaSearch111        from nucliadb_models.search import FindRequest, ResourceProperties112 113        request = FindRequest(114            query=query,115            page_size=k,116            show=[ResourceProperties.VALUES, ResourceProperties.EXTRA],117        )118        search = NucliaSearch()119        results = search.find(120            query=request, url=self.kb_url, api_key=self._config["TOKEN"]121        )122        paragraphs = []123        for resource in results.resources.values():124            for field in resource.fields.values():125                for paragraph_id, paragraph in field.paragraphs.items():126                    info = paragraph_id.split("/")127                    field_type = FIELD_TYPES.get(info[1], None)128                    field_id = info[2]129                    if not field_type:130                        continue131                    value = getattr(resource.data, field_type, {}).get(field_id, None)132                    paragraphs.append(133                        {134                            "text": paragraph.text,135                            "metadata": {136                                "extra": getattr(137                                    getattr(resource, "extra", {}), "metadata", None138                                ),139                                "value": value,140                            },141                            "order": paragraph.order,142                        }143                    )144        sorted_paragraphs = sorted(paragraphs, key=lambda x: x["order"])145        return [146            Document(page_content=paragraph["text"], metadata=paragraph["metadata"])147            for paragraph in sorted_paragraphs148        ]149 150    @classmethod151    def from_texts(152        cls: Type[VST],153        texts: List[str],154        embedding: Embeddings,155        metadatas: Optional[List[dict]] = None,156        **kwargs: Any,157    ) -> VST:158        """Return VectorStore initialized from texts and embeddings."""159        raise NotImplementedError160