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

Merge commit '63a30b25eb' into cohere

Browse source
This commit is contained in:
JackColquitt 2024-08-05 16:13:49 -07:00
commit 0a15e58217
14 changed files with 295 additions and 51 deletions
Download patch file Download diff file Expand all files Collapse all files

2
Makefile
View file

@ -1,6 +1,6 @@
# VERSION=$(shell git describe | sed 's/^v//')
VERSION=0.5.5
VERSION=0.5.6
all: container

40
docker-compose-azure.yaml
View file

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

39
docker-compose-claude.yaml
View file

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

40
docker-compose-ollama.yaml
View file

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

42
docker-compose-vertexai.yaml
View file

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

5
requirements.txt
View file

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

2
setup.py
View file

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

23
trustgraph/embeddings/vectorize/vectorize.py
View file

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

15
trustgraph/embeddings_client.py
View file

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

22
trustgraph/graph_rag.py
View file

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

24
trustgraph/kg/extract_definitions/extract.py
View file

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

29
trustgraph/kg/extract_relationships/extract.py
View file

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

15
trustgraph/llm_client.py
View file

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

48
trustgraph/retrieval/graph_rag/rag.py
View file

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