diff --git a/templates/components.jsonnet b/templates/components.jsonnet index ec7f862b..26368deb 100644 --- a/templates/components.jsonnet +++ b/templates/components.jsonnet @@ -25,6 +25,7 @@ "trustgraph-base": import "components/trustgraph.jsonnet", "vector-store-milvus": import "components/milvus.jsonnet", "vector-store-qdrant": import "components/qdrant.jsonnet", + "vector-store-pinecone": import "components/pinecone.jsonnet", "vertexai": import "components/vertexai.jsonnet", "null": {}, @@ -34,6 +35,7 @@ "cassandra": import "components/cassandra.jsonnet", "neo4j": import "components/neo4j.jsonnet", "qdrant": import "components/qdrant.jsonnet", + "pinecone": import "components/pinecone.jsonnet", "milvus": import "components/milvus.jsonnet", "trustgraph": import "components/trustgraph.jsonnet", diff --git a/templates/components/pinecone.jsonnet b/templates/components/pinecone.jsonnet new file mode 100644 index 00000000..3422952a --- /dev/null +++ b/templates/components/pinecone.jsonnet @@ -0,0 +1,153 @@ +local base = import "base/base.jsonnet"; +local images = import "values/images.jsonnet"; +local url = import "values/url.jsonnet"; +local cassandra_hosts = "cassandra"; + +{ + + "pinecone-cloud":: "aws", + "pinecone-region":: "us-east-1", + + "store-graph-embeddings" +: { + + create:: function(engine) + + local envSecrets = engine.envSecrets("pinecone-api-key") + .with_env_var("PINECONE_API_KEY", "pinecone-api-key"); + + local container = + engine.container("store-graph-embeddings") + .with_image(images.trustgraph) + .with_command([ + "ge-write-pinecone", + "-p", + url.pulsar, + ]) + .with_env_var_secrets(envSecrets) + .with_limits("0.5", "128M") + .with_reservations("0.1", "128M"); + + local containerSet = engine.containers( + "store-graph-embeddings", [ container ] + ); + + local service = + engine.internalService(containerSet) + .with_port(8080, 8080, "metrics"); + + engine.resources([ + envSecrets, + containerSet, + service, + ]) + + }, + + "query-graph-embeddings" +: { + + create:: function(engine) + + local envSecrets = engine.envSecrets("pinecone-api-key") + .with_env_var("PINECONE_API_KEY", "pinecone-api-key"); + + local container = + engine.container("query-graph-embeddings") + .with_image(images.trustgraph) + .with_command([ + "ge-query-pinecone", + "-p", + url.pulsar, + ]) + .with_env_var_secrets(envSecrets) + .with_limits("0.5", "128M") + .with_reservations("0.1", "128M"); + + local containerSet = engine.containers( + "query-graph-embeddings", [ container ] + ); + + local service = + engine.internalService(containerSet) + .with_port(8080, 8080, "metrics"); + + engine.resources([ + envSecrets, + containerSet, + service, + ]) + + }, + + "store-doc-embeddings" +: { + + create:: function(engine) + + local envSecrets = engine.envSecrets("pinecone-api-key") + .with_env_var("PINECONE_API_KEY", "pinecone-api-key"); + + local container = + engine.container("store-doc-embeddings") + .with_image(images.trustgraph) + .with_command([ + "de-write-pinecone", + "-p", + url.pulsar, + ]) + .with_env_var_secrets(envSecrets) + .with_limits("0.5", "128M") + .with_reservations("0.1", "128M"); + + local containerSet = engine.containers( + "store-doc-embeddings", [ container ] + ); + + local service = + engine.internalService(containerSet) + .with_port(8080, 8080, "metrics"); + + engine.resources([ + envSecrets, + containerSet, + service, + ]) + + }, + + "query-doc-embeddings" +: { + + create:: function(engine) + + local envSecrets = engine.envSecrets("pinecone-api-key") + .with_env_var("PINECONE_API_KEY", "pinecone-api-key"); + + local container = + engine.container("query-doc-embeddings") + .with_image(images.trustgraph) + .with_command([ + "de-query-pinecone", + "-p", + url.pulsar, + ]) + .with_env_var_secrets(envSecrets) + .with_limits("0.5", "128M") + .with_reservations("0.1", "128M"); + + local containerSet = engine.containers( + "query-doc-embeddings", [ container ] + ); + + local service = + engine.internalService(containerSet) + .with_port(8080, 8080, "metrics"); + + engine.resources([ + envSecrets, + containerSet, + service, + ]) + + + } + +} + diff --git a/trustgraph-flow/trustgraph/query/graph_embeddings/pinecone/service.py b/trustgraph-flow/trustgraph/query/graph_embeddings/pinecone/service.py index be855608..26acac80 100755 --- a/trustgraph-flow/trustgraph/query/graph_embeddings/pinecone/service.py +++ b/trustgraph-flow/trustgraph/query/graph_embeddings/pinecone/service.py @@ -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(): diff --git a/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py b/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py index fb974d48..980086e8 100755 --- a/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py +++ b/trustgraph-flow/trustgraph/storage/graph_embeddings/pinecone/write.py @@ -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(