Fix doc RAG processing, put graph-rag and doc-rag in appropriate component

files.

templates/components/document-rag.jsonnet

@@ -39,5 +39,35 @@ local prompts = import "prompts/mixtral.jsonnet";
 
     },
 
+    "document-embeddings" +: {
+    
+        create:: function(engine)
+
+            local container =
+                engine.container("document-embeddings")
+                    .with_image(images.trustgraph)
+                    .with_command([
+                        "document-embeddings",
+                        "-p",
+                        url.pulsar,
+                    ])
+                    .with_limits("1.0", "512M")
+                    .with_reservations("0.5", "512M");
+
+            local containerSet = engine.containers(
+                "document-embeddings", [ container ]
+            );
+
+            local service =
+                engine.internalService(containerSet)
+                .with_port(8000, 8000, "metrics");
+
+            engine.resources([
+                containerSet,
+                service,
+            ])
+
+    },
+
 }
 

templates/components/graph-rag.jsonnet

@@ -138,5 +138,35 @@ local url = import "values/url.jsonnet";
 
     },
 
+    "graph-embeddings" +: {
+    
+        create:: function(engine)
+
+            local container =
+                engine.container("graph-embeddings")
+                    .with_image(images.trustgraph)
+                    .with_command([
+                        "graph-embeddings",
+                        "-p",
+                        url.pulsar,
+                    ])
+                    .with_limits("1.0", "512M")
+                    .with_reservations("0.5", "512M");
+
+            local containerSet = engine.containers(
+                "graph-embeddings", [ container ]
+            );
+
+            local service =
+                engine.internalService(containerSet)
+                .with_port(8000, 8000, "metrics");
+
+            engine.resources([
+                containerSet,
+                service,
+            ])
+
+    },
+
 }
 

templates/components/trustgraph.jsonnet

@@ -119,66 +119,6 @@ local prompt = import "prompt-template.jsonnet";
 
     },
 
-    "graph-embeddings" +: {
-    
-        create:: function(engine)
-
-            local container =
-                engine.container("graph-embeddings")
-                    .with_image(images.trustgraph)
-                    .with_command([
-                        "graph-embeddings",
-                        "-p",
-                        url.pulsar,
-                    ])
-                    .with_limits("1.0", "512M")
-                    .with_reservations("0.5", "512M");
-
-            local containerSet = engine.containers(
-                "graph-embeddings", [ container ]
-            );
-
-            local service =
-                engine.internalService(containerSet)
-                .with_port(8000, 8000, "metrics");
-
-            engine.resources([
-                containerSet,
-                service,
-            ])
-
-    },
-
-    "document-embeddings" +: {
-    
-        create:: function(engine)
-
-            local container =
-                engine.container("document-embeddings")
-                    .with_image(images.trustgraph)
-                    .with_command([
-                        "document-embeddings",
-                        "-p",
-                        url.pulsar,
-                    ])
-                    .with_limits("1.0", "512M")
-                    .with_reservations("0.5", "512M");
-
-            local containerSet = engine.containers(
-                "document-embeddings", [ container ]
-            );
-
-            local service =
-                engine.internalService(containerSet)
-                .with_port(8000, 8000, "metrics");
-
-            engine.resources([
-                containerSet,
-                service,
-            ])
-
-    },
-
     "metering" +: {
     
         create:: function(engine)

trustgraph-base/trustgraph/api/api.py

@@ -131,6 +131,35 @@ class Api:
         except:
             raise ProtocolException(f"Response not formatted correctly")
 
+    def document_rag(self, question):
+
+        # The input consists of a question
+        input = {
+            "query": question
+        }
+
+        url = f"{self.url}document-rag"
+
+        # Invoke the API, input is passed as JSON
+        resp = requests.post(url, json=input)
+
+        # Should be a 200 status code
+        if resp.status_code != 200:
+            raise ProtocolException(f"Status code {resp.status_code}")
+
+        try:
+            # Parse the response as JSON
+            object = resp.json()
+        except:
+            raise ProtocolException(f"Expected JSON response")
+
+        self.check_error(resp)
+
+        try:
+            return object["response"]
+        except:
+            raise ProtocolException(f"Response not formatted correctly")
+
     def embeddings(self, text):
 
         # The input consists of a text block

trustgraph-base/trustgraph/clients/document_embeddings_client.py

@@ -38,8 +38,12 @@ class DocumentEmbeddingsClient(BaseClient):
             output_schema=DocumentEmbeddingsResponse,
         )
 
-    def request(self, vectors, limit=10, timeout=300):
+    def request(
+            self, vectors, user="trustgraph", collection="default",
+            limit=10, timeout=300
+    ):
         return self.call(
+            user=user, collection=collection,
             vectors=vectors, limit=limit, timeout=timeout
         ).documents
 

trustgraph-base/trustgraph/schema/documents.py

@@ -55,6 +55,8 @@ document_embeddings_store_queue = topic('document-embeddings-store')
 class DocumentEmbeddingsRequest(Record):
     vectors = Array(Array(Double()))
     limit = Integer()
+    user = String()
+    collection = String()
 
 class DocumentEmbeddingsResponse(Record):
     error = Error()

trustgraph-flow/trustgraph/document_rag.py

@@ -16,6 +16,44 @@ from . schema import document_embeddings_response_queue
 LABEL="http://www.w3.org/2000/01/rdf-schema#label"
 DEFINITION="http://www.w3.org/2004/02/skos/core#definition"
 
+class Query:
+
+    def __init__(self, rag, user, collection, verbose):
+        self.rag = rag
+        self.user = user
+        self.collection = collection
+        self.verbose = verbose
+
+    def get_vector(self, query):
+
+        if self.verbose:
+            print("Compute embeddings...", flush=True)
+
+        qembeds = self.rag.embeddings.request(query)
+
+        if self.verbose:
+            print("Done.", flush=True)
+
+        return qembeds
+
+    def get_docs(self, query):
+
+        vectors = self.get_vector(query)
+
+        if self.verbose:
+            print("Get entities...", flush=True)
+
+        docs = self.rag.de_client.request(
+            vectors, limit=self.rag.doc_limit
+        )
+
+        if self.verbose:
+            print("Docs:", flush=True)
+            for doc in docs:
+                print(doc, flush=True)
+
+        return docs
+
 class DocumentRag:
 
     def __init__(
@@ -55,7 +93,7 @@ class DocumentRag:
             print("Initialising...", flush=True)
 
         # FIXME: Configurable
-        self.entity_limit = 20
+        self.doc_limit = 20
 
         self.de_client = DocumentEmbeddingsClient(
             pulsar_host=pulsar_host,
@@ -81,42 +119,16 @@ class DocumentRag:
         if self.verbose:
             print("Initialised", flush=True)
 
-    def get_vector(self, query):
+    def query(self, query, user="trustgraph", collection="default"):
-
-        if self.verbose:
-            print("Compute embeddings...", flush=True)
-
-        qembeds = self.embeddings.request(query)
-
-        if self.verbose:
-            print("Done.", flush=True)
-
-        return qembeds
-
-    def get_docs(self, query):
-
-        vectors = self.get_vector(query)
-
-        if self.verbose:
-            print("Get entities...", flush=True)
-
-        docs = self.de_client.request(
-            vectors, self.entity_limit
-        )
-
-        if self.verbose:
-            print("Docs:", flush=True)
-            for doc in docs:
-                print(doc, flush=True)
-
-        return docs
-
-    def query(self, query):
 
         if self.verbose:
             print("Construct prompt...", flush=True)
 
-        docs = self.get_docs(query)
+        q = Query(
+            rag=self, user=user, collection=collection, verbose=self.verbose
+        )
+
+        docs = q.get_docs(query)
 
         if self.verbose:
             print("Invoke LLM...", flush=True)

trustgraph-flow/trustgraph/gateway/service.py

@@ -31,6 +31,7 @@ from . subscriber import Subscriber
 from . text_completion import TextCompletionRequestor
 from . prompt import PromptRequestor
 from . graph_rag import GraphRagRequestor
+from . document_rag import DocumentRagRequestor
 from . triples_query import TriplesQueryRequestor
 from . graph_embeddings_query import GraphEmbeddingsQueryRequestor
 from . embeddings import EmbeddingsRequestor
@@ -91,6 +92,10 @@ class Api:
                 pulsar_host=self.pulsar_host, timeout=self.timeout,
                 auth = self.auth,
             ),
+            "document-rag": DocumentRagRequestor(
+                pulsar_host=self.pulsar_host, timeout=self.timeout,
+                auth = self.auth,
+            ),
             "triples-query": TriplesQueryRequestor(
                 pulsar_host=self.pulsar_host, timeout=self.timeout,
                 auth = self.auth,
@@ -140,6 +145,10 @@ class Api:
                 endpoint_path = "/api/v1/graph-rag", auth=self.auth,
                 requestor = self.services["graph-rag"],
             ),
+            ServiceEndpoint(
+                endpoint_path = "/api/v1/document-rag", auth=self.auth,
+                requestor = self.services["document-rag"],
+            ),
             ServiceEndpoint(
                 endpoint_path = "/api/v1/triples-query", auth=self.auth,
                 requestor = self.services["triples-query"],

trustgraph-flow/trustgraph/storage/doc_embeddings/milvus/write.py

@@ -39,11 +39,16 @@ class Processor(Consumer):
 
         v = msg.value()
 
-        chunk = v.chunk.decode("utf-8")
+        for emb in v.chunks:
 
-        if v.chunk != "" and v.chunk is not None:
+            chunk = emb.chunk.decode("utf-8")
-            for vec in v.vectors:
+            if chunk == "" or chunk is None: continue
-                self.vecstore.insert(vec, chunk)
+
+            for vec in emb.vectors:
+
+                if chunk != "" and v.chunk is not None:
+                    for vec in v.vectors:
+                        self.vecstore.insert(vec, chunk)
 
     @staticmethod
     def add_args(parser):

trustgraph-flow/trustgraph/storage/doc_embeddings/pinecone/write.py

@@ -65,71 +65,74 @@ class Processor(Consumer):
 
         v = msg.value()
 
-        chunk = v.chunk.decode("utf-8")
+        for emb in v.chunks:
 
-        if chunk == "": return
+            chunk = emb.chunk.decode("utf-8")
+            if chunk == "" or chunk is None: continue
 
-        for vec in v.vectors:
+            for vec in emb.vectors:
 
-            dim = len(vec)
+                for vec in v.vectors:
-            collection = (
-                "d-" + v.metadata.user + "-" + str(dim)
-            )
 
-            if index_name != self.last_index_name:
+                    dim = len(vec)
+                    collection = (
+                        "d-" + v.metadata.user + "-" + str(dim)
+                    )
 
-                if not self.pinecone.has_index(index_name):
+                    if index_name != self.last_index_name:
 
-                    try:
+                        if not self.pinecone.has_index(index_name):
 
-                        self.pinecone.create_index(
+                            try:
-                            name = index_name,
-                            dimension = dim,
-                            metric = "cosine",
-                            spec = ServerlessSpec(
-                                cloud = self.cloud,
-                                region = self.region,
-                            )
-                        )
 
-                        for i in range(0, 1000):
+                                self.pinecone.create_index(
+                                    name = index_name,
+                                    dimension = dim,
+                                    metric = "cosine",
+                                    spec = ServerlessSpec(
+                                        cloud = self.cloud,
+                                        region = self.region,
+                                    )
+                                )
 
-                            if self.pinecone.describe_index(
+                                for i in range(0, 1000):
-                                    index_name
-                            ).status["ready"]:
-                                break
 
-                            time.sleep(1)
+                                    if self.pinecone.describe_index(
+                                            index_name
+                                    ).status["ready"]:
+                                        break
 
-                        if not self.pinecone.describe_index(
+                                    time.sleep(1)
-                                index_name
-                        ).status["ready"]:
-                            raise RuntimeError(
-                                "Gave up waiting for index creation"
-                            )
 
-                    except Exception as e:
+                                if not self.pinecone.describe_index(
-                        print("Pinecone index creation failed")
+                                        index_name
-                        raise e
+                                ).status["ready"]:
+                                    raise RuntimeError(
+                                        "Gave up waiting for index creation"
+                                    )
 
-                    print(f"Index {index_name} created", flush=True)
+                            except Exception as e:
+                                print("Pinecone index creation failed")
+                                raise e
 
-                self.last_index_name = index_name
+                            print(f"Index {index_name} created", flush=True)
 
-            index = self.pinecone.Index(index_name)
+                        self.last_index_name = index_name
 
-            records = [
+                    index = self.pinecone.Index(index_name)
-                {
-                    "id": id,
-                    "values": vec,
-                    "metadata": { "doc": chunk },
-                }
-            ]
 
-            index.upsert(
+                    records = [
-                vectors = records,
+                        {
-                namespace = v.metadata.collection,
+                            "id": id,
-            )
+                            "values": vec,
+                            "metadata": { "doc": chunk },
+                        }
+                    ]
+
+                    index.upsert(
+                        vectors = records,
+                        namespace = v.metadata.collection,
+                    )
 
     @staticmethod
     def add_args(parser):

trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py

@@ -44,47 +44,48 @@ class Processor(Consumer):
 
         v = msg.value()
 
-        chunk = v.chunk.decode("utf-8")
+        for emb in v.chunks:
 
-        if chunk == "": return
+            chunk = emb.chunk.decode("utf-8")
+            if chunk == "": return
 
-        for vec in v.vectors:
+            for vec in emb.vectors:
 
-            dim = len(vec)
+                dim = len(vec)
-            collection = (
+                collection = (
-                "d_" + v.metadata.user + "_" + v.metadata.collection + "_" +
+                    "d_" + v.metadata.user + "_" + v.metadata.collection + "_" +
-                str(dim)
+                    str(dim)
-            )
+                )
 
-            if collection != self.last_collection:
+                if collection != self.last_collection:
 
-                if not self.client.collection_exists(collection):
+                    if not self.client.collection_exists(collection):
 
-                    try:
+                        try:
-                        self.client.create_collection(
+                            self.client.create_collection(
-                            collection_name=collection,
+                                collection_name=collection,
-                            vectors_config=VectorParams(
+                                vectors_config=VectorParams(
-                                size=dim, distance=Distance.COSINE
+                                    size=dim, distance=Distance.COSINE
-                            ),
+                                ),
+                            )
+                        except Exception as e:
+                            print("Qdrant collection creation failed")
+                            raise e
+
+                    self.last_collection = collection
+
+                self.client.upsert(
+                    collection_name=collection,
+                    points=[
+                        PointStruct(
+                            id=str(uuid.uuid4()),
+                            vector=vec,
+                            payload={
+                                "doc": chunk,
+                            }
                         )
-                    except Exception as e:
+                    ]
-                        print("Qdrant collection creation failed")
+                )
-                        raise e
-
-                self.last_collection = collection
-
-            self.client.upsert(
-                collection_name=collection,
-                points=[
-                    PointStruct(
-                        id=str(uuid.uuid4()),
-                        vector=vec,
-                        payload={
-                            "doc": chunk,
-                        }
-                    )
-                ]
-            )
 
     @staticmethod
     def add_args(parser):