apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-07 12:22:11 +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=$(shell git describe | sed 's/^v//')
VERSION=0.5.5 VERSION=0.5.6
all: container all: container

40
docker-compose-azure.yaml
View file

@ -130,7 +130,7 @@ services:
restart: on-failure:100 restart: on-failure:100
pdf-decoder: pdf-decoder:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "pdf-decoder" - "pdf-decoder"
- "-p" - "-p"
@ -138,7 +138,7 @@ services:
restart: on-failure:100 restart: on-failure:100
chunker: chunker:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "chunker-recursive" - "chunker-recursive"
- "-p" - "-p"
@ -146,7 +146,7 @@ services:
restart: on-failure:100 restart: on-failure:100
vectorize: vectorize:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "embeddings-vectorize" - "embeddings-vectorize"
- "-p" - "-p"
@ -154,7 +154,7 @@ services:
restart: on-failure:100 restart: on-failure:100
embeddings: embeddings:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "embeddings-hf" - "embeddings-hf"
- "-p" - "-p"
@ -164,7 +164,7 @@ services:
restart: on-failure:100 restart: on-failure:100
kg-extract-definitions: kg-extract-definitions:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "kg-extract-definitions" - "kg-extract-definitions"
- "-p" - "-p"
@ -172,7 +172,7 @@ services:
restart: on-failure:100 restart: on-failure:100
kg-extract-relationships: kg-extract-relationships:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "kg-extract-relationships" - "kg-extract-relationships"
- "-p" - "-p"
@ -180,7 +180,7 @@ services:
restart: on-failure:100 restart: on-failure:100
store-graph-embeddings: store-graph-embeddings:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "ge-write-milvus" - "ge-write-milvus"
- "-p" - "-p"
@ -190,7 +190,7 @@ services:
restart: on-failure:100 restart: on-failure:100
store-triples: store-triples:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "triples-write-cassandra" - "triples-write-cassandra"
- "-p" - "-p"
@ -200,7 +200,7 @@ services:
restart: on-failure:100 restart: on-failure:100
text-completion: text-completion:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "text-completion-azure" - "text-completion-azure"
- "-p" - "-p"
@ -211,11 +211,31 @@ services:
- ${AZURE_ENDPOINT} - ${AZURE_ENDPOINT}
restart: on-failure:100 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: graph-rag:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "graph-rag" - "graph-rag"
- "-p" - "-p"
- "pulsar://pulsar:6650" - "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 restart: on-failure:100

39
docker-compose-claude.yaml
View file

@ -130,7 +130,7 @@ services:
restart: on-failure:100 restart: on-failure:100
pdf-decoder: pdf-decoder:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "pdf-decoder" - "pdf-decoder"
- "-p" - "-p"
@ -138,7 +138,7 @@ services:
restart: on-failure:100 restart: on-failure:100
chunker: chunker:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "chunker-recursive" - "chunker-recursive"
- "-p" - "-p"
@ -146,7 +146,7 @@ services:
restart: on-failure:100 restart: on-failure:100
vectorize: vectorize:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "embeddings-vectorize" - "embeddings-vectorize"
- "-p" - "-p"
@ -154,7 +154,7 @@ services:
restart: on-failure:100 restart: on-failure:100
embeddings: embeddings:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "embeddings-hf" - "embeddings-hf"
- "-p" - "-p"
@ -164,7 +164,7 @@ services:
restart: on-failure:100 restart: on-failure:100
kg-extract-definitions: kg-extract-definitions:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "kg-extract-definitions" - "kg-extract-definitions"
- "-p" - "-p"
@ -172,7 +172,7 @@ services:
restart: on-failure:100 restart: on-failure:100
kg-extract-relationships: kg-extract-relationships:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "kg-extract-relationships" - "kg-extract-relationships"
- "-p" - "-p"
@ -180,7 +180,7 @@ services:
restart: on-failure:100 restart: on-failure:100
store-graph-embeddings: store-graph-embeddings:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "ge-write-milvus" - "ge-write-milvus"
- "-p" - "-p"
@ -190,7 +190,7 @@ services:
restart: on-failure:100 restart: on-failure:100
store-triples: store-triples:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "triples-write-cassandra" - "triples-write-cassandra"
- "-p" - "-p"
@ -200,7 +200,7 @@ services:
restart: on-failure:100 restart: on-failure:100
text-completion: text-completion:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "text-completion-claude" - "text-completion-claude"
- "-p" - "-p"
@ -209,11 +209,28 @@ services:
- ${CLAUDE_KEY} - ${CLAUDE_KEY}
restart: on-failure:100 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: graph-rag:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "graph-rag" - "graph-rag"
- "-p" - "-p"
- "pulsar://pulsar:6650" - "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 restart: on-failure:100

40
docker-compose-ollama.yaml
View file

@ -130,7 +130,7 @@ services:
restart: on-failure:100 restart: on-failure:100
pdf-decoder: pdf-decoder:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "pdf-decoder" - "pdf-decoder"
- "-p" - "-p"
@ -138,7 +138,7 @@ services:
restart: on-failure:100 restart: on-failure:100
chunker: chunker:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "chunker-recursive" - "chunker-recursive"
- "-p" - "-p"
@ -146,7 +146,7 @@ services:
restart: on-failure:100 restart: on-failure:100
vectorize: vectorize:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "embeddings-vectorize" - "embeddings-vectorize"
- "-p" - "-p"
@ -154,7 +154,7 @@ services:
restart: on-failure:100 restart: on-failure:100
embeddings: embeddings:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "embeddings-hf" - "embeddings-hf"
- "-p" - "-p"
@ -164,7 +164,7 @@ services:
restart: on-failure:100 restart: on-failure:100
kg-extract-definitions: kg-extract-definitions:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "kg-extract-definitions" - "kg-extract-definitions"
- "-p" - "-p"
@ -172,7 +172,7 @@ services:
restart: on-failure:100 restart: on-failure:100
kg-extract-relationships: kg-extract-relationships:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "kg-extract-relationships" - "kg-extract-relationships"
- "-p" - "-p"
@ -180,7 +180,7 @@ services:
restart: on-failure:100 restart: on-failure:100
store-graph-embeddings: store-graph-embeddings:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "ge-write-milvus" - "ge-write-milvus"
- "-p" - "-p"
@ -190,7 +190,7 @@ services:
restart: on-failure:100 restart: on-failure:100
store-triples: store-triples:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "triples-write-cassandra" - "triples-write-cassandra"
- "-p" - "-p"
@ -200,7 +200,7 @@ services:
restart: on-failure:100 restart: on-failure:100
text-completion: text-completion:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "text-completion-ollama" - "text-completion-ollama"
- "-p" - "-p"
@ -210,12 +210,32 @@ services:
# - "-m" # - "-m"
# - "llama3.1:8b" # - "llama3.1:8b"
restart: on-failure:100 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: graph-rag:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "graph-rag" - "graph-rag"
- "-p" - "-p"
- "pulsar://pulsar:6650" - "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 restart: on-failure:100

42
docker-compose-vertexai.yaml
View file

@ -130,7 +130,7 @@ services:
restart: on-failure:100 restart: on-failure:100
pdf-decoder: pdf-decoder:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "pdf-decoder" - "pdf-decoder"
- "-p" - "-p"
@ -138,7 +138,7 @@ services:
restart: on-failure:100 restart: on-failure:100
chunker: chunker:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "chunker-recursive" - "chunker-recursive"
- "-p" - "-p"
@ -146,7 +146,7 @@ services:
restart: on-failure:100 restart: on-failure:100
vectorize: vectorize:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "embeddings-vectorize" - "embeddings-vectorize"
- "-p" - "-p"
@ -154,7 +154,7 @@ services:
restart: on-failure:100 restart: on-failure:100
embeddings: embeddings:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "embeddings-hf" - "embeddings-hf"
- "-p" - "-p"
@ -164,7 +164,7 @@ services:
restart: on-failure:100 restart: on-failure:100
kg-extract-definitions: kg-extract-definitions:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "kg-extract-definitions" - "kg-extract-definitions"
- "-p" - "-p"
@ -172,7 +172,7 @@ services:
restart: on-failure:100 restart: on-failure:100
kg-extract-relationships: kg-extract-relationships:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "kg-extract-relationships" - "kg-extract-relationships"
- "-p" - "-p"
@ -180,7 +180,7 @@ services:
restart: on-failure:100 restart: on-failure:100
store-graph-embeddings: store-graph-embeddings:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "ge-write-milvus" - "ge-write-milvus"
- "-p" - "-p"
@ -190,7 +190,7 @@ services:
restart: on-failure:100 restart: on-failure:100
store-triples: store-triples:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "triples-write-cassandra" - "triples-write-cassandra"
- "-p" - "-p"
@ -200,9 +200,9 @@ services:
restart: on-failure:100 restart: on-failure:100
text-completion: text-completion:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "text-completion-ollama" - "text-completion-vertexai"
- "-p" - "-p"
- "pulsar://pulsar:6650" - "pulsar://pulsar:6650"
- "-k" - "-k"
@ -212,12 +212,32 @@ services:
volumes: volumes:
- "./vertexai:/vertexai" - "./vertexai:/vertexai"
restart: on-failure:100 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: graph-rag:
image: docker.io/trustgraph/trustgraph-flow:0.5.5 image: docker.io/trustgraph/trustgraph-flow:0.5.6
command: command:
- "graph-rag" - "graph-rag"
- "-p" - "-p"
- "pulsar://pulsar:6650" - "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 restart: on-failure:100

5
requirements.txt
View file

@ -10,7 +10,12 @@ langchain-huggingface
langchain-text-splitters langchain-text-splitters
langchain-community langchain-community
huggingface-hub huggingface-hub
requests
cassandra-driver cassandra-driver
pulsar-client pulsar-client
pypdf
anthropic anthropic
google-cloud-aiplatform 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: with open("README.md", "r") as fh:
long_description = fh.read() long_description = fh.read()
version = "0.5.5" version = "0.5.6"
setuptools.setup( setuptools.setup(
name="trustgraph", 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, ChunkEmbeddings
from ... schema import chunk_ingest_queue, chunk_embeddings_ingest_queue 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 ... embeddings_client import EmbeddingsClient
from ... log_level import LogLevel from ... log_level import LogLevel
from ... base import ConsumerProducer from ... base import ConsumerProducer
@ -23,11 +24,19 @@ class Processor(ConsumerProducer):
input_queue = params.get("input_queue", default_input_queue) input_queue = params.get("input_queue", default_input_queue)
output_queue = params.get("output_queue", default_output_queue) output_queue = params.get("output_queue", default_output_queue)
subscriber = params.get("subscriber", default_subscriber) 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__( super(Processor, self).__init__(
**params | { **params | {
"input_queue": input_queue, "input_queue": input_queue,
"output_queue": output_queue, "output_queue": output_queue,
"embeddings_request_queue": emb_request_queue,
"embeddings_response_queue": emb_response_queue,
"subscriber": subscriber, "subscriber": subscriber,
"input_schema": Chunk, "input_schema": Chunk,
"output_schema": ChunkEmbeddings, "output_schema": ChunkEmbeddings,
@ -36,6 +45,8 @@ class Processor(ConsumerProducer):
self.embeddings = EmbeddingsClient( self.embeddings = EmbeddingsClient(
pulsar_host=self.pulsar_host, pulsar_host=self.pulsar_host,
input_queue=emb_request_queue,
output_queue=emb_response_queue,
subscriber=module + "-emb", subscriber=module + "-emb",
) )
@ -74,6 +85,18 @@ class Processor(ConsumerProducer):
default_output_queue, 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(): def run():
Processor.start(module, __doc__) Processor.start(module, __doc__)

15
trustgraph/embeddings_client.py
View file

@ -17,12 +17,21 @@ DEBUG=_pulsar.LoggerLevel.Debug
class EmbeddingsClient: class EmbeddingsClient:
def __init__( 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", pulsar_host="pulsar://pulsar:6650",
): ):
self.client = None self.client = None
if input_queue == None:
input_queue=embeddings_request_queue
if output_queue == None:
output_queue=embeddings_response_queue
if subscriber == None: if subscriber == None:
subscriber = str(uuid.uuid4()) subscriber = str(uuid.uuid4())
@ -32,13 +41,13 @@ class EmbeddingsClient:
) )
self.producer = self.client.create_producer( self.producer = self.client.create_producer(
topic=embeddings_request_queue, topic=input_queue,
schema=JsonSchema(EmbeddingsRequest), schema=JsonSchema(EmbeddingsRequest),
chunking_enabled=True, chunking_enabled=True,
) )
self.consumer = self.client.subscribe( self.consumer = self.client.subscribe(
embeddings_response_queue, subscriber, output_queue, subscriber,
schema=JsonSchema(EmbeddingsResponse), 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.trustgraph import TrustGraph
from trustgraph.llm_client import LlmClient from trustgraph.llm_client import LlmClient
from trustgraph.embeddings_client import EmbeddingsClient 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" LABEL="http://www.w3.org/2000/01/rdf-schema#label"
DEFINITION="http://www.w3.org/2004/02/skos/core#definition" DEFINITION="http://www.w3.org/2004/02/skos/core#definition"
@ -15,6 +17,10 @@ class GraphRag:
graph_hosts=None, graph_hosts=None,
pulsar_host="pulsar://pulsar:6650", pulsar_host="pulsar://pulsar:6650",
vector_store="http://milvus:19530", vector_store="http://milvus:19530",
completion_request_queue=None,
completion_response_queue=None,
emb_request_queue=None,
emb_response_queue=None,
verbose=False, verbose=False,
entity_limit=50, entity_limit=50,
triple_limit=30, triple_limit=30,
@ -24,6 +30,18 @@ class GraphRag:
self.verbose=verbose 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: if graph_hosts == None:
graph_hosts = ["cassandra"] graph_hosts = ["cassandra"]
@ -34,6 +52,8 @@ class GraphRag:
self.embeddings = EmbeddingsClient( self.embeddings = EmbeddingsClient(
pulsar_host=pulsar_host, pulsar_host=pulsar_host,
input_queue=emb_request_queue,
output_queue=emb_response_queue,
subscriber=module + "-emb", subscriber=module + "-emb",
) )
@ -47,6 +67,8 @@ class GraphRag:
self.llm = LlmClient( self.llm = LlmClient(
pulsar_host=pulsar_host, pulsar_host=pulsar_host,
input_queue=completion_request_queue,
output_queue=completion_response_queue,
subscriber=module + "-llm", 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 ChunkEmbeddings, Triple, Source, Value
from ... schema import chunk_embeddings_ingest_queue, triples_store_queue 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 ... log_level import LogLevel
from ... llm_client import LlmClient from ... llm_client import LlmClient
from ... prompts import to_definitions from ... prompts import to_definitions
@ -30,6 +32,12 @@ class Processor(ConsumerProducer):
input_queue = params.get("input_queue", default_input_queue) input_queue = params.get("input_queue", default_input_queue)
output_queue = params.get("output_queue", default_output_queue) output_queue = params.get("output_queue", default_output_queue)
subscriber = params.get("subscriber", default_subscriber) 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__( super(Processor, self).__init__(
**params | { **params | {
@ -38,11 +46,15 @@ class Processor(ConsumerProducer):
"subscriber": subscriber, "subscriber": subscriber,
"input_schema": ChunkEmbeddings, "input_schema": ChunkEmbeddings,
"output_schema": Triple, "output_schema": Triple,
"text_completion_request_queue": tc_request_queue,
"text_completion_response_queue": tc_response_queue,
} }
) )
self.llm = LlmClient( self.llm = LlmClient(
pulsar_host=self.pulsar_host, pulsar_host=self.pulsar_host,
input_queue=tc_request_queue,
output_queue=tc_response_queue,
subscriber = module + "-llm", subscriber = module + "-llm",
) )
@ -108,6 +120,18 @@ class Processor(ConsumerProducer):
default_output_queue, 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(): def run():
Processor.start(module, __doc__) 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 pulsar.schema import JsonSchema
from ... schema import ChunkEmbeddings, Triple, GraphEmbeddings, Source, Value 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 ... log_level import LogLevel
from ... llm_client import LlmClient from ... llm_client import LlmClient
from ... prompts import to_relationships from ... prompts import to_relationships
@ -35,6 +38,12 @@ class Processor(ConsumerProducer):
output_queue = params.get("output_queue", default_output_queue) output_queue = params.get("output_queue", default_output_queue)
vector_queue = params.get("vector_queue", default_vector_queue) vector_queue = params.get("vector_queue", default_vector_queue)
subscriber = params.get("subscriber", default_subscriber) 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__( super(Processor, self).__init__(
**params | { **params | {
@ -43,6 +52,8 @@ class Processor(ConsumerProducer):
"subscriber": subscriber, "subscriber": subscriber,
"input_schema": ChunkEmbeddings, "input_schema": ChunkEmbeddings,
"output_schema": Triple, "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, "input_queue": input_queue,
"output_queue": output_queue, "output_queue": output_queue,
"vector_queue": vector_queue, "vector_queue": vector_queue,
"text_completion_request_queue": tc_request_queue,
"text_completion_response_queue": tc_response_queue,
"subscriber": subscriber, "subscriber": subscriber,
"input_schema": ChunkEmbeddings.__name__, "input_schema": ChunkEmbeddings.__name__,
"output_schema": Triple.__name__, "output_schema": Triple.__name__,
@ -63,6 +76,8 @@ class Processor(ConsumerProducer):
self.llm = LlmClient( self.llm = LlmClient(
pulsar_host = self.pulsar_host, pulsar_host = self.pulsar_host,
input_queue=tc_request_queue,
output_queue=tc_response_queue,
subscriber = module + "-llm", subscriber = module + "-llm",
) )
@ -179,6 +194,18 @@ class Processor(ConsumerProducer):
help=f'Vector output queue (default: {default_vector_queue})' 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(): def run():
Processor.start(module, __doc__) Processor.start(module, __doc__)

15
trustgraph/llm_client.py
View file

@ -19,10 +19,19 @@ DEBUG=_pulsar.LoggerLevel.Debug
class LlmClient: class LlmClient:
def __init__( 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", 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: if subscriber == None:
subscriber = str(uuid.uuid4()) subscriber = str(uuid.uuid4())
@ -32,13 +41,13 @@ class LlmClient:
) )
self.producer = self.client.create_producer( self.producer = self.client.create_producer(
topic=text_completion_request_queue, topic=input_queue,
schema=JsonSchema(TextCompletionRequest), schema=JsonSchema(TextCompletionRequest),
chunking_enabled=True, chunking_enabled=True,
) )
self.consumer = self.client.subscribe( self.consumer = self.client.subscribe(
text_completion_response_queue, subscriber, output_queue, subscriber,
schema=JsonSchema(TextCompletionResponse), 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 GraphRagQuery, GraphRagResponse
from ... schema import graph_rag_request_queue, graph_rag_response_queue 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 ... log_level import LogLevel
from ... graph_rag import GraphRag from ... graph_rag import GraphRag
from ... base import ConsumerProducer from ... base import ConsumerProducer
@ -30,6 +34,18 @@ class Processor(ConsumerProducer):
entity_limit = params.get("entity_limit", 50) entity_limit = params.get("entity_limit", 50)
triple_limit = params.get("triple_limit", 30) triple_limit = params.get("triple_limit", 30)
max_subgraph_size = params.get("max_subgraph_size", 3000) 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__( super(Processor, self).__init__(
**params | { **params | {
@ -41,12 +57,20 @@ class Processor(ConsumerProducer):
"entity_limit": entity_limit, "entity_limit": entity_limit,
"triple_limit": triple_limit, "triple_limit": triple_limit,
"max_subgraph_size": max_subgraph_size, "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( self.rag = GraphRag(
pulsar_host=self.pulsar_host, pulsar_host=self.pulsar_host,
graph_hosts=graph_hosts.split(","), 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, vector_store=vector_store,
verbose=True, verbose=True,
entity_limit=entity_limit, entity_limit=entity_limit,
@ -114,6 +138,30 @@ class Processor(ConsumerProducer):
help=f'Max subgraph size (default: 3000)' 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(): def run():
Processor.start(module, __doc__) Processor.start(module, __doc__)