apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-16 16:51:02 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

Added Pinecone for GE write & query

Browse source
This commit is contained in:
Cyber MacGeddon 2024-11-22 19:44:09 +00:00
parent ae1264f5c4
commit 85181bda85
9 changed files with 362 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

6
trustgraph-flow/scripts/ge-query-pinecone Executable file
View file

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

6
trustgraph-flow/scripts/ge-write-pinecone Executable file
View file

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

3
trustgraph-flow/setup.py
View file

@ -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",

3
trustgraph-flow/trustgraph/query/graph_embeddings/pinecone/__init__.py Normal file
View file

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

7
trustgraph-flow/trustgraph/query/graph_embeddings/pinecone/__main__.py Executable file
View file

@ -0,0 +1,7 @@
#!/usr/bin/env python3
from . hf import run
if __name__ == '__main__':
run()

153
trustgraph-flow/trustgraph/query/graph_embeddings/pinecone/service.py Executable file
View file

@ -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__)

3
trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/__init__.py Normal file
View file

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

7
trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/__main__.py Executable file
View file

@ -0,0 +1,7 @@
#!/usr/bin/env python3
from . write import run
if __name__ == '__main__':
run()

174
trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py Executable file
View file

@ -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__)