codekingpro/portable-devtools

1import asyncio2import json3import logging4from typing import Any, Dict, List, Optional5 6from langchain_core.documents import Document7 8from langchain_community.document_loaders.base import BaseLoader9 10logger = logging.getLogger(__name__)11 12 13class SurrealDBLoader(BaseLoader):14    """Load SurrealDB documents."""15 16    def __init__(17        self,18        filter_criteria: Optional[Dict] = None,19        **kwargs: Any,20    ) -> None:21        try:22            from surrealdb import Surreal23        except ImportError as e:24            raise ImportError(25                """Cannot import from surrealdb.26                please install with `pip install surrealdb`."""27            ) from e28 29        self.dburl = kwargs.pop("dburl", "ws://localhost:8000/rpc")30 31        if self.dburl[0:2] == "ws":32            self.sdb = Surreal(self.dburl)33        else:34            raise ValueError("Only websocket connections are supported at this time.")35 36        self.filter_criteria = filter_criteria or {}37 38        if "table" in self.filter_criteria:39            raise ValueError(40                "key `table` is not a valid criteria for `filter_criteria` argument."41            )42 43        self.ns = kwargs.pop("ns", "langchain")44        self.db = kwargs.pop("db", "database")45        self.table = kwargs.pop("table", "documents")46        self.sdb = Surreal(self.dburl)47        self.kwargs = kwargs48 49    async def initialize(self) -> None:50        """51        Initialize connection to surrealdb database52        and authenticate if credentials are provided53        """54        await self.sdb.connect()55        if "db_user" in self.kwargs and "db_pass" in self.kwargs:56            user = self.kwargs.get("db_user")57            password = self.kwargs.get("db_pass")58            await self.sdb.signin({"user": user, "pass": password})59 60        await self.sdb.use(self.ns, self.db)61 62    def load(self) -> List[Document]:63        async def _load() -> List[Document]:64            await self.initialize()65            return await self.aload()66 67        return asyncio.run(_load())68 69    async def aload(self) -> List[Document]:70        """Load data into Document objects."""71 72        query = "SELECT * FROM type::table($table)"73        if self.filter_criteria is not None and len(self.filter_criteria) > 0:74            query += " WHERE "75            for idx, key in enumerate(self.filter_criteria):76                query += f""" {"AND" if idx > 0 else ""} {key} = ${key}"""77 78        metadata = {79            "ns": self.ns,80            "db": self.db,81            "table": self.table,82        }83        results = await self.sdb.query(84            query, {"table": self.table, **self.filter_criteria}85        )86 87        return [88            (89                Document(90                    page_content=json.dumps(result),91                    metadata={"id": result["id"], **result["metadata"], **metadata},92                )93            )94            for result in results[0]["result"]95        ]96