naas-abi-core 1.4.1__py3-none-any.whl

naas_abi_core/services/triple_store/adaptors/secondary/TripleStoreService__SecondaryAdaptor__Filesystem.py

@@ -0,0 +1,223 @@
+import os
+from threading import Lock
+from typing import Any, Dict, List, Tuple
+
+import rdflib
+from naas_abi_core.services.triple_store.adaptors.secondary.base.TripleStoreService__SecondaryAdaptor__FileBase import (
+    TripleStoreService__SecondaryAdaptor__FileBase,
+)
+from naas_abi_core.services.triple_store.TripleStorePorts import (
+    Exceptions,
+    ITripleStorePort,
+    OntologyEvent,
+)
+from rdflib import RDFS, BNode, Graph, URIRef, query
+
+
+class TripleStoreService__SecondaryAdaptor__Filesystem(
+    ITripleStorePort, TripleStoreService__SecondaryAdaptor__FileBase
+):
+    __store_path: str
+    __triples_path: str
+
+    __live_graph: Graph
+
+    __lock: Lock
+
+    def __init__(self, store_path: str, triples_path: str = "triples"):
+        self.__store_path = store_path
+        self.__triples_path = triples_path
+        self.__lock = Lock()
+
+        os.makedirs(os.path.join(self.__store_path, self.__triples_path), exist_ok=True)
+
+        self.__live_graph = self.load()
+
+    def __merge_graphs(self, graphs: List[Graph]) -> Graph:
+        merged_graph = Graph()
+        for graph in graphs:
+            merged_graph += graph
+
+        return merged_graph
+
+    def hash_triples_path(self, hash_value: str) -> str:
+        return os.path.join(
+            self.__store_path,
+            "triples",
+            f"{hash_value}.ttl" if not hash_value.endswith(".ttl") else hash_value,
+        )
+
+    ## File System Methods
+
+    def insert(self, triples: Graph):
+        with self.__lock:
+            triples_by_subject: Dict[Any, List[Tuple[Any, Any]]] = (
+                self.triples_by_subject(triples)
+            )
+
+            for subject in triples_by_subject:
+                subject_hash = self.iri_hash(subject)
+
+                if not os.path.exists(self.hash_triples_path(subject_hash)):
+                    graph = Graph()
+
+                    for prefix, namespace in triples.namespaces():
+                        graph.bind(prefix, namespace)
+                else:
+                    graph = Graph().parse(
+                        self.hash_triples_path(subject_hash), format="turtle"
+                    )
+
+                for p, o in triples_by_subject[subject]:
+                    graph.add((subject, p, o))
+
+                graph.serialize(
+                    destination=self.hash_triples_path(subject_hash), format="turtle"
+                )
+
+            for prefix, namespace in triples.namespaces():
+                self.__live_graph.bind(prefix, namespace)
+
+            # Update the live graph
+            self.__live_graph += triples
+
+    def remove(self, triples: Graph):
+        with self.__lock:
+            triples_by_subject: Dict[Any, List[Tuple[Any, Any]]] = (
+                self.triples_by_subject(triples)
+            )
+
+            for subject in triples_by_subject:
+                subject_hash = self.iri_hash(subject)
+
+                if os.path.exists(self.hash_triples_path(subject_hash)):
+                    graph = Graph().parse(
+                        self.hash_triples_path(subject_hash), format="turtle"
+                    )
+
+                    for p, o in triples_by_subject[subject]:
+                        graph.remove((subject, p, o))
+
+                    graph.serialize(
+                        destination=self.hash_triples_path(subject_hash),
+                        format="turtle",
+                    )
+
+            # Update the live graph
+            self.__live_graph -= triples
+
+    ## Ontology Methods
+
+    def get(self) -> Graph:
+        return self.__live_graph
+
+    def get_subject_graph(self, subject: URIRef) -> Graph:
+        subject_hash = self.iri_hash(subject)
+
+        if os.path.exists(self.hash_triples_path(subject_hash)):
+            return Graph().parse(self.hash_triples_path(subject_hash), format="turtle")
+
+        raise Exceptions.SubjectNotFoundError(f"Subject {subject} not found")
+
+    def load(self) -> Graph:
+        with self.__lock:
+            triples = Graph()
+
+            for file in os.listdir(os.path.join(self.__store_path, "triples")):
+                try:
+                    g = Graph().parse(self.hash_triples_path(file), format="turtle")
+                except Exception as e:
+                    from naas_abi_core import logger
+
+                    logger.error(
+                        f"Error loading triples from {self.hash_triples_path(file)}: {e}"
+                    )
+                    raise e
+
+                for prefix, namespace in g.namespaces():
+                    triples.bind(prefix, namespace)
+
+                triples += g
+
+            return triples
+
+    def query(self, query: str) -> query.Result:
+        with self.__lock:
+            aggregate_graph = self.get()
+
+            return aggregate_graph.query(query)
+
+    def query_view(self, view: str, query: str) -> rdflib.query.Result:
+        with self.__lock:
+            if os.path.exists(os.path.join(self.__store_path, "views", view)):
+                aggregate_graph = Graph()
+
+                for file in os.listdir(os.path.join(self.__store_path, "views", view)):
+                    g = Graph().parse(
+                        os.path.join(self.__store_path, "views", view, file),
+                        format="turtle",
+                    )
+
+                    for prefix, namespace in g.namespaces():
+                        aggregate_graph.bind(prefix, namespace)
+
+                    aggregate_graph += g
+
+                return aggregate_graph.query(query)
+            else:
+                raise Exceptions.ViewNotFoundError(f"View {view} not found")
+
+    def handle_view_event(
+        self,
+        view: Tuple[URIRef | None, URIRef | None, URIRef | None],
+        event: OntologyEvent,
+        triple: Tuple[URIRef | None, URIRef | None, URIRef | None],
+    ):
+        s, _, o = triple
+
+        assert isinstance(s, BNode) or isinstance(s, URIRef), type(s)
+        assert isinstance(o, BNode) or isinstance(o, URIRef), type(o)
+
+        partition_hash = self.iri_hash(o)
+
+        if os.path.exists(self.hash_triples_path(partition_hash)):
+            graph = Graph().parse(
+                self.hash_triples_path(partition_hash), format="turtle"
+            )
+
+            label = graph.value(subject=o, predicate=RDFS.label)
+            object_id = str(o).split("/")[-1].split("#")[-1]
+
+            dir_name = f"{label}_{object_id}"
+
+            os.makedirs(
+                os.path.join(self.__store_path, "views", dir_name), exist_ok=True
+            )
+
+            if event == OntologyEvent.INSERT:
+                try:
+                    # Create symbolic link
+                    os.symlink(
+                        os.path.join("..", "..", "triples", f"{self.iri_hash(s)}.ttl"),
+                        os.path.join(
+                            self.__store_path,
+                            "views",
+                            dir_name,
+                            f"{self.iri_hash(s)}.ttl",
+                        ),
+                    )
+                except FileExistsError:
+                    pass
+            elif event == OntologyEvent.DELETE:
+                # Remove symbolic link
+                try:
+                    os.remove(
+                        os.path.join(
+                            self.__store_path,
+                            "views",
+                            dir_name,
+                            f"{self.iri_hash(s)}.ttl",
+                        )
+                    )
+                except FileNotFoundError:
+                    pass

naas_abi_core/services/triple_store/adaptors/secondary/TripleStoreService__SecondaryAdaptor__ObjectStorage.py

@@ -0,0 +1,234 @@
+import queue
+import time
+from threading import Lock, Thread
+from typing import Any, Dict, List, Tuple
+
+import rdflib
+from naas_abi_core import logger
+from naas_abi_core.services.object_storage.ObjectStoragePort import (
+    Exceptions as ObjectStorageExceptions,
+)
+from naas_abi_core.services.object_storage.ObjectStorageService import (
+    ObjectStorageService,
+)
+from naas_abi_core.services.triple_store.adaptors.secondary.base.TripleStoreService__SecondaryAdaptor__FileBase import (
+    TripleStoreService__SecondaryAdaptor__FileBase,
+)
+from naas_abi_core.services.triple_store.TripleStorePorts import (
+    Exceptions,
+    ITripleStorePort,
+    OntologyEvent,
+)
+from naas_abi_core.utils.Workers import Job, WorkerPool
+from rdflib import Graph, Node, URIRef, query
+
+
+class TripleStoreService__SecondaryAdaptor__ObjectStorage(
+    ITripleStorePort, TripleStoreService__SecondaryAdaptor__FileBase
+):
+    __object_storage_service: ObjectStorageService
+
+    __triples_prefix: str
+
+    __live_graph: Graph
+
+    __lock: Lock
+
+    def __init__(
+        self,
+        object_storage_service: ObjectStorageService,
+        triples_prefix: str = "triples",
+    ):
+        logger.debug("Initializing TripleStoreService__SecondaryAdaptor__ObjectStorage")
+        self.__object_storage_service = object_storage_service
+        self.__triples_prefix = triples_prefix
+
+        self.__lock = Lock()
+
+        self.__insert_pool = WorkerPool(num_workers=50)
+
+        self.__live_graph = self.load()
+
+    def load_triples(self, subject_hash: str) -> Graph:
+        obj: bytes = self.__object_storage_service.get_object(
+            prefix=self.__triples_prefix, key=f"{subject_hash}.ttl"
+        )
+
+        content: str = obj.decode("utf-8")
+
+        return Graph().parse(data=str(content), format="turtle")
+
+    def store(self, name: str, triples: Graph):
+        serialized_triples = triples.serialize(format="turtle").encode("utf-8")
+        self.__object_storage_service.put_object(
+            prefix=self.__triples_prefix, key=f"{name}.ttl", content=serialized_triples
+        )
+
+    def insert(self, triples: Graph):
+        with self.__lock:
+            triples_by_subject: Dict[Node, List[Tuple[Node, Node]]] = (
+                self.triples_by_subject(triples)
+            )
+
+            def __insert(
+                subject: URIRef, triples_by_subject: Dict[Node, List[Tuple[Node, Node]]]
+            ):
+                subject_hash = self.iri_hash(subject)
+
+                try:
+                    graph = self.load_triples(subject_hash)
+                except ObjectStorageExceptions.ObjectNotFound:
+                    graph = Graph()
+                    for prefix, namespace in triples.namespaces():
+                        graph.bind(prefix, namespace)
+
+                for p, o in triples_by_subject[subject]:
+                    graph.add((subject, p, o))
+
+                self.store(subject_hash, graph)
+
+            jobs: List[Job] = [
+                Job(
+                    queue=None,
+                    func=__insert,
+                    subject=subject,
+                    triples_by_subject=triples_by_subject,
+                )
+                for subject in triples_by_subject
+            ]
+
+            result_queue: queue.Queue[Job] = self.__insert_pool.submit_all(jobs)
+
+            while result_queue.qsize() < result_queue.maxsize:
+                logger.debug(f"Inserting {result_queue.qsize()}/{result_queue.maxsize}")
+                time.sleep(0.1)
+
+            for prefix, namespace in triples.namespaces():
+                self.__live_graph.bind(prefix, namespace)
+
+            self.__live_graph += triples
+
+    def remove(self, triples: Graph):
+        with self.__lock:
+            triples_by_subject: Dict[Any, List[Tuple[Any, Any]]] = (
+                self.triples_by_subject(triples)
+            )
+
+            for subject in triples_by_subject:
+                subject_hash = self.iri_hash(subject)
+
+                try:
+                    graph = self.load_triples(subject_hash)
+                    for p, o in triples_by_subject[subject]:
+                        graph.add((subject, p, o))
+
+                    self.store(subject_hash, graph)
+                except ObjectStorageExceptions.ObjectNotFound:
+                    pass
+
+            for prefix, namespace in triples.namespaces():
+                self.__live_graph.bind(prefix, namespace)
+
+            self.__live_graph -= triples
+
+    def get_subject_graph(self, subject: str | URIRef) -> Graph:
+        subject_hash = (
+            self.iri_hash(URIRef(subject))
+            if isinstance(subject, str)
+            else self.iri_hash(subject)
+        )
+
+        try:
+            graph = self.load_triples(subject_hash)
+            return graph
+        except ObjectStorageExceptions.ObjectNotFound:
+            raise Exceptions.SubjectNotFoundError(f"Subject {subject} not found")
+
+    def load(self) -> Graph:
+        with self.__lock:
+            logger.debug("Loading triples from object storage")
+            triples = Graph()
+
+            # Queue to stream files as they are discovered
+            files_queue: queue.Queue[str] = queue.Queue()
+            worker_pool = WorkerPool(num_workers=50)
+            result_queue: queue.Queue[Job] = queue.Queue()
+            job_nbr = 0
+
+            def list_objects_worker(files_queue: queue.Queue):
+                try:
+                    self.__object_storage_service.list_objects(
+                        prefix=self.__triples_prefix, queue=files_queue
+                    )
+                except ObjectStorageExceptions.ObjectNotFound:
+                    pass
+
+            list_object_thread = Thread(
+                target=list_objects_worker, args=(files_queue,), daemon=True
+            )
+            list_object_thread.start()
+
+            # Process files as they are discovered
+            while not files_queue.empty() or list_object_thread.is_alive():
+                try:
+                    file: str = files_queue.get(timeout=1.0)
+
+                    file_hash = file.split("/")[-1].split(".")[0]
+                    worker_pool.submit(
+                        Job(
+                            queue=result_queue,
+                            func=self.load_triples,
+                            subject_hash=file_hash,
+                        )
+                    )
+                    job_nbr += 1
+                except queue.Empty:
+                    continue
+
+            # Consume completed jobs.
+            completed_jobs = 0
+            while completed_jobs < job_nbr:
+                logger.debug(f"{completed_jobs}/{job_nbr}")
+                try:
+                    job = result_queue.get(timeout=1.0)
+                    if job.is_completed():
+                        result_graph = job.get_result()
+                        triples += result_graph
+
+                        for prefix, namespace in result_graph.namespaces():
+                            triples.bind(prefix, namespace)
+
+                        completed_jobs += 1
+                except queue.Empty:
+                    continue
+
+            logger.debug(f"Loaded {len(triples)} triples")
+
+            logger.debug("Joining list objects thread")
+            list_object_thread.join()
+
+            logger.debug("Shutting down worker pool")
+            worker_pool.shutdown()
+
+            logger.debug("Done")
+
+            return triples
+
+    def get(self) -> Graph:
+        return self.__live_graph
+
+    def query(self, query: str) -> query.Result:
+        with self.__lock:
+            return self.get().query(query)
+
+    def query_view(self, view: str, query: str) -> rdflib.query.Result:
+        with self.__lock:
+            return self.get().query(query)
+
+    def handle_view_event(
+        self,
+        view: Tuple[URIRef | None, URIRef | None, URIRef | None],
+        event: OntologyEvent,
+        triple: Tuple[URIRef | None, URIRef | None, URIRef | None],
+    ):
+        pass

naas_abi_core/services/triple_store/adaptors/secondary/base/TripleStoreService__SecondaryAdaptor__FileBase.py

@@ -0,0 +1,18 @@
+import hashlib
+from rdflib import Graph, Node, URIRef
+from typing import Dict, List, Tuple
+
+
+class TripleStoreService__SecondaryAdaptor__FileBase:
+    def iri_hash(self, iri: URIRef) -> str:
+        return hashlib.sha256(iri.encode("utf-8")).hexdigest()
+
+    def triples_by_subject(self, triples: Graph) -> Dict[Node, List[Tuple[Node, Node]]]:
+        triples_by_subject: Dict[Node, List[Tuple[Node, Node]]] = {}
+
+        for s, p, o in triples.triples((None, None, None)):
+            if s not in triples_by_subject:
+                triples_by_subject[s] = []
+            triples_by_subject[s].append((p, o))
+
+        return triples_by_subject

naas_abi_core/services/vector_store/IVectorStorePort.py

@@ -0,0 +1,101 @@
+from abc import ABC, abstractmethod
+from typing import List, Dict, Any, Optional
+from dataclasses import dataclass
+import numpy as np
+
+
+@dataclass
+class VectorDocument:
+    id: str
+    vector: np.ndarray
+    metadata: Dict[str, Any]
+    payload: Optional[Dict[str, Any]] = None
+
+
+@dataclass
+class SearchResult:
+    id: str
+    score: float
+    vector: Optional[np.ndarray] = None
+    metadata: Optional[Dict[str, Any]] = None
+    payload: Optional[Dict[str, Any]] = None
+
+
+class IVectorStorePort(ABC):
+    @abstractmethod
+    def initialize(self) -> None:
+        pass
+
+    @abstractmethod
+    def create_collection(
+        self,
+        collection_name: str,
+        dimension: int,
+        distance_metric: str = "cosine",
+        **kwargs
+    ) -> None:
+        pass
+
+    @abstractmethod
+    def delete_collection(self, collection_name: str) -> None:
+        pass
+
+    @abstractmethod
+    def list_collections(self) -> List[str]:
+        pass
+
+    @abstractmethod
+    def store_vectors(
+        self,
+        collection_name: str,
+        documents: List[VectorDocument]
+    ) -> None:
+        pass
+
+    @abstractmethod
+    def search(
+        self,
+        collection_name: str,
+        query_vector: np.ndarray,
+        k: int = 10,
+        filter: Optional[Dict[str, Any]] = None,
+        include_vectors: bool = False,
+        include_metadata: bool = True
+    ) -> List[SearchResult]:
+        pass
+
+    @abstractmethod
+    def get_vector(
+        self,
+        collection_name: str,
+        vector_id: str,
+        include_vector: bool = True
+    ) -> Optional[VectorDocument]:
+        pass
+
+    @abstractmethod
+    def update_vector(
+        self,
+        collection_name: str,
+        vector_id: str,
+        vector: Optional[np.ndarray] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        payload: Optional[Dict[str, Any]] = None
+    ) -> None:
+        pass
+
+    @abstractmethod
+    def delete_vectors(
+        self,
+        collection_name: str,
+        vector_ids: List[str]
+    ) -> None:
+        pass
+
+    @abstractmethod
+    def count_vectors(self, collection_name: str) -> int:
+        pass
+
+    @abstractmethod
+    def close(self) -> None:
+        pass