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

Add templates

Browse source
This commit is contained in:
Cyber MacGeddon 2024-11-22 20:20:18 +00:00
parent 85181bda85
commit f3604acc97
4 changed files with 167 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -7,6 +7,7 @@ entities. Pinecone implementation.
from pinecone import Pinecone, ServerlessSpec
from pinecone.grpc import PineconeGRPC, GRPCClientConfig
import uuid
import os
from .... schema import GraphEmbeddingsRequest, GraphEmbeddingsResponse
from .... schema import Error, Value
@ -19,6 +20,7 @@ module = ".".join(__name__.split(".")[1:-1])
default_input_queue = graph_embeddings_request_queue
default_output_queue = graph_embeddings_response_queue
default_subscriber = module
default_api_key = os.getenv("PINECONE_API_KEY", "not-specified")
class Processor(ConsumerProducer):
@ -29,7 +31,7 @@ class Processor(ConsumerProducer):
subscriber = params.get("subscriber", default_subscriber)
self.url = params.get("url", None)
self.api_key = params.get("api_key", None)
self.api_key = params.get("api_key", default_api_key)
if self.url:
@ -138,13 +140,13 @@ class Processor(ConsumerProducer):
parser.add_argument(
'-a', '--api-key',
required=True,
help=f'Pinecone API key.'
default=default_api_key,
help='Pinecone API key. (default from PINECONE_API_KEY)'
)
parser.add_argument(
'-u', '--url',
help=f'Pinecone URL. If unspecified, serverless is used'
help='Pinecone URL. If unspecified, serverless is used'
)
def run():

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

@ -8,6 +8,7 @@ from pinecone.grpc import PineconeGRPC, GRPCClientConfig
import time
import uuid
import os
from .... schema import GraphEmbeddings
from .... schema import graph_embeddings_store_queue
@ -18,6 +19,7 @@ module = ".".join(__name__.split(".")[1:-1])
default_input_queue = graph_embeddings_store_queue
default_subscriber = module
default_api_key = os.getenv("PINECONE_API_KEY", "not-specified")
default_cloud = "aws"
default_region = "us-east-1"
@ -31,7 +33,7 @@ class Processor(Consumer):
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)
self.api_key = params.get("api_key", default_api_key)
if self.api_key is None:
raise RuntimeError("Pinecone API key must be specified")
@ -147,13 +149,13 @@ class Processor(Consumer):
parser.add_argument(
'-a', '--api-key',
required=True,
help=f'Pinecone API key.'
default=default_api_key,
help='Pinecone API key. (default from PINECONE_API_KEY)'
)
parser.add_argument(
'-u', '--url',
help=f'Pinecone URL. If unspecified, serverless is used'
help='Pinecone URL. If unspecified, serverless is used'
)
parser.add_argument(