trustgraph-flow 0.19.4__tar.gz → 0.19.6__tar.gz

{trustgraph-flow-0.19.4 → trustgraph-flow-0.19.6}/PKG-INFO

@@ -1,12 +1,12 @@
 Metadata-Version: 2.1
 Name: trustgraph-flow
-Version: 0.19.4
+Version: 0.19.6
 Summary: TrustGraph provides a means to run a pipeline of flexible AI processing components in a flexible means to achieve a processing pipeline.
 Home-page: https://github.com/trustgraph-ai/trustgraph
 Author: trustgraph.ai
 Author-email: security@trustgraph.ai
 License: UNKNOWN
-Download-URL: https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v0.19.4.tar.gz
+Download-URL: https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v0.19.6.tar.gz
 Platform: UNKNOWN
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)

trustgraph-flow-0.19.6/scripts/document-embeddings

@@ -0,0 +1,6 @@
+#!/usr/bin/env python3
+
+from trustgraph.embeddings.document_embeddings import run
+
+run()
+

trustgraph-flow-0.19.6/scripts/graph-embeddings

@@ -0,0 +1,6 @@
+#!/usr/bin/env python3
+
+from trustgraph.embeddings.graph_embeddings import run
+
+run()
+

{trustgraph-flow-0.19.4 → trustgraph-flow-0.19.6}/setup.py

@@ -63,29 +63,30 @@ setuptools.setup(
         "falkordb",
     ],
     scripts=[
-        "scripts/api-gateway",
         "scripts/agent-manager-react",
+        "scripts/api-gateway",
         "scripts/chunker-recursive",
         "scripts/chunker-token",
         "scripts/de-query-milvus",
-        "scripts/de-query-qdrant",
         "scripts/de-query-pinecone",
+        "scripts/de-query-qdrant",
         "scripts/de-write-milvus",
-        "scripts/de-write-qdrant",
         "scripts/de-write-pinecone",
+        "scripts/de-write-qdrant",
+        "scripts/document-embeddings",
         "scripts/document-rag",
         "scripts/embeddings-ollama",
-        "scripts/embeddings-vectorize",
         "scripts/ge-query-milvus",
         "scripts/ge-query-pinecone",
         "scripts/ge-query-qdrant",
         "scripts/ge-write-milvus",
         "scripts/ge-write-pinecone",
         "scripts/ge-write-qdrant",
+        "scripts/graph-embeddings",
         "scripts/graph-rag",
         "scripts/kg-extract-definitions",
-        "scripts/kg-extract-topics",
         "scripts/kg-extract-relationships",
+        "scripts/kg-extract-topics",
         "scripts/metering",
         "scripts/object-extract-row",
         "scripts/oe-write-milvus",
@@ -103,13 +104,13 @@ setuptools.setup(
         "scripts/text-completion-ollama",
         "scripts/text-completion-openai",
         "scripts/triples-query-cassandra",
-        "scripts/triples-query-neo4j",
-        "scripts/triples-query-memgraph",
         "scripts/triples-query-falkordb",
+        "scripts/triples-query-memgraph",
+        "scripts/triples-query-neo4j",
         "scripts/triples-write-cassandra",
-        "scripts/triples-write-neo4j",
-        "scripts/triples-write-memgraph",
         "scripts/triples-write-falkordb",
+        "scripts/triples-write-memgraph",
+        "scripts/triples-write-neo4j",
         "scripts/wikipedia-lookup",
     ]
 )

{trustgraph-flow-0.19.4 → trustgraph-flow-0.19.6}/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):
-
-        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)
+    def query(self, query, user="trustgraph", collection="default"):
 
         if self.verbose:
-            print("Get entities...", flush=True)
+            print("Construct prompt...", flush=True)
 
-        docs = self.de_client.request(
-            vectors, self.entity_limit
+        q = Query(
+            rag=self, user=user, collection=collection, verbose=self.verbose
         )
 
-        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)
+        docs = q.get_docs(query)
 
         if self.verbose:
             print("Invoke LLM...", flush=True)

trustgraph-flow-0.19.6/trustgraph/embeddings/document_embeddings/__init__.py

@@ -0,0 +1,3 @@
+
+from . embeddings import *
+

{trustgraph-flow-0.19.4/trustgraph/embeddings/vectorize → trustgraph-flow-0.19.6/trustgraph/embeddings/document_embeddings}/__main__.py

@@ -1,5 +1,5 @@
 
-from . vectorize import run
+from . embeddings import run
 
 if __name__ == '__main__':
     run()

trustgraph-flow-0.19.6/trustgraph/embeddings/document_embeddings/embeddings.py

@@ -0,0 +1,109 @@
+
+"""
+Document embeddings, calls the embeddings service to get embeddings for a
+chunk of text.  Input is chunk of text plus metadata.
+Output is chunk plus embedding.
+"""
+
+from ... schema import Chunk, ChunkEmbeddings, DocumentEmbeddings
+from ... schema import chunk_ingest_queue
+from ... schema import document_embeddings_store_queue
+from ... schema import embeddings_request_queue, embeddings_response_queue
+from ... clients.embeddings_client import EmbeddingsClient
+from ... log_level import LogLevel
+from ... base import ConsumerProducer
+
+module = ".".join(__name__.split(".")[1:-1])
+
+default_input_queue = chunk_ingest_queue
+default_output_queue = document_embeddings_store_queue
+default_subscriber = module
+
+class Processor(ConsumerProducer):
+
+    def __init__(self, **params):
+
+        input_queue = params.get("input_queue", default_input_queue)
+        output_queue = params.get("output_queue", default_output_queue)
+        subscriber = params.get("subscriber", default_subscriber)
+        emb_request_queue = params.get(
+            "embeddings_request_queue", embeddings_request_queue
+        )
+        emb_response_queue = params.get(
+            "embeddings_response_queue", embeddings_response_queue
+        )
+
+        super(Processor, self).__init__(
+            **params | {
+                "input_queue": input_queue,
+                "output_queue": output_queue,
+                "embeddings_request_queue": emb_request_queue,
+                "embeddings_response_queue": emb_response_queue,
+                "subscriber": subscriber,
+                "input_schema": Chunk,
+                "output_schema": DocumentEmbeddings,
+            }
+        )
+
+        self.embeddings = EmbeddingsClient(
+            pulsar_host=self.pulsar_host,
+            input_queue=emb_request_queue,
+            output_queue=emb_response_queue,
+            subscriber=module + "-emb",
+        )
+
+    def handle(self, msg):
+
+        v = msg.value()
+        print(f"Indexing {v.metadata.id}...", flush=True)
+
+        try:
+
+            vectors = self.embeddings.request(v.chunk)
+
+            embeds = [
+                ChunkEmbeddings(
+                    chunk=v.chunk,
+                    vectors=vectors,
+                )
+            ]
+
+            r = DocumentEmbeddings(
+                metadata=v.metadata,
+                chunks=embeds,
+            )
+
+            self.producer.send(r)
+
+        except Exception as e:
+            print("Exception:", e, flush=True)
+
+            # Retry
+            raise e
+
+        print("Done.", flush=True)
+
+    @staticmethod
+    def add_args(parser):
+
+        ConsumerProducer.add_args(
+            parser, default_input_queue, default_subscriber,
+            default_output_queue,
+        )
+
+        parser.add_argument(
+            '--embeddings-request-queue',
+            default=embeddings_request_queue,
+            help=f'Embeddings request queue (default: {embeddings_request_queue})',
+        )
+
+        parser.add_argument(
+            '--embeddings-response-queue',
+            default=embeddings_response_queue,
+            help=f'Embeddings request queue (default: {embeddings_response_queue})',
+        )
+
+def run():
+
+    Processor.start(module, __doc__)
+

trustgraph-flow-0.19.6/trustgraph/embeddings/graph_embeddings/__init__.py

@@ -0,0 +1,3 @@
+
+from . embeddings import *
+

trustgraph-flow-0.19.6/trustgraph/embeddings/graph_embeddings/__main__.py

@@ -0,0 +1,6 @@
+
+from . embeddings import run
+
+if __name__ == '__main__':
+    run()
+

trustgraph-flow-0.19.4/trustgraph/embeddings/vectorize/vectorize.py → trustgraph-flow-0.19.6/trustgraph/embeddings/graph_embeddings/embeddings.py

@@ -1,7 +1,8 @@
 
 """
-Vectorizer, calls the embeddings service to get embeddings for a chunk.
-Input is text chunk, output is chunk and vectors.
+Graph embeddings, calls the embeddings service to get embeddings for a
+set of entity contexts.  Input is entity plus textual context.
+Output is entity plus embedding.
 """
 
 from ... schema import EntityContexts, EntityEmbeddings, GraphEmbeddings
@@ -51,11 +52,6 @@ class Processor(ConsumerProducer):
             subscriber=module + "-emb",
         )
 
-    def emit(self, rec, vectors):
-
-        r = GraphEmbeddings(metadata=metadata, chunk=chunk, vectors=vectors)
-        self.producer.send(r)
-
     def handle(self, msg):
 
         v = msg.value()

trustgraph-flow-0.19.6/trustgraph/flow_version.py

@@ -0,0 +1 @@
+__version__ = "0.19.6"

trustgraph-flow-0.19.6/trustgraph/gateway/document_rag.py

@@ -0,0 +1,30 @@
+
+from .. schema import DocumentRagQuery, DocumentRagResponse
+from .. schema import document_rag_request_queue
+from .. schema import document_rag_response_queue
+
+from . endpoint import ServiceEndpoint
+from . requestor import ServiceRequestor
+
+class DocumentRagRequestor(ServiceRequestor):
+    def __init__(self, pulsar_host, timeout, auth):
+
+        super(DocumentRagRequestor, self).__init__(
+            pulsar_host=pulsar_host,
+            request_queue=document_rag_request_queue,
+            response_queue=document_rag_response_queue,
+            request_schema=DocumentRagQuery,
+            response_schema=DocumentRagResponse,
+            timeout=timeout,
+        )
+
+    def to_request(self, body):
+        return DocumentRagQuery(
+            query=body["query"],
+            user=body.get("user", "trustgraph"),
+            collection=body.get("collection", "default"),
+        )
+
+    def from_response(self, message):
+        return { "response": message.response }, True
+

{trustgraph-flow-0.19.4 → trustgraph-flow-0.19.6}/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-0.19.4 → trustgraph-flow-0.19.6}/trustgraph/storage/doc_embeddings/milvus/write.py

@@ -3,15 +3,16 @@
 Accepts entity/vector pairs and writes them to a Milvus store.
 """
 
-from .... schema import ChunkEmbeddings
-from .... schema import chunk_embeddings_ingest_queue
-from .... log_level import LogLevel
 from .... direct.milvus_doc_embeddings import DocVectors
+
+from .... schema import DocumentEmbeddings
+from .... schema import document_embeddings_store_queue
+from .... log_level import LogLevel
 from .... base import Consumer
 
 module = ".".join(__name__.split(".")[1:-1])
 
-default_input_queue = chunk_embeddings_ingest_queue
+default_input_queue = document_embeddings_store_queue
 default_subscriber = module
 default_store_uri = 'http://localhost:19530'
 
@@ -27,7 +28,7 @@ class Processor(Consumer):
             **params | {
                 "input_queue": input_queue,
                 "subscriber": subscriber,
-                "input_schema": ChunkEmbeddings,
+                "input_schema": DocumentEmbeddings,
                 "store_uri": store_uri,
             }
         )
@@ -38,11 +39,16 @@ class Processor(Consumer):
 
         v = msg.value()
 
-        chunk = v.chunk.decode("utf-8")
+        for emb in v.chunks:
+
+            chunk = emb.chunk.decode("utf-8")
+            if chunk == "" or chunk is None: continue
+
+            for vec in emb.vectors:
 
-        if v.chunk != "" and v.chunk is not None:
-            for vec in v.vectors:
-                self.vecstore.insert(vec, chunk)
+                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-0.19.6/trustgraph/storage/doc_embeddings/pinecone/write.py

@@ -0,0 +1,170 @@
+
+"""
+Accepts entity/vector pairs and writes them to a Qdrant store.
+"""
+
+from qdrant_client import QdrantClient
+from qdrant_client.models import PointStruct
+from qdrant_client.models import Distance, VectorParams
+
+import time
+import uuid
+import os
+
+from .... schema import DocumentEmbeddings
+from .... schema import document_embeddings_store_queue
+from .... log_level import LogLevel
+from .... base import Consumer
+
+module = ".".join(__name__.split(".")[1:-1])
+
+default_input_queue = document_embeddings_store_queue
+default_subscriber = module
+default_api_key = os.getenv("PINECONE_API_KEY", "not-specified")
+default_cloud = "aws"
+default_region = "us-east-1"
+
+class Processor(Consumer):
+
+    def __init__(self, **params):
+
+        input_queue = params.get("input_queue", default_input_queue)
+        subscriber = params.get("subscriber", default_subscriber)
+
+        self.url = params.get("url", None)
+        self.cloud = params.get("cloud", default_cloud)
+        self.region = params.get("region", default_region)
+        self.api_key = params.get("api_key", default_api_key)
+
+        if self.api_key is None:
+            raise RuntimeError("Pinecone API key must be specified")
+
+        if self.url:
+
+            self.pinecone = PineconeGRPC(
+                api_key = self.api_key,
+                host = self.url
+            )
+
+        else:
+
+            self.pinecone = Pinecone(api_key = self.api_key)
+
+        super(Processor, self).__init__(
+            **params | {
+                "input_queue": input_queue,
+                "subscriber": subscriber,
+                "input_schema": DocumentEmbeddings,
+                "url": self.url,
+            }
+        )
+
+        self.last_index_name = None
+
+    def handle(self, msg):
+
+        v = msg.value()
+
+        for emb in v.chunks:
+
+            chunk = emb.chunk.decode("utf-8")
+            if chunk == "" or chunk is None: continue
+
+            for vec in emb.vectors:
+
+                for vec in v.vectors:
+
+                    dim = len(vec)
+                    collection = (
+                        "d-" + v.metadata.user + "-" + str(dim)
+                    )
+
+                    if index_name != self.last_index_name:
+
+                        if not self.pinecone.has_index(index_name):
+
+                            try:
+
+                                self.pinecone.create_index(
+                                    name = index_name,
+                                    dimension = dim,
+                                    metric = "cosine",
+                                    spec = ServerlessSpec(
+                                        cloud = self.cloud,
+                                        region = self.region,
+                                    )
+                                )
+
+                                for i in range(0, 1000):
+
+                                    if self.pinecone.describe_index(
+                                            index_name
+                                    ).status["ready"]:
+                                        break
+
+                                    time.sleep(1)
+
+                                if not self.pinecone.describe_index(
+                                        index_name
+                                ).status["ready"]:
+                                    raise RuntimeError(
+                                        "Gave up waiting for index creation"
+                                    )
+
+                            except Exception as e:
+                                print("Pinecone index creation failed")
+                                raise e
+
+                            print(f"Index {index_name} created", flush=True)
+
+                        self.last_index_name = index_name
+
+                    index = self.pinecone.Index(index_name)
+
+                    records = [
+                        {
+                            "id": id,
+                            "values": vec,
+                            "metadata": { "doc": chunk },
+                        }
+                    ]
+
+                    index.upsert(
+                        vectors = records,
+                        namespace = v.metadata.collection,
+                    )
+
+    @staticmethod
+    def add_args(parser):
+
+        Consumer.add_args(
+            parser, default_input_queue, default_subscriber,
+        )
+
+        parser.add_argument(
+            '-a', '--api-key',
+            default=default_api_key,
+            help='Pinecone API key. (default from PINECONE_API_KEY)'
+        )
+
+        parser.add_argument(
+            '-u', '--url',
+            help='Pinecone URL.  If unspecified, serverless is used'
+        )
+
+        parser.add_argument(
+            '--cloud',
+            default=default_cloud,
+            help=f'Pinecone cloud, (default: {default_cloud}'
+        )
+
+        parser.add_argument(
+            '--region',
+            default=default_region,
+            help=f'Pinecone region, (default: {default_region}'
+        )
+
+def run():
+
+    Processor.start(module, __doc__)
+