diff --git a/trustgraph-flow/scripts/ge-query-pinecone b/trustgraph-flow/scripts/ge-query-pinecone new file mode 100755 index 00000000..b75aec78 --- /dev/null +++ b/trustgraph-flow/scripts/ge-query-pinecone @@ -0,0 +1,6 @@ +#!/usr/bin/env python3 + +from trustgraph.query.graph_embeddings.pinecone import run + +run() + diff --git a/trustgraph-flow/scripts/ge-write-pinecone b/trustgraph-flow/scripts/ge-write-pinecone new file mode 100755 index 00000000..802a8377 --- /dev/null +++ b/trustgraph-flow/scripts/ge-write-pinecone @@ -0,0 +1,6 @@ +#!/usr/bin/env python3 + +from trustgraph.storage.graph_embeddings.pinecone import run + +run() + diff --git a/trustgraph-flow/setup.py b/trustgraph-flow/setup.py index 44901119..1650122f 100644 --- a/trustgraph-flow/setup.py +++ b/trustgraph-flow/setup.py @@ -60,6 +60,7 @@ setuptools.setup( "jsonschema", "aiohttp", "aiopulsar-py", + "pinecone[grpc]", ], scripts=[ "scripts/api-gateway", @@ -74,8 +75,10 @@ setuptools.setup( "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-rag", "scripts/kg-extract-definitions", diff --git a/trustgraph-flow/trustgraph/query/graph_embeddings/pinecone/__init__.py b/trustgraph-flow/trustgraph/query/graph_embeddings/pinecone/__init__.py new file mode 100644 index 00000000..ba844705 --- /dev/null +++ b/trustgraph-flow/trustgraph/query/graph_embeddings/pinecone/__init__.py @@ -0,0 +1,3 @@ + +from . service import * + diff --git a/trustgraph-flow/trustgraph/query/graph_embeddings/pinecone/__main__.py b/trustgraph-flow/trustgraph/query/graph_embeddings/pinecone/__main__.py new file mode 100755 index 00000000..89684e3e --- /dev/null +++ b/trustgraph-flow/trustgraph/query/graph_embeddings/pinecone/__main__.py @@ -0,0 +1,7 @@ +#!/usr/bin/env python3 + +from . hf import run + +if __name__ == '__main__': + run() + diff --git a/trustgraph-flow/trustgraph/query/graph_embeddings/pinecone/service.py b/trustgraph-flow/trustgraph/query/graph_embeddings/pinecone/service.py new file mode 100755 index 00000000..be855608 --- /dev/null +++ b/trustgraph-flow/trustgraph/query/graph_embeddings/pinecone/service.py @@ -0,0 +1,153 @@ + +""" +Graph embeddings query service. Input is vector, output is list of +entities. Pinecone implementation. +""" + +from pinecone import Pinecone, ServerlessSpec +from pinecone.grpc import PineconeGRPC, GRPCClientConfig +import uuid + +from .... schema import GraphEmbeddingsRequest, GraphEmbeddingsResponse +from .... schema import Error, Value +from .... schema import graph_embeddings_request_queue +from .... schema import graph_embeddings_response_queue +from .... base import ConsumerProducer + +module = ".".join(__name__.split(".")[1:-1]) + +default_input_queue = graph_embeddings_request_queue +default_output_queue = graph_embeddings_response_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) + + self.url = params.get("url", None) + self.api_key = params.get("api_key", None) + + 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, + "output_queue": output_queue, + "subscriber": subscriber, + "input_schema": GraphEmbeddingsRequest, + "output_schema": GraphEmbeddingsResponse, + "url": self.url, + } + ) + + def create_value(self, ent): + if ent.startswith("http://") or ent.startswith("https://"): + return Value(value=ent, is_uri=True) + else: + return Value(value=ent, is_uri=False) + + def handle(self, msg): + + try: + + v = msg.value() + + # Sender-produced ID + id = msg.properties()["id"] + + print(f"Handling input {id}...", flush=True) + + entities = set() + + for vec in v.vectors: + + dim = len(vec) + + index_name = ( + "t-" + v.user + "-" + str(dim) + ) + + index = self.pinecone.Index(index_name) + + results = index.query( + namespace=v.collection, + vector=vec, + top_k=v.limit, + include_values=False, + include_metadata=True + ) + + for r in results.matches: + ent = r.metadata["entity"] + entities.add(ent) + + # Convert set to list + entities = list(entities) + + ents2 = [] + + for ent in entities: + ents2.append(self.create_value(ent)) + + entities = ents2 + + print("Send response...", flush=True) + r = GraphEmbeddingsResponse(entities=entities, error=None) + self.producer.send(r, properties={"id": id}) + + print("Done.", flush=True) + + except Exception as e: + + print(f"Exception: {e}") + + print("Send error response...", flush=True) + + r = GraphEmbeddingsResponse( + error=Error( + type = "llm-error", + message = str(e), + ), + entities=None, + ) + + self.producer.send(r, properties={"id": id}) + + self.consumer.acknowledge(msg) + + @staticmethod + def add_args(parser): + + ConsumerProducer.add_args( + parser, default_input_queue, default_subscriber, + default_output_queue, + ) + + parser.add_argument( + '-a', '--api-key', + required=True, + help=f'Pinecone API key.' + ) + + parser.add_argument( + '-u', '--url', + help=f'Pinecone URL. If unspecified, serverless is used' + ) + +def run(): + + Processor.start(module, __doc__) + diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/__init__.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/__init__.py new file mode 100644 index 00000000..d891d55f --- /dev/null +++ b/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/__init__.py @@ -0,0 +1,3 @@ + +from . write import * + diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/__main__.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/__main__.py new file mode 100755 index 00000000..c05d8c6d --- /dev/null +++ b/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/__main__.py @@ -0,0 +1,7 @@ +#!/usr/bin/env python3 + +from . write import run + +if __name__ == '__main__': + run() + diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py new file mode 100755 index 00000000..fb974d48 --- /dev/null +++ b/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py @@ -0,0 +1,174 @@ + +""" +Accepts entity/vector pairs and writes them to a Pinecone store. +""" + +from pinecone import Pinecone, ServerlessSpec +from pinecone.grpc import PineconeGRPC, GRPCClientConfig + +import time +import uuid + +from .... schema import GraphEmbeddings +from .... schema import graph_embeddings_store_queue +from .... log_level import LogLevel +from .... base import Consumer + +module = ".".join(__name__.split(".")[1:-1]) + +default_input_queue = graph_embeddings_store_queue +default_subscriber = module +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", None) + + 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": GraphEmbeddings, + "url": self.url, + } + ) + + self.last_index_name = None + + def handle(self, msg): + + v = msg.value() + + # Make sure have an ID + # if v.metadata: + # if v.metadata.id: + # id = v.metadata.id + # else: + # id = str(uuid.uuid4()) + # else: + # str(uuid.uuid4()) + + id = str(uuid.uuid4()) + + if v.entity.value == "" or v.entity.value is None: return + + for vec in v.vectors: + + dim = len(vec) + + index_name = ( + "t-" + 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": { "entity": v.entity.value }, + } + ] + + 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', + required=True, + help=f'Pinecone API key.' + ) + + parser.add_argument( + '-u', '--url', + help=f'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__) +