diff --git a/Makefile b/Makefile index 348aa2aa..0d96919f 100644 --- a/Makefile +++ b/Makefile @@ -1,6 +1,6 @@ # VERSION=$(shell git describe | sed 's/^v//') -VERSION=0.5.5 +VERSION=0.5.6 all: container diff --git a/docker-compose-azure.yaml b/docker-compose-azure.yaml index 3dc3653b..0149cc4c 100644 --- a/docker-compose-azure.yaml +++ b/docker-compose-azure.yaml @@ -130,7 +130,7 @@ services: restart: on-failure:100 pdf-decoder: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "pdf-decoder" - "-p" @@ -138,7 +138,7 @@ services: restart: on-failure:100 chunker: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "chunker-recursive" - "-p" @@ -146,7 +146,7 @@ services: restart: on-failure:100 vectorize: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "embeddings-vectorize" - "-p" @@ -154,7 +154,7 @@ services: restart: on-failure:100 embeddings: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "embeddings-hf" - "-p" @@ -164,7 +164,7 @@ services: restart: on-failure:100 kg-extract-definitions: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "kg-extract-definitions" - "-p" @@ -172,7 +172,7 @@ services: restart: on-failure:100 kg-extract-relationships: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "kg-extract-relationships" - "-p" @@ -180,7 +180,7 @@ services: restart: on-failure:100 store-graph-embeddings: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "ge-write-milvus" - "-p" @@ -190,7 +190,7 @@ services: restart: on-failure:100 store-triples: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "triples-write-cassandra" - "-p" @@ -200,7 +200,7 @@ services: restart: on-failure:100 text-completion: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "text-completion-azure" - "-p" @@ -211,11 +211,31 @@ services: - ${AZURE_ENDPOINT} restart: on-failure:100 + text-completion-rag: + image: docker.io/trustgraph/trustgraph-flow:0.5.6 + command: + - "text-completion-azure" + - "-p" + - "pulsar://pulsar:6650" + - "-k" + - ${AZURE_TOKEN} + - "-e" + - ${AZURE_ENDPOINT} + - "-i" + - "non-persistent://tg/request/text-completion-rag" + - "-o" + - "non-persistent://tg/response/text-completion-rag-response" + restart: on-failure:100 + graph-rag: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "graph-rag" - "-p" - "pulsar://pulsar:6650" + - "--text-completion-request-queue" + - "non-persistent://tg/request/text-completion-rag" + - "--text-completion-response-queue" + - "non-persistent://tg/response/text-completion-rag-response" restart: on-failure:100 diff --git a/docker-compose-claude.yaml b/docker-compose-claude.yaml index 26e205ea..35b422db 100644 --- a/docker-compose-claude.yaml +++ b/docker-compose-claude.yaml @@ -130,7 +130,7 @@ services: restart: on-failure:100 pdf-decoder: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "pdf-decoder" - "-p" @@ -138,7 +138,7 @@ services: restart: on-failure:100 chunker: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "chunker-recursive" - "-p" @@ -146,7 +146,7 @@ services: restart: on-failure:100 vectorize: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "embeddings-vectorize" - "-p" @@ -154,7 +154,7 @@ services: restart: on-failure:100 embeddings: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "embeddings-hf" - "-p" @@ -164,7 +164,7 @@ services: restart: on-failure:100 kg-extract-definitions: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "kg-extract-definitions" - "-p" @@ -172,7 +172,7 @@ services: restart: on-failure:100 kg-extract-relationships: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "kg-extract-relationships" - "-p" @@ -180,7 +180,7 @@ services: restart: on-failure:100 store-graph-embeddings: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "ge-write-milvus" - "-p" @@ -190,7 +190,7 @@ services: restart: on-failure:100 store-triples: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "triples-write-cassandra" - "-p" @@ -200,7 +200,7 @@ services: restart: on-failure:100 text-completion: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "text-completion-claude" - "-p" @@ -209,11 +209,28 @@ services: - ${CLAUDE_KEY} restart: on-failure:100 + text-completion-rag: + image: docker.io/trustgraph/trustgraph-flow:0.5.6 + command: + - "text-completion-ollama" + - "-p" + - "pulsar://pulsar:6650" + - "-k" + - ${CLAUDE_KEY} + - "-i" + - "non-persistent://tg/request/text-completion-rag" + - "-o" + - "non-persistent://tg/response/text-completion-rag-response" + restart: on-failure:100 + graph-rag: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "graph-rag" - "-p" - "pulsar://pulsar:6650" + - "--text-completion-request-queue" + - "non-persistent://tg/request/text-completion-rag" + - "--text-completion-response-queue" + - "non-persistent://tg/response/text-completion-rag-response" restart: on-failure:100 - diff --git a/docker-compose-ollama.yaml b/docker-compose-ollama.yaml index a2ec6f1a..0ae7288a 100644 --- a/docker-compose-ollama.yaml +++ b/docker-compose-ollama.yaml @@ -130,7 +130,7 @@ services: restart: on-failure:100 pdf-decoder: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "pdf-decoder" - "-p" @@ -138,7 +138,7 @@ services: restart: on-failure:100 chunker: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "chunker-recursive" - "-p" @@ -146,7 +146,7 @@ services: restart: on-failure:100 vectorize: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "embeddings-vectorize" - "-p" @@ -154,7 +154,7 @@ services: restart: on-failure:100 embeddings: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "embeddings-hf" - "-p" @@ -164,7 +164,7 @@ services: restart: on-failure:100 kg-extract-definitions: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "kg-extract-definitions" - "-p" @@ -172,7 +172,7 @@ services: restart: on-failure:100 kg-extract-relationships: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "kg-extract-relationships" - "-p" @@ -180,7 +180,7 @@ services: restart: on-failure:100 store-graph-embeddings: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "ge-write-milvus" - "-p" @@ -190,7 +190,7 @@ services: restart: on-failure:100 store-triples: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "triples-write-cassandra" - "-p" @@ -200,7 +200,7 @@ services: restart: on-failure:100 text-completion: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "text-completion-ollama" - "-p" @@ -210,12 +210,32 @@ services: # - "-m" # - "llama3.1:8b" restart: on-failure:100 + + text-completion-rag: + image: docker.io/trustgraph/trustgraph-flow:0.5.6 + command: + - "text-completion-ollama" + - "-p" + - "pulsar://pulsar:6650" + - "-r" + - "http://${OLLAMA_HOST}:11434/" + - "-i" + - "non-persistent://tg/request/text-completion-rag" + - "-o" + - "non-persistent://tg/response/text-completion-rag-response" + - "-m" + - "llama3.1:8b" + restart: on-failure:100 graph-rag: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "graph-rag" - "-p" - "pulsar://pulsar:6650" + - "--text-completion-request-queue" + - "non-persistent://tg/request/text-completion-rag" + - "--text-completion-response-queue" + - "non-persistent://tg/response/text-completion-rag-response" restart: on-failure:100 diff --git a/docker-compose-vertexai.yaml b/docker-compose-vertexai.yaml index c55a04b5..4417cea7 100644 --- a/docker-compose-vertexai.yaml +++ b/docker-compose-vertexai.yaml @@ -130,7 +130,7 @@ services: restart: on-failure:100 pdf-decoder: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "pdf-decoder" - "-p" @@ -138,7 +138,7 @@ services: restart: on-failure:100 chunker: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "chunker-recursive" - "-p" @@ -146,7 +146,7 @@ services: restart: on-failure:100 vectorize: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "embeddings-vectorize" - "-p" @@ -154,7 +154,7 @@ services: restart: on-failure:100 embeddings: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "embeddings-hf" - "-p" @@ -164,7 +164,7 @@ services: restart: on-failure:100 kg-extract-definitions: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "kg-extract-definitions" - "-p" @@ -172,7 +172,7 @@ services: restart: on-failure:100 kg-extract-relationships: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "kg-extract-relationships" - "-p" @@ -180,7 +180,7 @@ services: restart: on-failure:100 store-graph-embeddings: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "ge-write-milvus" - "-p" @@ -190,7 +190,7 @@ services: restart: on-failure:100 store-triples: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "triples-write-cassandra" - "-p" @@ -200,9 +200,9 @@ services: restart: on-failure:100 text-completion: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - - "text-completion-ollama" + - "text-completion-vertexai" - "-p" - "pulsar://pulsar:6650" - "-k" @@ -212,12 +212,32 @@ services: volumes: - "./vertexai:/vertexai" restart: on-failure:100 + + text-completion-rag: + image: docker.io/trustgraph/trustgraph-flow:0.5.6 + command: + - "text-completion-vertexai" + - "-p" + - "pulsar://pulsar:6650" + - "-k" + - "/vertexai/private.json" + - "-r" + - "us-west1" + - "-i" + - "non-persistent://tg/request/text-completion-rag" + - "-o" + - "non-persistent://tg/response/text-completion-rag-response" + restart: on-failure:100 graph-rag: - image: docker.io/trustgraph/trustgraph-flow:0.5.5 + image: docker.io/trustgraph/trustgraph-flow:0.5.6 command: - "graph-rag" - "-p" - "pulsar://pulsar:6650" + - "--text-completion-request-queue" + - "non-persistent://tg/request/text-completion-rag" + - "--text-completion-response-queue" + - "non-persistent://tg/response/text-completion-rag-response" restart: on-failure:100 diff --git a/requirements.txt b/requirements.txt index 7bfcf3ba..a0c184ca 100644 --- a/requirements.txt +++ b/requirements.txt @@ -10,7 +10,12 @@ langchain-huggingface langchain-text-splitters langchain-community huggingface-hub +requests cassandra-driver pulsar-client +pypdf anthropic google-cloud-aiplatform +pyyaml +prometheus-client +pyarrow diff --git a/setup.py b/setup.py index 1fec809b..d0f4626e 100644 --- a/setup.py +++ b/setup.py @@ -4,7 +4,7 @@ import os with open("README.md", "r") as fh: long_description = fh.read() -version = "0.5.5" +version = "0.5.6" setuptools.setup( name="trustgraph", diff --git a/trustgraph/embeddings/vectorize/vectorize.py b/trustgraph/embeddings/vectorize/vectorize.py index bac0e932..1500b254 100755 --- a/trustgraph/embeddings/vectorize/vectorize.py +++ b/trustgraph/embeddings/vectorize/vectorize.py @@ -6,6 +6,7 @@ Input is text chunk, output is chunk and vectors. from ... schema import Chunk, ChunkEmbeddings from ... schema import chunk_ingest_queue, chunk_embeddings_ingest_queue +from ... schema import embeddings_request_queue, embeddings_response_queue from ... embeddings_client import EmbeddingsClient from ... log_level import LogLevel from ... base import ConsumerProducer @@ -23,11 +24,19 @@ class Processor(ConsumerProducer): 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": ChunkEmbeddings, @@ -36,6 +45,8 @@ class Processor(ConsumerProducer): self.embeddings = EmbeddingsClient( pulsar_host=self.pulsar_host, + input_queue=emb_request_queue, + output_queue=emb_response_queue, subscriber=module + "-emb", ) @@ -74,6 +85,18 @@ class Processor(ConsumerProducer): 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__) diff --git a/trustgraph/embeddings_client.py b/trustgraph/embeddings_client.py index 9b7fa81a..6daa2ff5 100644 --- a/trustgraph/embeddings_client.py +++ b/trustgraph/embeddings_client.py @@ -17,12 +17,21 @@ DEBUG=_pulsar.LoggerLevel.Debug class EmbeddingsClient: def __init__( - self, log_level=ERROR, subscriber=None, + self, log_level=ERROR, + input_queue=None, + output_queue=None, + subscriber=None, pulsar_host="pulsar://pulsar:6650", ): self.client = None + if input_queue == None: + input_queue=embeddings_request_queue + + if output_queue == None: + output_queue=embeddings_response_queue + if subscriber == None: subscriber = str(uuid.uuid4()) @@ -32,13 +41,13 @@ class EmbeddingsClient: ) self.producer = self.client.create_producer( - topic=embeddings_request_queue, + topic=input_queue, schema=JsonSchema(EmbeddingsRequest), chunking_enabled=True, ) self.consumer = self.client.subscribe( - embeddings_response_queue, subscriber, + output_queue, subscriber, schema=JsonSchema(EmbeddingsResponse), ) diff --git a/trustgraph/graph_rag.py b/trustgraph/graph_rag.py index 0b7ccbdc..967bd68a 100644 --- a/trustgraph/graph_rag.py +++ b/trustgraph/graph_rag.py @@ -4,6 +4,8 @@ from trustgraph.triple_vectors import TripleVectors from trustgraph.trustgraph import TrustGraph from trustgraph.llm_client import LlmClient from trustgraph.embeddings_client import EmbeddingsClient +from . schema import text_completion_request_queue +from . schema import text_completion_response_queue LABEL="http://www.w3.org/2000/01/rdf-schema#label" DEFINITION="http://www.w3.org/2004/02/skos/core#definition" @@ -15,6 +17,10 @@ class GraphRag: graph_hosts=None, pulsar_host="pulsar://pulsar:6650", vector_store="http://milvus:19530", + completion_request_queue=None, + completion_response_queue=None, + emb_request_queue=None, + emb_response_queue=None, verbose=False, entity_limit=50, triple_limit=30, @@ -24,6 +30,18 @@ class GraphRag: self.verbose=verbose + if completion_request_queue == None: + completion_request_queue = text_completion_request_queue + + if completion_response_queue == None: + completion_response_queue = text_completion_response_queue + + if emb_request_queue == None: + emb_request_queue = embeddings_request_queue + + if emb_response_queue == None: + emb_response_queue = embeddings_response_queue + if graph_hosts == None: graph_hosts = ["cassandra"] @@ -34,6 +52,8 @@ class GraphRag: self.embeddings = EmbeddingsClient( pulsar_host=pulsar_host, + input_queue=emb_request_queue, + output_queue=emb_response_queue, subscriber=module + "-emb", ) @@ -47,6 +67,8 @@ class GraphRag: self.llm = LlmClient( pulsar_host=pulsar_host, + input_queue=completion_request_queue, + output_queue=completion_response_queue, subscriber=module + "-llm", ) diff --git a/trustgraph/kg/extract_definitions/extract.py b/trustgraph/kg/extract_definitions/extract.py index 334fb32d..d090b919 100755 --- a/trustgraph/kg/extract_definitions/extract.py +++ b/trustgraph/kg/extract_definitions/extract.py @@ -9,6 +9,8 @@ import json from ... schema import ChunkEmbeddings, Triple, Source, Value from ... schema import chunk_embeddings_ingest_queue, triples_store_queue +from ... schema import text_completion_request_queue +from ... schema import text_completion_response_queue from ... log_level import LogLevel from ... llm_client import LlmClient from ... prompts import to_definitions @@ -30,6 +32,12 @@ class Processor(ConsumerProducer): input_queue = params.get("input_queue", default_input_queue) output_queue = params.get("output_queue", default_output_queue) subscriber = params.get("subscriber", default_subscriber) + tc_request_queue = params.get( + "text_completion_request_queue", text_completion_request_queue + ) + tc_response_queue = params.get( + "text_completion_response_queue", text_completion_response_queue + ) super(Processor, self).__init__( **params | { @@ -38,11 +46,15 @@ class Processor(ConsumerProducer): "subscriber": subscriber, "input_schema": ChunkEmbeddings, "output_schema": Triple, + "text_completion_request_queue": tc_request_queue, + "text_completion_response_queue": tc_response_queue, } ) self.llm = LlmClient( pulsar_host=self.pulsar_host, + input_queue=tc_request_queue, + output_queue=tc_response_queue, subscriber = module + "-llm", ) @@ -108,6 +120,18 @@ class Processor(ConsumerProducer): default_output_queue, ) + parser.add_argument( + '--text-completion-request-queue', + default=text_completion_request_queue, + help=f'Text completion request queue (default: {text_completion_request_queue})', + ) + + parser.add_argument( + '--text-completion-response-queue', + default=text_completion_response_queue, + help=f'Text completion response queue (default: {text_completion_response_queue})', + ) + def run(): Processor.start(module, __doc__) diff --git a/trustgraph/kg/extract_relationships/extract.py b/trustgraph/kg/extract_relationships/extract.py index b3e34549..e2c63c4f 100755 --- a/trustgraph/kg/extract_relationships/extract.py +++ b/trustgraph/kg/extract_relationships/extract.py @@ -11,7 +11,10 @@ import os from pulsar.schema import JsonSchema from ... schema import ChunkEmbeddings, Triple, GraphEmbeddings, Source, Value -from ... schema import chunk_embeddings_ingest_queue, triples_store_queue, graph_embeddings_store_queue +from ... schema import chunk_embeddings_ingest_queue, triples_store_queue +from ... schema import graph_embeddings_store_queue +from ... schema import text_completion_request_queue +from ... schema import text_completion_response_queue from ... log_level import LogLevel from ... llm_client import LlmClient from ... prompts import to_relationships @@ -35,6 +38,12 @@ class Processor(ConsumerProducer): output_queue = params.get("output_queue", default_output_queue) vector_queue = params.get("vector_queue", default_vector_queue) subscriber = params.get("subscriber", default_subscriber) + tc_request_queue = params.get( + "text_completion_request_queue", text_completion_request_queue + ) + tc_response_queue = params.get( + "text_completion_response_queue", text_completion_response_queue + ) super(Processor, self).__init__( **params | { @@ -43,6 +52,8 @@ class Processor(ConsumerProducer): "subscriber": subscriber, "input_schema": ChunkEmbeddings, "output_schema": Triple, + "text_completion_request_queue": tc_request_queue, + "text_completion_response_queue": tc_response_queue, } ) @@ -55,6 +66,8 @@ class Processor(ConsumerProducer): "input_queue": input_queue, "output_queue": output_queue, "vector_queue": vector_queue, + "text_completion_request_queue": tc_request_queue, + "text_completion_response_queue": tc_response_queue, "subscriber": subscriber, "input_schema": ChunkEmbeddings.__name__, "output_schema": Triple.__name__, @@ -63,6 +76,8 @@ class Processor(ConsumerProducer): self.llm = LlmClient( pulsar_host = self.pulsar_host, + input_queue=tc_request_queue, + output_queue=tc_response_queue, subscriber = module + "-llm", ) @@ -179,6 +194,18 @@ class Processor(ConsumerProducer): help=f'Vector output queue (default: {default_vector_queue})' ) + parser.add_argument( + '--text-completion-request-queue', + default=text_completion_request_queue, + help=f'Text completion request queue (default: {text_completion_request_queue})', + ) + + parser.add_argument( + '--text-completion-response-queue', + default=text_completion_response_queue, + help=f'Text completion response queue (default: {text_completion_response_queue})', + ) + def run(): Processor.start(module, __doc__) diff --git a/trustgraph/llm_client.py b/trustgraph/llm_client.py index 3ed9a29a..062fb323 100644 --- a/trustgraph/llm_client.py +++ b/trustgraph/llm_client.py @@ -19,10 +19,19 @@ DEBUG=_pulsar.LoggerLevel.Debug class LlmClient: def __init__( - self, log_level=ERROR, subscriber=None, + self, log_level=ERROR, + subscriber=None, + input_queue=None, + output_queue=None, pulsar_host="pulsar://pulsar:6650", ): + if input_queue == None: + input_queue = text_completion_request_queue + + if output_queue == None: + output_queue = text_completion_response_queue + if subscriber == None: subscriber = str(uuid.uuid4()) @@ -32,13 +41,13 @@ class LlmClient: ) self.producer = self.client.create_producer( - topic=text_completion_request_queue, + topic=input_queue, schema=JsonSchema(TextCompletionRequest), chunking_enabled=True, ) self.consumer = self.client.subscribe( - text_completion_response_queue, subscriber, + output_queue, subscriber, schema=JsonSchema(TextCompletionResponse), ) diff --git a/trustgraph/retrieval/graph_rag/rag.py b/trustgraph/retrieval/graph_rag/rag.py index 80cbfafb..e43aa0d8 100755 --- a/trustgraph/retrieval/graph_rag/rag.py +++ b/trustgraph/retrieval/graph_rag/rag.py @@ -6,6 +6,10 @@ Input is query, output is response. from ... schema import GraphRagQuery, GraphRagResponse from ... schema import graph_rag_request_queue, graph_rag_response_queue +from ... schema import text_completion_request_queue +from ... schema import text_completion_response_queue +from ... schema import embeddings_request_queue +from ... schema import embeddings_response_queue from ... log_level import LogLevel from ... graph_rag import GraphRag from ... base import ConsumerProducer @@ -30,6 +34,18 @@ class Processor(ConsumerProducer): entity_limit = params.get("entity_limit", 50) triple_limit = params.get("triple_limit", 30) max_subgraph_size = params.get("max_subgraph_size", 3000) + tc_request_queue = params.get( + "text_completion_request_queue", text_completion_request_queue + ) + tc_response_queue = params.get( + "text_completion_response_queue", text_completion_response_queue + ) + 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 | { @@ -41,12 +57,20 @@ class Processor(ConsumerProducer): "entity_limit": entity_limit, "triple_limit": triple_limit, "max_subgraph_size": max_subgraph_size, + "text_completion_request_queue": tc_request_queue, + "text_completion_response_queue": tc_response_queue, + "embeddings_request_queue": emb_request_queue, + "embeddings_response_queue": emb_response_queue, } ) self.rag = GraphRag( pulsar_host=self.pulsar_host, graph_hosts=graph_hosts.split(","), + completion_request_queue=tc_request_queue, + completion_response_queue=tc_response_queue, + emb_request_queue=emb_request_queue, + emb_response_queue=emb_response_queue, vector_store=vector_store, verbose=True, entity_limit=entity_limit, @@ -114,6 +138,30 @@ class Processor(ConsumerProducer): help=f'Max subgraph size (default: 3000)' ) + parser.add_argument( + '--text-completion-request-queue', + default=text_completion_request_queue, + help=f'Text completion request queue (default: {text_completion_request_queue})', + ) + + parser.add_argument( + '--text-completion-response-queue', + default=text_completion_response_queue, + help=f'Text completion response queue (default: {text_completion_response_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__)