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

Feature/configure flows (#345)

Browse source 
- Keeps processing in different flows separate so that data can go to different stores / collections etc.
- Potentially supports different processing flows
- Tidies the processing API with common base-classes for e.g. LLMs, and automatic configuration of 'clients' to use the right queue names in a flow
This commit is contained in:
cybermaggedon 2025-04-22 20:21:38 +01:00 • committed by Cyber MacGeddon
parent dc0ce1041b
commit 31328317fd
125 changed files with 3751 additions and 2628 deletions
Download patch file Download diff file Expand all files Collapse all files

25
trustgraph-flow/trustgraph/agent/react/agent_manager.py
View file

@ -8,12 +8,11 @@ logger = logging.getLogger(__name__)
class AgentManager:
def __init__(self, context, tools, additional_context=None):
self.context = context
def __init__(self, tools, additional_context=None):
self.tools = tools
self.additional_context = additional_context
def reason(self, question, history):
async def reason(self, question, history, context):
tools = self.tools
@ -56,10 +55,7 @@ class AgentManager:
logger.info(f"prompt: {variables}")
obj = self.context.prompt.request(
"agent-react",
variables
)
obj = await context("prompt-request").agent_react(variables)
print(json.dumps(obj, indent=4), flush=True)
@ -85,9 +81,13 @@ class AgentManager:
return a
async def react(self, question, history, think, observe):
async def react(self, question, history, think, observe, context):
act = self.reason(question, history)
act = await self.reason(
question = question,
history = history,
context = context,
)
logger.info(f"act: {act}")
if isinstance(act, Final):
@ -104,7 +104,12 @@ class AgentManager:
else:
raise RuntimeError(f"No action for {act.name}!")
resp = action.implementation.invoke(**act.arguments)
print("TOOL>>>", act)
resp = await action.implementation(context).invoke(
**act.arguments
)
print("RSETUL", resp)
resp = resp.strip()

188
trustgraph-flow/trustgraph/agent/react/service.py
View file

@ -6,103 +6,68 @@ import json
import re
import sys
from pulsar.schema import JsonSchema
from ... base import AgentService, TextCompletionClientSpec, PromptClientSpec
from ... base import GraphRagClientSpec
from ... base import ConsumerProducer
from ... schema import Error
from ... schema import AgentRequest, AgentResponse, AgentStep
from ... schema import agent_request_queue, agent_response_queue
from ... schema import prompt_request_queue as pr_request_queue
from ... schema import prompt_response_queue as pr_response_queue
from ... schema import graph_rag_request_queue as gr_request_queue
from ... schema import graph_rag_response_queue as gr_response_queue
from ... clients.prompt_client import PromptClient
from ... clients.llm_client import LlmClient
from ... clients.graph_rag_client import GraphRagClient
from ... schema import AgentRequest, AgentResponse, AgentStep, Error
from . tools import KnowledgeQueryImpl, TextCompletionImpl
from . agent_manager import AgentManager
from . types import Final, Action, Tool, Argument
module = ".".join(__name__.split(".")[1:-1])
default_ident = "agent-manager"
default_max_iterations = 10
default_input_queue = agent_request_queue
default_output_queue = agent_response_queue
default_subscriber = module
default_max_iterations = 15
class Processor(ConsumerProducer):
class Processor(AgentService):
def __init__(self, **params):
id = params.get("id")
self.max_iterations = int(
params.get("max_iterations", default_max_iterations)
)
tools = {}
input_queue = params.get("input_queue", default_input_queue)
output_queue = params.get("output_queue", default_output_queue)
subscriber = params.get("subscriber", default_subscriber)
prompt_request_queue = params.get(
"prompt_request_queue", pr_request_queue
)
prompt_response_queue = params.get(
"prompt_response_queue", pr_response_queue
)
graph_rag_request_queue = params.get(
"graph_rag_request_queue", gr_request_queue
)
graph_rag_response_queue = params.get(
"graph_rag_response_queue", gr_response_queue
)
self.config_key = params.get("config_type", "agent")
super(Processor, self).__init__(
**params | {
"input_queue": input_queue,
"output_queue": output_queue,
"subscriber": subscriber,
"input_schema": AgentRequest,
"output_schema": AgentResponse,
"prompt_request_queue": prompt_request_queue,
"prompt_response_queue": prompt_response_queue,
"graph_rag_request_queue": gr_request_queue,
"graph_rag_response_queue": gr_response_queue,
"id": id,
"max_iterations": self.max_iterations,
"config_type": self.config_key,
}
)
self.prompt = PromptClient(
subscriber=subscriber,
input_queue=prompt_request_queue,
output_queue=prompt_response_queue,
pulsar_host = self.pulsar_host,
pulsar_api_key=self.pulsar_api_key,
)
self.graph_rag = GraphRagClient(
subscriber=subscriber,
input_queue=graph_rag_request_queue,
output_queue=graph_rag_response_queue,
pulsar_host = self.pulsar_host,
pulsar_api_key=self.pulsar_api_key,
)
# Need to be able to feed requests to myself
self.recursive_input = self.client.create_producer(
topic=input_queue,
schema=JsonSchema(AgentRequest),
)
self.agent = AgentManager(
context=self,
tools=[],
additional_context="",
)
async def on_config(self, version, config):
self.config_handlers.append(self.on_tools_config)
self.register_specification(
TextCompletionClientSpec(
request_name = "text-completion-request",
response_name = "text-completion-response",
)
)
self.register_specification(
GraphRagClientSpec(
request_name = "graph-rag-request",
response_name = "graph-rag-response",
)
)
self.register_specification(
PromptClientSpec(
request_name = "prompt-request",
response_name = "prompt-response",
)
)
async def on_tools_config(self, config, version):
print("Loading configuration version", version)
@ -138,9 +103,9 @@ class Processor(ConsumerProducer):
impl_id = data.get("type")
if impl_id == "knowledge-query":
impl = KnowledgeQueryImpl(self)
impl = KnowledgeQueryImpl
elif impl_id == "text-completion":
impl = TextCompletionImpl(self)
impl = TextCompletionImpl
else:
raise RuntimeError(
f"Tool-kind {impl_id} not known"
@ -155,7 +120,6 @@ class Processor(ConsumerProducer):
)
self.agent = AgentManager(
context=self,
tools=tools,
additional_context=additional
)
@ -164,19 +128,14 @@ class Processor(ConsumerProducer):
except Exception as e:
print("Exception:", e, flush=True)
print("on_tools_config Exception:", e, flush=True)
print("Configuration reload failed", flush=True)
async def handle(self, msg):
async def agent_request(self, request, respond, next, flow):
try:
v = msg.value()
# Sender-produced ID
id = msg.properties()["id"]
if v.history:
if request.history:
history = [
Action(
thought=h.thought,
@ -184,12 +143,12 @@ class Processor(ConsumerProducer):
arguments=h.arguments,
observation=h.observation
)
for h in v.history
for h in request.history
]
else:
history = []
print(f"Question: {v.question}", flush=True)
print(f"Question: {request.question}", flush=True)
if len(history) >= self.max_iterations:
raise RuntimeError("Too many agent iterations")
@ -207,7 +166,7 @@ class Processor(ConsumerProducer):
observation=None,
)
await self.send(r, properties={"id": id})
await respond(r)
async def observe(x):
@ -220,15 +179,21 @@ class Processor(ConsumerProducer):
observation=x,
)
await self.send(r, properties={"id": id})
await respond(r)
act = await self.agent.react(v.question, history, think, observe)
act = await self.agent.react(
question = request.question,
history = history,
think = think,
observe = observe,
context = flow,
)
print(f"Action: {act}", flush=True)
print("Send response...", flush=True)
if isinstance(act, Final):
if type(act) == Final:
print("Send final response...", flush=True)
r = AgentResponse(
answer=act.final,
@ -236,18 +201,20 @@ class Processor(ConsumerProducer):
thought=None,
)
await self.send(r, properties={"id": id})
await respond(r)
print("Done.", flush=True)
return
print("Send next...", flush=True)
history.append(act)
r = AgentRequest(
question=v.question,
plan=v.plan,
state=v.state,
question=request.question,
plan=request.plan,
state=request.state,
history=[
AgentStep(
thought=h.thought,
@ -259,7 +226,7 @@ class Processor(ConsumerProducer):
]
)
self.recursive_input.send(r, properties={"id": id})
await next(r)
print("Done.", flush=True)
@ -267,7 +234,7 @@ class Processor(ConsumerProducer):
except Exception as e:
print(f"Exception: {e}")
print(f"agent_request Exception: {e}")
print("Send error response...", flush=True)
@ -279,39 +246,12 @@ class Processor(ConsumerProducer):
response=None,
)
await self.send(r, properties={"id": id})
await respond(r)
@staticmethod
def add_args(parser):
ConsumerProducer.add_args(
parser, default_input_queue, default_subscriber,
default_output_queue,
)
parser.add_argument(
'--prompt-request-queue',
default=pr_request_queue,
help=f'Prompt request queue (default: {pr_request_queue})',
)
parser.add_argument(
'--prompt-response-queue',
default=pr_response_queue,
help=f'Prompt response queue (default: {pr_response_queue})',
)
parser.add_argument(
'--graph-rag-request-queue',
default=gr_request_queue,
help=f'Graph RAG request queue (default: {gr_request_queue})',
)
parser.add_argument(
'--graph-rag-response-queue',
default=gr_response_queue,
help=f'Graph RAG response queue (default: {gr_response_queue})',
)
AgentService.add_args(parser)
parser.add_argument(
'--max-iterations',
@ -327,5 +267,5 @@ class Processor(ConsumerProducer):
def run():
Processor.launch(module, __doc__)
Processor.launch(default_ident, __doc__)

16
trustgraph-flow/trustgraph/agent/react/tools.py
View file

@ -4,16 +4,22 @@
class KnowledgeQueryImpl:
def __init__(self, context):
self.context = context
def invoke(self, **arguments):
return self.context.graph_rag.request(arguments.get("question"))
async def invoke(self, **arguments):
client = self.context("graph-rag-request")
print("Graph RAG question...", flush=True)
return await client.rag(
arguments.get("question")
)
# This tool implementation knows how to do text completion. This uses
# the prompt service, rather than talking to TextCompletion directly.
class TextCompletionImpl:
def __init__(self, context):
self.context = context
def invoke(self, **arguments):
return self.context.prompt.request(
"question", { "question": arguments.get("question") }
async def invoke(self, **arguments):
client = self.context("prompt-request")
print("Prompt question...", flush=True)
return await client.question(
arguments.get("question")
)

61
trustgraph-flow/trustgraph/chunking/recursive/chunker.py
View file

@ -7,40 +7,27 @@ as text as separate output objects.
from langchain_text_splitters import RecursiveCharacterTextSplitter
from prometheus_client import Histogram
from ... schema import TextDocument, Chunk, Metadata
from ... schema import text_ingest_queue, chunk_ingest_queue
from ... log_level import LogLevel
from ... base import ConsumerProducer
from ... schema import TextDocument, Chunk
from ... base import FlowProcessor, ConsumerSpec, ProducerSpec
module = ".".join(__name__.split(".")[1:-1])
default_ident = "chunker"
default_input_queue = text_ingest_queue
default_output_queue = chunk_ingest_queue
default_subscriber = module
class Processor(ConsumerProducer):
class Processor(FlowProcessor):
def __init__(self, **params):
input_queue = params.get("input_queue", default_input_queue)
output_queue = params.get("output_queue", default_output_queue)
subscriber = params.get("subscriber", default_subscriber)
id = params.get("id", default_ident)
chunk_size = params.get("chunk_size", 2000)
chunk_overlap = params.get("chunk_overlap", 100)
super(Processor, self).__init__(
**params | {
"input_queue": input_queue,
"output_queue": output_queue,
"subscriber": subscriber,
"input_schema": TextDocument,
"output_schema": Chunk,
}
**params | { "id": id }
)
if not hasattr(__class__, "chunk_metric"):
__class__.chunk_metric = Histogram(
'chunk_size', 'Chunk size',
["id", "flow"],
buckets=[100, 160, 250, 400, 650, 1000, 1600,
2500, 4000, 6400, 10000, 16000]
)
@ -52,7 +39,24 @@ class Processor(ConsumerProducer):
is_separator_regex=False,
)
async def handle(self, msg):
self.register_specification(
ConsumerSpec(
name = "input",
schema = TextDocument,
handler = self.on_message,
)
)
self.register_specification(
ProducerSpec(
name = "output",
schema = Chunk,
)
)
print("Chunker initialised", flush=True)
async def on_message(self, msg, consumer, flow):
v = msg.value()
print(f"Chunking {v.metadata.id}...", flush=True)
@ -63,24 +67,25 @@ class Processor(ConsumerProducer):
for ix, chunk in enumerate(texts):
print("Chunk", len(chunk.page_content), flush=True)
r = Chunk(
metadata=v.metadata,
chunk=chunk.page_content.encode("utf-8"),
)
__class__.chunk_metric.observe(len(chunk.page_content))
__class__.chunk_metric.labels(
id=consumer.id, flow=consumer.flow
).observe(len(chunk.page_content))
await self.send(r)
await flow("output").send(r)
print("Done.", flush=True)
@staticmethod
def add_args(parser):
ConsumerProducer.add_args(
parser, default_input_queue, default_subscriber,
default_output_queue,
)
FlowProcessor.add_args(parser)
parser.add_argument(
'-z', '--chunk-size',
@ -98,5 +103,5 @@ class Processor(ConsumerProducer):
def run():
Processor.launch(module, __doc__)
Processor.launch(default_ident, __doc__)

61
trustgraph-flow/trustgraph/chunking/token/chunker.py
View file

@ -7,40 +7,27 @@ as text as separate output objects.
from langchain_text_splitters import TokenTextSplitter
from prometheus_client import Histogram
from ... schema import TextDocument, Chunk, Metadata
from ... schema import text_ingest_queue, chunk_ingest_queue
from ... log_level import LogLevel
from ... base import ConsumerProducer
from ... schema import TextDocument, Chunk
from ... base import FlowProcessor
module = ".".join(__name__.split(".")[1:-1])
default_ident = "chunker"
default_input_queue = text_ingest_queue
default_output_queue = chunk_ingest_queue
default_subscriber = module
class Processor(ConsumerProducer):
class Processor(FlowProcessor):
def __init__(self, **params):
input_queue = params.get("input_queue", default_input_queue)
output_queue = params.get("output_queue", default_output_queue)
subscriber = params.get("subscriber", default_subscriber)
id = params.get("id")
chunk_size = params.get("chunk_size", 250)
chunk_overlap = params.get("chunk_overlap", 15)
super(Processor, self).__init__(
**params | {
"input_queue": input_queue,
"output_queue": output_queue,
"subscriber": subscriber,
"input_schema": TextDocument,
"output_schema": Chunk,
}
**params | { "id": id }
)
if not hasattr(__class__, "chunk_metric"):
__class__.chunk_metric = Histogram(
'chunk_size', 'Chunk size',
["id", "flow"],
buckets=[100, 160, 250, 400, 650, 1000, 1600,
2500, 4000, 6400, 10000, 16000]
)
@ -51,7 +38,24 @@ class Processor(ConsumerProducer):
chunk_overlap=chunk_overlap,
)
async def handle(self, msg):
self.register_specification(
ConsumerSpec(
name = "input",
schema = TextDocument,
handler = self.on_message,
)
)
self.register_specification(
ProducerSpec(
name = "output",
schema = Chunk,
)
)
print("Chunker initialised", flush=True)
async def on_message(self, msg, consumer, flow):
v = msg.value()
print(f"Chunking {v.metadata.id}...", flush=True)
@ -62,24 +66,25 @@ class Processor(ConsumerProducer):
for ix, chunk in enumerate(texts):
print("Chunk", len(chunk.page_content), flush=True)
r = Chunk(
metadata=v.metadata,
chunk=chunk.page_content.encode("utf-8"),
)
__class__.chunk_metric.observe(len(chunk.page_content))
__class__.chunk_metric.labels(
id=consumer.id, flow=consumer.flow
).observe(len(chunk.page_content))
await self.send(r)
await flow("output").send(r)
print("Done.", flush=True)
@staticmethod
def add_args(parser):
ConsumerProducer.add_args(
parser, default_input_queue, default_subscriber,
default_output_queue,
)
FlowProcessor.add_args(parser)
parser.add_argument(
'-z', '--chunk-size',
@ -97,5 +102,5 @@ class Processor(ConsumerProducer):
def run():
Processor.launch(module, __doc__)
Processor.launch(default_ident, __doc__)

215
trustgraph-flow/trustgraph/config/service/config.py Normal file
View file

@ -0,0 +1,215 @@
from trustgraph.schema import ConfigResponse
from trustgraph.schema import ConfigValue, Error
# This behaves just like a dict, should be easier to add persistent storage
# later
class ConfigurationItems(dict):
pass
class Configuration(dict):
# FIXME: The state is held internally. This only works if there's
# one config service. Should be more than one, and use a
# back-end state store.
def __init__(self, push):
# Version counter
self.version = 0
# External function to respond to update
self.push = push
def __getitem__(self, key):
if key not in self:
self[key] = ConfigurationItems()
return dict.__getitem__(self, key)
async def handle_get(self, v):
for k in v.keys:
if k.type not in self or k.key not in self[k.type]:
return ConfigResponse(
version = None,
values = None,
directory = None,
config = None,
error = Error(
type = "key-error",
message = f"Key error"
)
)
values = [
ConfigValue(
type = k.type,
key = k.key,
value = self[k.type][k.key]
)
for k in v.keys
]
return ConfigResponse(
version = self.version,
values = values,
directory = None,
config = None,
error = None,
)
async def handle_list(self, v):
if v.type not in self:
return ConfigResponse(
version = None,
values = None,
directory = None,
config = None,
error = Error(
type = "key-error",
message = "No such type",
),
)
return ConfigResponse(
version = self.version,
values = None,
directory = list(self[v.type].keys()),
config = None,
error = None,
)
async def handle_getvalues(self, v):
if v.type not in self:
return ConfigResponse(
version = None,
values = None,
directory = None,
config = None,
error = Error(
type = "key-error",
message = f"Key error"
)
)
values = [
ConfigValue(
type = v.type,
key = k,
value = self[v.type][k],
)
for k in self[v.type]
]
return ConfigResponse(
version = self.version,
values = values,
directory = None,
config = None,
error = None,
)
async def handle_delete(self, v):
for k in v.keys:
if k.type not in self or k.key not in self[k.type]:
return ConfigResponse(
version = None,
values = None,
directory = None,
config = None,
error = Error(
type = "key-error",
message = f"Key error"
)
)
for k in v.keys:
del self[k.type][k.key]
self.version += 1
await self.push()
return ConfigResponse(
version = None,
value = None,
directory = None,
values = None,
config = None,
error = None,
)
async def handle_put(self, v):
for k in v.values:
self[k.type][k.key] = k.value
self.version += 1
await self.push()
return ConfigResponse(
version = None,
value = None,
directory = None,
values = None,
error = None,
)
async def handle_config(self, v):
return ConfigResponse(
version = self.version,
value = None,
directory = None,
values = None,
config = self,
error = None,
)
async def handle(self, msg):
print("Handle message ", msg.operation)
if msg.operation == "get":
resp = await self.handle_get(msg)
elif msg.operation == "list":
resp = await self.handle_list(msg)
elif msg.operation == "getvalues":
resp = await self.handle_getvalues(msg)
elif msg.operation == "delete":
resp = await self.handle_delete(msg)
elif msg.operation == "put":
resp = await self.handle_put(msg)
elif msg.operation == "config":
resp = await self.handle_config(msg)
else:
resp = ConfigResponse(
value=None,
directory=None,
values=None,
error=Error(
type = "bad-operation",
message = "Bad operation"
)
)
return resp

324
trustgraph-flow/trustgraph/config/service/service.py
View file

@ -1,287 +1,118 @@
"""
Config service. Fetchs an extract from the Wikipedia page
using the API.
Config service. Manages system global configuration state
"""
from pulsar.schema import JsonSchema
from trustgraph.schema import ConfigRequest, ConfigResponse, ConfigPush
from trustgraph.schema import ConfigValue, Error
from trustgraph.schema import Error
from trustgraph.schema import config_request_queue, config_response_queue
from trustgraph.schema import config_push_queue
from trustgraph.log_level import LogLevel
from trustgraph.base import ConsumerProducer
from trustgraph.base import AsyncProcessor, Consumer, Producer
module = ".".join(__name__.split(".")[1:-1])
from . config import Configuration
from ... base import ProcessorMetrics, ConsumerMetrics, ProducerMetrics
from ... base import Consumer, Producer
default_input_queue = config_request_queue
default_output_queue = config_response_queue
default_ident = "config-svc"
default_request_queue = config_request_queue
default_response_queue = config_response_queue
default_push_queue = config_push_queue
default_subscriber = module
# This behaves just like a dict, should be easier to add persistent storage
# later
class ConfigurationItems(dict):
pass
class Configuration(dict):
def __getitem__(self, key):
if key not in self:
self[key] = ConfigurationItems()
return dict.__getitem__(self, key)
class Processor(ConsumerProducer):
class Processor(AsyncProcessor):
def __init__(self, **params):
input_queue = params.get("input_queue", default_input_queue)
output_queue = params.get("output_queue", default_output_queue)
request_queue = params.get("request_queue", default_request_queue)
response_queue = params.get("response_queue", default_response_queue)
push_queue = params.get("push_queue", default_push_queue)
subscriber = params.get("subscriber", default_subscriber)
id = params.get("id")
request_schema = ConfigRequest
response_schema = ConfigResponse
push_schema = ConfigResponse
super(Processor, self).__init__(
**params | {
"input_queue": input_queue,
"output_queue": output_queue,
"push_queue": output_queue,
"subscriber": subscriber,
"input_schema": ConfigRequest,
"output_schema": ConfigResponse,
"push_schema": ConfigPush,
"request_schema": request_schema.__name__,
"response_schema": response_schema.__name__,
"push_schema": push_schema.__name__,
}
)
self.push_prod = self.client.create_producer(
topic=push_queue,
schema=JsonSchema(ConfigPush),
request_metrics = ConsumerMetrics(id + "-request")
response_metrics = ProducerMetrics(id + "-response")
push_metrics = ProducerMetrics(id + "-push")
self.push_pub = Producer(
client = self.client,
topic = push_queue,
schema = ConfigPush,
metrics = push_metrics,
)
# FIXME: The state is held internally. This only works if there's
# one config service. Should be more than one, and use a
# back-end state store.
self.config = Configuration()
self.response_pub = Producer(
client = self.client,
topic = response_queue,
schema = ConfigResponse,
metrics = response_metrics,
)
# Version counter
self.version = 0
self.subs = Consumer(
taskgroup = self.taskgroup,
client = self.client,
flow = None,
topic = request_queue,
subscriber = id,
schema = request_schema,
handler = self.on_message,
metrics = request_metrics,
)
self.config = Configuration(self.push)
print("Service initialised.")
async def start(self):
await self.push()
await self.subs.start()
async def handle_get(self, v, id):
for k in v.keys:
if k.type not in self.config or k.key not in self.config[k.type]:
return ConfigResponse(
version = None,
values = None,
directory = None,
config = None,
error = Error(
type = "key-error",
message = f"Key error"
)
)
values = [
ConfigValue(
type = k.type,
key = k.key,
value = self.config[k.type][k.key]
)
for k in v.keys
]
return ConfigResponse(
version = self.version,
values = values,
directory = None,
config = None,
error = None,
)
async def handle_list(self, v, id):
if v.type not in self.config:
return ConfigResponse(
version = None,
values = None,
directory = None,
config = None,
error = Error(
type = "key-error",
message = "No such type",
),
)
return ConfigResponse(
version = self.version,
values = None,
directory = list(self.config[v.type].keys()),
config = None,
error = None,
)
async def handle_getvalues(self, v, id):
if v.type not in self.config:
return ConfigResponse(
version = None,
values = None,
directory = None,
config = None,
error = Error(
type = "key-error",
message = f"Key error"
)
)
values = [
ConfigValue(
type = v.type,
key = k,
value = self.config[v.type][k],
)
for k in self.config[v.type]
]
return ConfigResponse(
version = self.version,
values = values,
directory = None,
config = None,
error = None,
)
async def handle_delete(self, v, id):
for k in v.keys:
if k.type not in self.config or k.key not in self.config[k.type]:
return ConfigResponse(
version = None,
values = None,
directory = None,
config = None,
error = Error(
type = "key-error",
message = f"Key error"
)
)
for k in v.keys:
del self.config[k.type][k.key]
self.version += 1
await self.push()
return ConfigResponse(
version = None,
value = None,
directory = None,
values = None,
config = None,
error = None,
)
async def handle_put(self, v, id):
for k in v.values:
self.config[k.type][k.key] = k.value
self.version += 1
await self.push()
return ConfigResponse(
version = None,
value = None,
directory = None,
values = None,
error = None,
)
async def handle_config(self, v, id):
return ConfigResponse(
version = self.version,
value = None,
directory = None,
values = None,
config = self.config,
error = None,
)
async def push(self):
resp = ConfigPush(
version = self.version,
version = self.config.version,
value = None,
directory = None,
values = None,
config = self.config,
error = None,
)
self.push_prod.send(resp)
print("Pushed.")
await self.push_pub.send(resp)
print("Pushed version ", self.config.version)
async def handle(self, msg):
v = msg.value()
# Sender-produced ID
id = msg.properties()["id"]
print(f"Handling {id}...", flush=True)
async def on_message(self, msg, consumer, flow):
try:
if v.operation == "get":
v = msg.value()
resp = await self.handle_get(v, id)
# Sender-produced ID
id = msg.properties()["id"]
elif v.operation == "list":
print(f"Handling {id}...", flush=True)
resp = await self.handle_list(v, id)
resp = await self.config.handle(v)
elif v.operation == "getvalues":
resp = await self.handle_getvalues(v, id)
elif v.operation == "delete":
resp = await self.handle_delete(v, id)
elif v.operation == "put":
resp = await self.handle_put(v, id)
elif v.operation == "config":
resp = await self.handle_config(v, id)
else:
resp = ConfigResponse(
value=None,
directory=None,
values=None,
error=Error(
type = "bad-operation",
message = "Bad operation"
)
)
await self.send(resp, properties={"id": id})
self.consumer.acknowledge(msg)
await self.response_pub.send(resp, properties={"id": id})
except Exception as e:
resp = ConfigResponse(
error=Error(
type = "unexpected-error",
@ -289,24 +120,33 @@ class Processor(ConsumerProducer):
),
text=None,
)
await self.send(resp, properties={"id": id})
self.consumer.acknowledge(msg)
await self.response_pub.send(resp, properties={"id": id})
@staticmethod
def add_args(parser):
ConsumerProducer.add_args(
parser, default_input_queue, default_subscriber,
default_output_queue,
AsyncProcessor.add_args(parser)
parser.add_argument(
'-q', '--request-queue',
default=default_request_queue,
help=f'Request queue (default: {default_request_queue})'
)
parser.add_argument(
'-q', '--push-queue',
'-r', '--response-queue',
default=default_response_queue,
help=f'Response queue {default_response_queue}',
)
parser.add_argument(
'--push-queue',
default=default_push_queue,
help=f'Config push queue (default: {default_push_queue})'
)
def run():
Processor.launch(module, __doc__)
Processor.launch(default_ident, __doc__)

23
trustgraph-flow/trustgraph/decoding/mistral_ocr/processor.py
View file

@ -17,12 +17,10 @@ from mistralai.models import OCRResponse
from ... schema import Document, TextDocument, Metadata
from ... schema import document_ingest_queue, text_ingest_queue
from ... log_level import LogLevel
from ... base import ConsumerProducer
from ... base import InputOutputProcessor
module = ".".join(__name__.split(".")[1:-1])
module = "ocr"
default_input_queue = document_ingest_queue
default_output_queue = text_ingest_queue
default_subscriber = module
default_api_key = os.getenv("MISTRAL_TOKEN")
@ -71,19 +69,17 @@ def get_combined_markdown(ocr_response: OCRResponse) -> str:
return "\n\n".join(markdowns)
class Processor(ConsumerProducer):
class Processor(InputOutputProcessor):
def __init__(self, **params):
input_queue = params.get("input_queue", default_input_queue)
output_queue = params.get("output_queue", default_output_queue)
id = params.get("id")
subscriber = params.get("subscriber", default_subscriber)
api_key = params.get("api_key", default_api_key)
super(Processor, self).__init__(
**params | {
"input_queue": input_queue,
"output_queue": output_queue,
"id": id,
"subscriber": subscriber,
"input_schema": Document,
"output_schema": TextDocument,
@ -151,7 +147,7 @@ class Processor(ConsumerProducer):
return markdown
async def handle(self, msg):
async def on_message(self, msg, consumer):
print("PDF message received")
@ -166,17 +162,14 @@ class Processor(ConsumerProducer):
text=markdown.encode("utf-8"),
)
await self.send(r)
await consumer.q.output.send(r)
print("Done.", flush=True)
@staticmethod
def add_args(parser):
ConsumerProducer.add_args(
parser, default_input_queue, default_subscriber,
default_output_queue,
)
InputOutputProcessor.add_args(parser, default_subscriber)
parser.add_argument(
'-k', '--api-key',

54
trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py
View file

@ -9,39 +9,43 @@ import base64
from langchain_community.document_loaders import PyPDFLoader
from ... schema import Document, TextDocument, Metadata
from ... schema import document_ingest_queue, text_ingest_queue
from ... log_level import LogLevel
from ... base import ConsumerProducer
from ... base import FlowProcessor, ConsumerSpec, ProducerSpec
module = ".".join(__name__.split(".")[1:-1])
default_ident = "pdf-decoder"
default_input_queue = document_ingest_queue
default_output_queue = text_ingest_queue
default_subscriber = module
class Processor(ConsumerProducer):
class Processor(FlowProcessor):
def __init__(self, **params):
input_queue = params.get("input_queue", default_input_queue)
output_queue = params.get("output_queue", default_output_queue)
subscriber = params.get("subscriber", default_subscriber)
id = params.get("id", default_ident)
super(Processor, self).__init__(
**params | {
"input_queue": input_queue,
"output_queue": output_queue,
"subscriber": subscriber,
"input_schema": Document,
"output_schema": TextDocument,
"id": id,
}
)
print("PDF inited")
self.register_specification(
ConsumerSpec(
name = "input",
schema = Document,
handler = self.on_message,
)
)
async def handle(self, msg):
self.register_specification(
ProducerSpec(
name = "output",
schema = TextDocument,
)
)
print("PDF message received")
print("PDF inited", flush=True)
async def on_message(self, msg, consumer, flow):
print("PDF message received", flush=True)
v = msg.value()
@ -59,24 +63,22 @@ class Processor(ConsumerProducer):
for ix, page in enumerate(pages):
print("page", ix, flush=True)
r = TextDocument(
metadata=v.metadata,
text=page.page_content.encode("utf-8"),
)
await self.send(r)
await flow("output").send(r)
print("Done.", flush=True)
@staticmethod
def add_args(parser):
ConsumerProducer.add_args(
parser, default_input_queue, default_subscriber,
default_output_queue,
)
FlowProcessor.add_args(parser)
def run():
Processor.launch(module, __doc__)
Processor.launch(default_ident, __doc__)

153
trustgraph-flow/trustgraph/document_rag.py
View file

@ -1,153 +0,0 @@
from . clients.document_embeddings_client import DocumentEmbeddingsClient
from . clients.triples_query_client import TriplesQueryClient
from . clients.embeddings_client import EmbeddingsClient
from . clients.prompt_client import PromptClient
from . schema import DocumentEmbeddingsRequest, DocumentEmbeddingsResponse
from . schema import TriplesQueryRequest, TriplesQueryResponse
from . schema import prompt_request_queue
from . schema import prompt_response_queue
from . schema import embeddings_request_queue
from . schema import embeddings_response_queue
from . schema import document_embeddings_request_queue
from . schema import document_embeddings_response_queue
LABEL="http://www.w3.org/2000/01/rdf-schema#label"
DEFINITION="http://www.w3.org/2004/02/skos/core#definition"
class Query:
def __init__(
self, rag, user, collection, verbose,
doc_limit=20
):
self.rag = rag
self.user = user
self.collection = collection
self.verbose = verbose
self.doc_limit = doc_limit
def get_vector(self, query):
if self.verbose:
print("Compute embeddings...", flush=True)
qembeds = self.rag.embeddings.request(query)
if self.verbose:
print("Done.", flush=True)
return qembeds
def get_docs(self, query):
vectors = self.get_vector(query)
if self.verbose:
print("Get entities...", flush=True)
docs = self.rag.de_client.request(
vectors, limit=self.doc_limit
)
if self.verbose:
print("Docs:", flush=True)
for doc in docs:
print(doc, flush=True)
return docs
class DocumentRag:
def __init__(
self,
pulsar_host="pulsar://pulsar:6650",
pulsar_api_key=None,
pr_request_queue=None,
pr_response_queue=None,
emb_request_queue=None,
emb_response_queue=None,
de_request_queue=None,
de_response_queue=None,
verbose=False,
module="test",
):
self.verbose=verbose
if pr_request_queue is None:
pr_request_queue = prompt_request_queue
if pr_response_queue is None:
pr_response_queue = prompt_response_queue
if emb_request_queue is None:
emb_request_queue = embeddings_request_queue
if emb_response_queue is None:
emb_response_queue = embeddings_response_queue
if de_request_queue is None:
de_request_queue = document_embeddings_request_queue
if de_response_queue is None:
de_response_queue = document_embeddings_response_queue
if self.verbose:
print("Initialising...", flush=True)
self.de_client = DocumentEmbeddingsClient(
pulsar_host=pulsar_host,
subscriber=module + "-de",
input_queue=de_request_queue,
output_queue=de_response_queue,
pulsar_api_key=pulsar_api_key,
)
self.embeddings = EmbeddingsClient(
pulsar_host=pulsar_host,
input_queue=emb_request_queue,
output_queue=emb_response_queue,
subscriber=module + "-emb",
pulsar_api_key=pulsar_api_key,
)
self.lang = PromptClient(
pulsar_host=pulsar_host,
input_queue=pr_request_queue,
output_queue=pr_response_queue,
subscriber=module + "-de-prompt",
pulsar_api_key=pulsar_api_key,
)
if self.verbose:
print("Initialised", flush=True)
def query(
self, query, user="trustgraph", collection="default",
doc_limit=20,
):
if self.verbose:
print("Construct prompt...", flush=True)
q = Query(
rag=self, user=user, collection=collection, verbose=self.verbose,
doc_limit=doc_limit
)
docs = q.get_docs(query)
if self.verbose:
print("Invoke LLM...", flush=True)
print(docs)
print(query)
resp = self.lang.request_document_prompt(query, docs)
if self.verbose:
print("Done", flush=True)
return resp

93
trustgraph-flow/trustgraph/embeddings/document_embeddings/embeddings.py
View file

@ -6,61 +6,63 @@ Output is chunk plus embedding.
"""
from ... schema import Chunk, ChunkEmbeddings, DocumentEmbeddings
from ... schema import chunk_ingest_queue
from ... schema import document_embeddings_store_queue
from ... schema import embeddings_request_queue, embeddings_response_queue
from ... clients.embeddings_client import EmbeddingsClient
from ... log_level import LogLevel
from ... base import ConsumerProducer
from ... schema import EmbeddingsRequest, EmbeddingsResponse
module = ".".join(__name__.split(".")[1:-1])
from ... base import FlowProcessor, RequestResponseSpec, ConsumerSpec
from ... base import ProducerSpec
default_input_queue = chunk_ingest_queue
default_output_queue = document_embeddings_store_queue
default_subscriber = module
default_ident = "document-embeddings"
class Processor(ConsumerProducer):
class Processor(FlowProcessor):
def __init__(self, **params):
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
)
id = params.get("id")
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": DocumentEmbeddings,
"id": id,
}
)
self.embeddings = EmbeddingsClient(
pulsar_host=self.pulsar_host,
pulsar_api_key=self.pulsar_api_key,
input_queue=emb_request_queue,
output_queue=emb_response_queue,
subscriber=module + "-emb",
self.register_specification(
ConsumerSpec(
name = "input",
schema = Chunk,
handler = self.on_message,
)
)
async def handle(self, msg):
self.register_specification(
RequestResponseSpec(
request_name = "embeddings-request",
request_schema = EmbeddingsRequest,
response_name = "embeddings-response",
response_schema = EmbeddingsResponse,
)
)
self.register_specification(
ProducerSpec(
name = "output",
schema = DocumentEmbeddings
)
)
async def on_message(self, msg, consumer, flow):
v = msg.value()
print(f"Indexing {v.metadata.id}...", flush=True)
try:
vectors = self.embeddings.request(v.chunk)
resp = await flow("embeddings-request").request(
EmbeddingsRequest(
text = v.chunk
)
)
vectors = resp.vectors
embeds = [
ChunkEmbeddings(
@ -74,7 +76,7 @@ class Processor(ConsumerProducer):
chunks=embeds,
)
await self.send(r)
await flow("output").send(r)
except Exception as e:
print("Exception:", e, flush=True)
@ -87,24 +89,9 @@ class Processor(ConsumerProducer):
@staticmethod
def add_args(parser):
ConsumerProducer.add_args(
parser, default_input_queue, default_subscriber,
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})',
)
FlowProcessor.add_args(parser)
def run():
Processor.launch(module, __doc__)
Processor.launch(default_ident, __doc__)

60
trustgraph-flow/trustgraph/embeddings/fastembed/processor.py
View file

@ -1,81 +1,43 @@
"""
Embeddings service, applies an embeddings model selected from HuggingFace.
Embeddings service, applies an embeddings model using fastembed
Input is text, output is embeddings vector.
"""
from ... schema import EmbeddingsRequest, EmbeddingsResponse
from ... schema import embeddings_request_queue, embeddings_response_queue
from ... log_level import LogLevel
from ... base import ConsumerProducer
from ... base import EmbeddingsService
from fastembed import TextEmbedding
import os
module = ".".join(__name__.split(".")[1:-1])
default_ident = "embeddings"
default_input_queue = embeddings_request_queue
default_output_queue = embeddings_response_queue
default_subscriber = module
default_model="sentence-transformers/all-MiniLM-L6-v2"
class Processor(ConsumerProducer):
class Processor(EmbeddingsService):
def __init__(self, **params):
input_queue = params.get("input_queue", default_input_queue)
output_queue = params.get("output_queue", default_output_queue)
subscriber = params.get("subscriber", default_subscriber)
model = params.get("model", default_model)
super(Processor, self).__init__(
**params | {
"input_queue": input_queue,
"output_queue": output_queue,
"subscriber": subscriber,
"input_schema": EmbeddingsRequest,
"output_schema": EmbeddingsResponse,
"model": model,
}
**params | { "model": model }
)
print("Get model...", flush=True)
self.embeddings = TextEmbedding(model_name = model)
async def handle(self, msg):
async def on_embeddings(self, text):
v = msg.value()
# Sender-produced ID
id = msg.properties()["id"]
print(f"Handling input {id}...", flush=True)
text = v.text
vecs = self.embeddings.embed([text])
vecs = [
return [
v.tolist()
for v in vecs
]
print("Send response...", flush=True)
r = EmbeddingsResponse(
vectors=list(vecs),
error=None,
)
await self.send(r, properties={"id": id})
print("Done.", flush=True)
@staticmethod
def add_args(parser):
ConsumerProducer.add_args(
parser, default_input_queue, default_subscriber,
default_output_queue,
)
EmbeddingsService.add_args(parser)
parser.add_argument(
'-m', '--model',
@ -85,5 +47,5 @@ class Processor(ConsumerProducer):
def run():
Processor.launch(module, __doc__)
Processor.launch(default_ident, __doc__)

86
trustgraph-flow/trustgraph/embeddings/graph_embeddings/embeddings.py
View file

@ -6,53 +6,48 @@ Output is entity plus embedding.
"""
from ... schema import EntityContexts, EntityEmbeddings, GraphEmbeddings
from ... schema import entity_contexts_ingest_queue
from ... schema import graph_embeddings_store_queue
from ... schema import embeddings_request_queue, embeddings_response_queue
from ... clients.embeddings_client import EmbeddingsClient
from ... log_level import LogLevel
from ... base import ConsumerProducer
from ... schema import EmbeddingsRequest, EmbeddingsResponse
module = ".".join(__name__.split(".")[1:-1])
from ... base import FlowProcessor, EmbeddingsClientSpec, ConsumerSpec
from ... base import ProducerSpec
default_input_queue = entity_contexts_ingest_queue
default_output_queue = graph_embeddings_store_queue
default_subscriber = module
default_ident = "graph-embeddings"
class Processor(ConsumerProducer):
class Processor(FlowProcessor):
def __init__(self, **params):
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
)
id = params.get("id")
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": EntityContexts,
"output_schema": GraphEmbeddings,
"id": id,
}
)
self.embeddings = EmbeddingsClient(
pulsar_host=self.pulsar_host,
input_queue=emb_request_queue,
output_queue=emb_response_queue,
subscriber=module + "-emb",
self.register_specification(
ConsumerSpec(
name = "input",
schema = EntityContexts,
handler = self.on_message,
)
)
async def handle(self, msg):
self.register_specification(
EmbeddingsClientSpec(
request_name = "embeddings-request",
response_name = "embeddings-response",
)
)
self.register_specification(
ProducerSpec(
name = "output",
schema = GraphEmbeddings
)
)
async def on_message(self, msg, consumer, flow):
v = msg.value()
print(f"Indexing {v.metadata.id}...", flush=True)
@ -63,7 +58,9 @@ class Processor(ConsumerProducer):
for entity in v.entities:
vectors = self.embeddings.request(entity.context)
vectors = await flow("embeddings-request").embed(
text = entity.context
)
entities.append(
EntityEmbeddings(
@ -77,7 +74,7 @@ class Processor(ConsumerProducer):
entities=entities,
)
await self.send(r)
await flow("output").send(r)
except Exception as e:
print("Exception:", e, flush=True)
@ -90,24 +87,9 @@ class Processor(ConsumerProducer):
@staticmethod
def add_args(parser):
ConsumerProducer.add_args(
parser, default_input_queue, default_subscriber,
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})',
)
FlowProcessor.add_args(parser)
def run():
Processor.launch(module, __doc__)
Processor.launch(default_ident, __doc__)

2
trustgraph-flow/trustgraph/embeddings/ollama/processor.py
View file

@ -11,7 +11,7 @@ from ... base import ConsumerProducer
from ollama import Client
import os
module = ".".join(__name__.split(".")[1:-1])
module = "embeddings"
default_input_queue = embeddings_request_queue
default_output_queue = embeddings_response_queue

2
trustgraph-flow/trustgraph/external/wikipedia/service.py vendored
View file

@ -11,7 +11,7 @@ from trustgraph.log_level import LogLevel
from trustgraph.base import ConsumerProducer
import requests
module = ".".join(__name__.split(".")[1:-1])
module = "wikipedia"
default_input_queue = encyclopedia_lookup_request_queue
default_output_queue = encyclopedia_lookup_response_queue

155
trustgraph-flow/trustgraph/extract/kg/definitions/extract.py
View file

@ -5,84 +5,62 @@ get entity definitions which are output as graph edges along with
entity/context definitions for embedding.
"""
import json
import urllib.parse
from pulsar.schema import JsonSchema
from .... schema import Chunk, Triple, Triples, Metadata, Value
from .... schema import EntityContext, EntityContexts
from .... schema import chunk_ingest_queue, triples_store_queue
from .... schema import entity_contexts_ingest_queue
from .... schema import prompt_request_queue
from .... schema import prompt_response_queue
from .... log_level import LogLevel
from .... clients.prompt_client import PromptClient
from .... schema import PromptRequest, PromptResponse
from .... rdf import TRUSTGRAPH_ENTITIES, DEFINITION, RDF_LABEL, SUBJECT_OF
from .... base import ConsumerProducer
from .... base import FlowProcessor, ConsumerSpec, ProducerSpec
from .... base import PromptClientSpec
DEFINITION_VALUE = Value(value=DEFINITION, is_uri=True)
RDF_LABEL_VALUE = Value(value=RDF_LABEL, is_uri=True)
SUBJECT_OF_VALUE = Value(value=SUBJECT_OF, is_uri=True)
module = ".".join(__name__.split(".")[1:-1])
default_ident = "kg-extract-definitions"
default_input_queue = chunk_ingest_queue
default_output_queue = triples_store_queue
default_entity_context_queue = entity_contexts_ingest_queue
default_subscriber = module
class Processor(ConsumerProducer):
class Processor(FlowProcessor):
def __init__(self, **params):
input_queue = params.get("input_queue", default_input_queue)
output_queue = params.get("output_queue", default_output_queue)
ec_queue = params.get(
"entity_context_queue",
default_entity_context_queue
)
subscriber = params.get("subscriber", default_subscriber)
pr_request_queue = params.get(
"prompt_request_queue", prompt_request_queue
)
pr_response_queue = params.get(
"prompt_response_queue", prompt_response_queue
)
id = params.get("id")
super(Processor, self).__init__(
**params | {
"input_queue": input_queue,
"output_queue": output_queue,
"subscriber": subscriber,
"input_schema": Chunk,
"output_schema": Triples,
"prompt_request_queue": pr_request_queue,
"prompt_response_queue": pr_response_queue,
"id": id,
}
)
self.ec_prod = self.client.create_producer(
topic=ec_queue,
schema=JsonSchema(EntityContexts),
self.register_specification(
ConsumerSpec(
name = "input",
schema = Chunk,
handler = self.on_message
)
)
__class__.pubsub_metric.info({
"input_queue": input_queue,
"output_queue": output_queue,
"entity_context_queue": ec_queue,
"prompt_request_queue": pr_request_queue,
"prompt_response_queue": pr_response_queue,
"subscriber": subscriber,
"input_schema": Chunk.__name__,
"output_schema": Triples.__name__,
"vector_schema": EntityContexts.__name__,
})
self.register_specification(
PromptClientSpec(
request_name = "prompt-request",
response_name = "prompt-response",
)
)
self.prompt = PromptClient(
pulsar_host=self.pulsar_host,
pulsar_api_key=self.pulsar_api_key,
input_queue=pr_request_queue,
output_queue=pr_response_queue,
subscriber = module + "-prompt",
self.register_specification(
ProducerSpec(
name = "triples",
schema = Triples
)
)
self.register_specification(
ProducerSpec(
name = "entity-contexts",
schema = EntityContexts
)
)
def to_uri(self, text):
@ -93,36 +71,47 @@ class Processor(ConsumerProducer):
return uri
def get_definitions(self, chunk):
return self.prompt.request_definitions(chunk)
async def emit_edges(self, metadata, triples):
async def emit_triples(self, pub, metadata, triples):
t = Triples(
metadata=metadata,
triples=triples,
)
await self.send(t)
await pub.send(t)
async def emit_ecs(self, metadata, entities):
async def emit_ecs(self, pub, metadata, entities):
t = EntityContexts(
metadata=metadata,
entities=entities,
)
self.ec_prod.send(t)
await pub.send(t)
async def handle(self, msg):
async def on_message(self, msg, consumer, flow):
v = msg.value()
print(f"Indexing {v.metadata.id}...", flush=True)
chunk = v.chunk.decode("utf-8")
print(chunk, flush=True)
try:
defs = self.get_definitions(chunk)
try:
defs = await flow("prompt-request").extract_definitions(
text = chunk
)
print("Response", defs, flush=True)
if type(defs) != list:
raise RuntimeError("Expecting array in prompt response")
except Exception as e:
print("Prompt exception:", e, flush=True)
raise e
triples = []
entities = []
@ -134,8 +123,8 @@ class Processor(ConsumerProducer):
for defn in defs:
s = defn.name
o = defn.definition
s = defn["entity"]
o = defn["definition"]
if s == "": continue
if o == "": continue
@ -166,13 +155,13 @@ class Processor(ConsumerProducer):
ec = EntityContext(
entity=s_value,
context=defn.definition,
context=defn["definition"],
)
entities.append(ec)
await self.emit_edges(
await self.emit_triples(
flow("triples"),
Metadata(
id=v.metadata.id,
metadata=[],
@ -183,6 +172,7 @@ class Processor(ConsumerProducer):
)
await self.emit_ecs(
flow("entity-contexts"),
Metadata(
id=v.metadata.id,
metadata=[],
@ -200,30 +190,9 @@ class Processor(ConsumerProducer):
@staticmethod
def add_args(parser):
ConsumerProducer.add_args(
parser, default_input_queue, default_subscriber,
default_output_queue,
)
parser.add_argument(
'-e', '--entity-context-queue',
default=default_entity_context_queue,
help=f'Entity context queue (default: {default_entity_context_queue})'
)
parser.add_argument(
'--prompt-request-queue',
default=prompt_request_queue,
help=f'Prompt request queue (default: {prompt_request_queue})',
)
parser.add_argument(
'--prompt-completion-response-queue',
default=prompt_response_queue,
help=f'Prompt response queue (default: {prompt_response_queue})',
)
FlowProcessor.add_args(parser)
def run():
Processor.launch(module, __doc__)
Processor.launch(default_ident, __doc__)

122
trustgraph-flow/trustgraph/extract/kg/relationships/extract.py
View file

@ -5,59 +5,54 @@ relationship analysis to get entity relationship edges which are output as
graph edges.
"""
import json
import urllib.parse
from .... schema import Chunk, Triple, Triples
from .... schema import Metadata, Value
from .... schema import chunk_ingest_queue, triples_store_queue
from .... schema import prompt_request_queue
from .... schema import prompt_response_queue
from .... log_level import LogLevel
from .... clients.prompt_client import PromptClient
from .... schema import PromptRequest, PromptResponse
from .... rdf import RDF_LABEL, TRUSTGRAPH_ENTITIES, SUBJECT_OF
from .... base import ConsumerProducer
from .... base import FlowProcessor, ConsumerSpec, ProducerSpec
from .... base import PromptClientSpec
RDF_LABEL_VALUE = Value(value=RDF_LABEL, is_uri=True)
SUBJECT_OF_VALUE = Value(value=SUBJECT_OF, is_uri=True)
module = ".".join(__name__.split(".")[1:-1])
default_ident = "kg-extract-relationships"
default_input_queue = chunk_ingest_queue
default_output_queue = triples_store_queue
default_subscriber = module
class Processor(ConsumerProducer):
class Processor(FlowProcessor):
def __init__(self, **params):
input_queue = params.get("input_queue", default_input_queue)
output_queue = params.get("output_queue", default_output_queue)
subscriber = params.get("subscriber", default_subscriber)
pr_request_queue = params.get(
"prompt_request_queue", prompt_request_queue
)
pr_response_queue = params.get(
"prompt_response_queue", prompt_response_queue
)
id = params.get("id")
super(Processor, self).__init__(
**params | {
"input_queue": input_queue,
"output_queue": output_queue,
"subscriber": subscriber,
"input_schema": Chunk,
"output_schema": Triples,
"prompt_request_queue": pr_request_queue,
"prompt_response_queue": pr_response_queue,
"id": id,
}
)
self.prompt = PromptClient(
pulsar_host=self.pulsar_host,
pulsar_api_key=self.pulsar_api_key,
input_queue=pr_request_queue,
output_queue=pr_response_queue,
subscriber = module + "-prompt",
self.register_specification(
ConsumerSpec(
name = "input",
schema = Chunk,
handler = self.on_message
)
)
self.register_specification(
PromptClientSpec(
request_name = "prompt-request",
response_name = "prompt-response",
)
)
self.register_specification(
ProducerSpec(
name = "triples",
schema = Triples
)
)
def to_uri(self, text):
@ -68,28 +63,39 @@ class Processor(ConsumerProducer):
return uri
def get_relationships(self, chunk):
return self.prompt.request_relationships(chunk)
async def emit_edges(self, metadata, triples):
async def emit_triples(self, pub, metadata, triples):
t = Triples(
metadata=metadata,
triples=triples,
)
await self.send(t)
await pub.send(t)
async def handle(self, msg):
async def on_message(self, msg, consumer, flow):
v = msg.value()
print(f"Indexing {v.metadata.id}...", flush=True)
chunk = v.chunk.decode("utf-8")
print(chunk, flush=True)
try:
rels = self.get_relationships(chunk)
try:
rels = await flow("prompt-request").extract_relationships(
text = chunk
)
print("Response", rels, flush=True)
if type(rels) != list:
raise RuntimeError("Expecting array in prompt response")
except Exception as e:
print("Prompt exception:", e, flush=True)
raise e
triples = []
@ -100,9 +106,9 @@ class Processor(ConsumerProducer):
for rel in rels:
s = rel.s
p = rel.p
o = rel.o
s = rel["subject"]
p = rel["predicate"]
o = rel["object"]
if s == "": continue
if p == "": continue
@ -118,7 +124,7 @@ class Processor(ConsumerProducer):
p_uri = self.to_uri(p)
p_value = Value(value=str(p_uri), is_uri=True)
if rel.o_entity:
if rel["object-entity"]:
o_uri = self.to_uri(o)
o_value = Value(value=str(o_uri), is_uri=True)
else:
@ -144,7 +150,7 @@ class Processor(ConsumerProducer):
o=Value(value=str(p), is_uri=False)
))
if rel.o_entity:
if rel["object-entity"]:
# Label for o
triples.append(Triple(
s=o_value,
@ -159,7 +165,7 @@ class Processor(ConsumerProducer):
o=Value(value=v.metadata.id, is_uri=True)
))
if rel.o_entity:
if rel["object-entity"]:
# 'Subject of' for o
triples.append(Triple(
s=o_value,
@ -168,6 +174,7 @@ class Processor(ConsumerProducer):
))
await self.emit_edges(
flow("triples"),
Metadata(
id=v.metadata.id,
metadata=[],
@ -185,24 +192,9 @@ class Processor(ConsumerProducer):
@staticmethod
def add_args(parser):
ConsumerProducer.add_args(
parser, default_input_queue, default_subscriber,
default_output_queue,
)
parser.add_argument(
'--prompt-request-queue',
default=prompt_request_queue,
help=f'Prompt request queue (default: {prompt_request_queue})',
)
parser.add_argument(
'--prompt-response-queue',
default=prompt_response_queue,
help=f'Prompt response queue (default: {prompt_response_queue})',
)
FlowProcessor.add_args(parser)
def run():
Processor.launch(module, __doc__)
Processor.launch(default_ident, __doc__)

2
trustgraph-flow/trustgraph/extract/kg/topics/extract.py
View file

@ -18,7 +18,7 @@ from .... base import ConsumerProducer
DEFINITION_VALUE = Value(value=DEFINITION, is_uri=True)
module = ".".join(__name__.split(".")[1:-1])
module = "kg-extract-topics"
default_input_queue = chunk_ingest_queue
default_output_queue = triples_store_queue

1
trustgraph-flow/trustgraph/gateway/agent.py
View file

@ -39,4 +39,3 @@ class AgentRequestor(ServiceRequestor):
# The 2nd boolean expression indicates whether we're done responding
return resp, (message.answer is not None)

7
trustgraph-flow/trustgraph/gateway/document_embeddings_load.py
View file

@ -1,6 +1,5 @@
import asyncio
from pulsar.schema import JsonSchema
import uuid
from aiohttp import WSMsgType
@ -26,12 +25,12 @@ class DocumentEmbeddingsLoadEndpoint(SocketEndpoint):
self.publisher = Publisher(
self.pulsar_client, document_embeddings_store_queue,
schema=JsonSchema(DocumentEmbeddings)
schema=DocumentEmbeddings
)
async def start(self):
self.publisher.start()
await self.publisher.start()
async def listener(self, ws, running):
@ -59,6 +58,6 @@ class DocumentEmbeddingsLoadEndpoint(SocketEndpoint):
],
)
self.publisher.send(None, elt)
await self.publisher.send(None, elt)
running.stop()

11
trustgraph-flow/trustgraph/gateway/document_embeddings_stream.py
View file

@ -1,7 +1,6 @@
import asyncio
import queue
from pulsar.schema import JsonSchema
import uuid
from .. schema import DocumentEmbeddings
@ -27,7 +26,7 @@ class DocumentEmbeddingsStreamEndpoint(SocketEndpoint):
self.subscriber = Subscriber(
self.pulsar_client, document_embeddings_store_queue,
"api-gateway", "api-gateway",
schema=JsonSchema(DocumentEmbeddings),
schema=DocumentEmbeddings,
)
async def listener(self, ws, running):
@ -44,17 +43,17 @@ class DocumentEmbeddingsStreamEndpoint(SocketEndpoint):
async def start(self):
self.subscriber.start()
await self.subscriber.start()
async def async_thread(self, ws, running):
id = str(uuid.uuid4())
q = self.subscriber.subscribe_all(id)
q = await self.subscriber.subscribe_all(id)
while running.get():
try:
resp = await asyncio.to_thread(q.get, timeout=0.5)
resp = await asyncio.wait_for(q.get(), timeout=0.5)
await ws.send_json(serialize_document_embeddings(resp))
except TimeoutError:
@ -67,7 +66,7 @@ class DocumentEmbeddingsStreamEndpoint(SocketEndpoint):
print(f"Exception: {str(e)}", flush=True)
break
self.subscriber.unsubscribe_all(id)
await self.subscriber.unsubscribe_all(id)
running.stop()

4
trustgraph-flow/trustgraph/gateway/endpoint.py
View file

@ -1,13 +1,9 @@
import asyncio
from pulsar.schema import JsonSchema
from aiohttp import web
import uuid
import logging
from .. base import Publisher
from .. base import Subscriber
logger = logging.getLogger("endpoint")
logger.setLevel(logging.INFO)

7
trustgraph-flow/trustgraph/gateway/graph_embeddings_load.py
View file

@ -1,6 +1,5 @@
import asyncio
from pulsar.schema import JsonSchema
import uuid
from aiohttp import WSMsgType
@ -26,12 +25,12 @@ class GraphEmbeddingsLoadEndpoint(SocketEndpoint):
self.publisher = Publisher(
self.pulsar_client, graph_embeddings_store_queue,
schema=JsonSchema(GraphEmbeddings)
schema=GraphEmbeddings
)
async def start(self):
self.publisher.start()
await self.publisher.start()
async def listener(self, ws, running):
@ -60,6 +59,6 @@ class GraphEmbeddingsLoadEndpoint(SocketEndpoint):
]
)
self.publisher.send(None, elt)
await self.publisher.send(None, elt)
running.stop()

11
trustgraph-flow/trustgraph/gateway/graph_embeddings_stream.py
View file

@ -1,7 +1,6 @@
import asyncio
import queue
from pulsar.schema import JsonSchema
import uuid
from .. schema import GraphEmbeddings
@ -26,7 +25,7 @@ class GraphEmbeddingsStreamEndpoint(SocketEndpoint):
self.subscriber = Subscriber(
self.pulsar_client, graph_embeddings_store_queue,
"api-gateway", "api-gateway",
schema=JsonSchema(GraphEmbeddings)
schema=GraphEmbeddings
)
async def listener(self, ws, running):
@ -41,17 +40,17 @@ class GraphEmbeddingsStreamEndpoint(SocketEndpoint):
async def start(self):
self.subscriber.start()
await self.subscriber.start()
async def async_thread(self, ws, running):
id = str(uuid.uuid4())
q = self.subscriber.subscribe_all(id)
q = await self.subscriber.subscribe_all(id)
while running.get():
try:
resp = await asyncio.to_thread(q.get, timeout=0.5)
resp = await asyncio.wait_for(q.get, timeout=0.5)
await ws.send_json(serialize_graph_embeddings(resp))
except TimeoutError:
@ -64,7 +63,7 @@ class GraphEmbeddingsStreamEndpoint(SocketEndpoint):
print(f"Exception: {str(e)}", flush=True)
break
self.subscriber.unsubscribe_all(id)
await self.subscriber.unsubscribe_all(id)
running.stop()

1
trustgraph-flow/trustgraph/gateway/metrics.py
View file

@ -7,7 +7,6 @@
import aiohttp
from aiohttp import web
import asyncio
from pulsar.schema import JsonSchema
import uuid
import logging

1
trustgraph-flow/trustgraph/gateway/mux.py
View file

@ -1,7 +1,6 @@
import asyncio
import queue
from pulsar.schema import JsonSchema
import uuid
from aiohttp import web, WSMsgType

22
trustgraph-flow/trustgraph/gateway/requestor.py
View file

@ -1,6 +1,5 @@
import asyncio
from pulsar.schema import JsonSchema
import uuid
import logging
@ -23,21 +22,21 @@ class ServiceRequestor:
self.pub = Publisher(
pulsar_client, request_queue,
schema=JsonSchema(request_schema),
schema=request_schema,
)
self.sub = Subscriber(
pulsar_client, response_queue,
subscription, consumer_name,
JsonSchema(response_schema)
response_schema
)
self.timeout = timeout
async def start(self):
self.pub.start()
self.sub.start()
await self.pub.start()
await self.sub.start()
def to_request(self, request):
raise RuntimeError("Not defined")
@ -51,18 +50,15 @@ class ServiceRequestor:
try:
q = self.sub.subscribe(id)
q = await self.sub.subscribe(id)
await asyncio.to_thread(
self.pub.send, id, self.to_request(request)
)
await self.pub.send(id, self.to_request(request))
while True:
try:
resp = await asyncio.to_thread(
q.get,
timeout=self.timeout
resp = await asyncio.wait_for(
q.get(), timeout=self.timeout
)
except Exception as e:
raise RuntimeError("Timeout")
@ -99,5 +95,5 @@ class ServiceRequestor:
return err
finally:
self.sub.unsubscribe(id)
await self.sub.unsubscribe(id)

9
trustgraph-flow/trustgraph/gateway/sender.py
View file

@ -2,7 +2,6 @@
# Like ServiceRequestor, but just fire-and-forget instead of request/response
import asyncio
from pulsar.schema import JsonSchema
import uuid
import logging
@ -21,12 +20,12 @@ class ServiceSender:
self.pub = Publisher(
pulsar_client, request_queue,
schema=JsonSchema(request_schema),
schema=request_schema,
)
async def start(self):
self.pub.start()
await self.pub.start()
def to_request(self, request):
raise RuntimeError("Not defined")
@ -35,9 +34,7 @@ class ServiceSender:
try:
await asyncio.to_thread(
self.pub.send, None, self.to_request(request)
)
await self.pub.send(None, self.to_request(request))
if responder:
await responder({}, True)

3
trustgraph-flow/trustgraph/gateway/service.py
View file

@ -3,7 +3,7 @@ API gateway. Offers HTTP services which are translated to interaction on the
Pulsar bus.
"""
module = ".".join(__name__.split(".")[1:-1])
module = "api-gateway"
# FIXME: Subscribes to Pulsar unnecessarily, should only do it when there
# are active listeners
@ -19,7 +19,6 @@ import os
import base64
import pulsar
from pulsar.schema import JsonSchema
from prometheus_client import start_http_server
from .. log_level import LogLevel

7
trustgraph-flow/trustgraph/gateway/triples_load.py
View file

@ -1,6 +1,5 @@
import asyncio
from pulsar.schema import JsonSchema
import uuid
from aiohttp import WSMsgType
@ -24,12 +23,12 @@ class TriplesLoadEndpoint(SocketEndpoint):
self.publisher = Publisher(
self.pulsar_client, triples_store_queue,
schema=JsonSchema(Triples)
schema=Triples
)
async def start(self):
self.publisher.start()
await self.publisher.start()
async def listener(self, ws, running):
@ -51,7 +50,7 @@ class TriplesLoadEndpoint(SocketEndpoint):
triples=to_subgraph(data["triples"]),
)
self.publisher.send(None, elt)
await self.publisher.send(None, elt)
running.stop()

5
trustgraph-flow/trustgraph/gateway/triples_stream.py
View file

@ -1,7 +1,6 @@
import asyncio
import queue
from pulsar.schema import JsonSchema
import uuid
from .. schema import Triples
@ -24,7 +23,7 @@ class TriplesStreamEndpoint(SocketEndpoint):
self.subscriber = Subscriber(
self.pulsar_client, triples_store_queue,
"api-gateway", "api-gateway",
schema=JsonSchema(Triples)
schema=Triples
)
async def listener(self, ws, running):
@ -39,7 +38,7 @@ class TriplesStreamEndpoint(SocketEndpoint):
async def start(self):
self.subscriber.start()
await self.subscriber.start()
async def async_thread(self, ws, running):

295
trustgraph-flow/trustgraph/graph_rag.py
View file

@ -1,295 +0,0 @@
from . clients.graph_embeddings_client import GraphEmbeddingsClient
from . clients.triples_query_client import TriplesQueryClient
from . clients.embeddings_client import EmbeddingsClient
from . clients.prompt_client import PromptClient
from . schema import GraphEmbeddingsRequest, GraphEmbeddingsResponse
from . schema import TriplesQueryRequest, TriplesQueryResponse
from . schema import prompt_request_queue
from . schema import prompt_response_queue
from . schema import embeddings_request_queue
from . schema import embeddings_response_queue
from . schema import graph_embeddings_request_queue
from . schema import graph_embeddings_response_queue
from . schema import triples_request_queue
from . schema import triples_response_queue
LABEL="http://www.w3.org/2000/01/rdf-schema#label"
DEFINITION="http://www.w3.org/2004/02/skos/core#definition"
class Query:
def __init__(
self, rag, user, collection, verbose,
entity_limit=50, triple_limit=30, max_subgraph_size=1000,
max_path_length=2,
):
self.rag = rag
self.user = user
self.collection = collection
self.verbose = verbose
self.entity_limit = entity_limit
self.triple_limit = triple_limit
self.max_subgraph_size = max_subgraph_size
self.max_path_length = max_path_length
def get_vector(self, query):
if self.verbose:
print("Compute embeddings...", flush=True)
qembeds = self.rag.embeddings.request(query)
if self.verbose:
print("Done.", flush=True)
return qembeds
def get_entities(self, query):
vectors = self.get_vector(query)
if self.verbose:
print("Get entities...", flush=True)
entities = self.rag.ge_client.request(
user=self.user, collection=self.collection,
vectors=vectors, limit=self.entity_limit,
)
entities = [
e.value
for e in entities
]
if self.verbose:
print("Entities:", flush=True)
for ent in entities:
print(" ", ent, flush=True)
return entities
def maybe_label(self, e):
if e in self.rag.label_cache:
return self.rag.label_cache[e]
res = self.rag.triples_client.request(
user=self.user, collection=self.collection,
s=e, p=LABEL, o=None, limit=1,
)
if len(res) == 0:
self.rag.label_cache[e] = e
return e
self.rag.label_cache[e] = res[0].o.value
return self.rag.label_cache[e]
def follow_edges(self, ent, subgraph, path_length):
# Not needed?
if path_length <= 0:
return
# Stop spanning around if the subgraph is already maxed out
if len(subgraph) >= self.max_subgraph_size:
return
res = self.rag.triples_client.request(
user=self.user, collection=self.collection,
s=ent, p=None, o=None,
limit=self.triple_limit
)
for triple in res:
subgraph.add(
(triple.s.value, triple.p.value, triple.o.value)
)
if path_length > 1:
self.follow_edges(triple.o.value, subgraph, path_length-1)
res = self.rag.triples_client.request(
user=self.user, collection=self.collection,
s=None, p=ent, o=None,
limit=self.triple_limit
)
for triple in res:
subgraph.add(
(triple.s.value, triple.p.value, triple.o.value)
)
res = self.rag.triples_client.request(
user=self.user, collection=self.collection,
s=None, p=None, o=ent,
limit=self.triple_limit,
)
for triple in res:
subgraph.add(
(triple.s.value, triple.p.value, triple.o.value)
)
if path_length > 1:
self.follow_edges(triple.s.value, subgraph, path_length-1)
def get_subgraph(self, query):
entities = self.get_entities(query)
if self.verbose:
print("Get subgraph...", flush=True)
subgraph = set()
for ent in entities:
self.follow_edges(ent, subgraph, self.max_path_length)
subgraph = list(subgraph)
return subgraph
def get_labelgraph(self, query):
subgraph = self.get_subgraph(query)
sg2 = []
for edge in subgraph:
if edge[1] == LABEL:
continue
s = self.maybe_label(edge[0])
p = self.maybe_label(edge[1])
o = self.maybe_label(edge[2])
sg2.append((s, p, o))
sg2 = sg2[0:self.max_subgraph_size]
if self.verbose:
print("Subgraph:", flush=True)
for edge in sg2:
print(" ", str(edge), flush=True)
if self.verbose:
print("Done.", flush=True)
return sg2
class GraphRag:
def __init__(
self,
pulsar_host="pulsar://pulsar:6650",
pulsar_api_key=None,
pr_request_queue=None,
pr_response_queue=None,
emb_request_queue=None,
emb_response_queue=None,
ge_request_queue=None,
ge_response_queue=None,
tpl_request_queue=None,
tpl_response_queue=None,
verbose=False,
module="test",
):
self.verbose=verbose
if pr_request_queue is None:
pr_request_queue = prompt_request_queue
if pr_response_queue is None:
pr_response_queue = prompt_response_queue
if emb_request_queue is None:
emb_request_queue = embeddings_request_queue
if emb_response_queue is None:
emb_response_queue = embeddings_response_queue
if ge_request_queue is None:
ge_request_queue = graph_embeddings_request_queue
if ge_response_queue is None:
ge_response_queue = graph_embeddings_response_queue
if tpl_request_queue is None:
tpl_request_queue = triples_request_queue
if tpl_response_queue is None:
tpl_response_queue = triples_response_queue
if self.verbose:
print("Initialising...", flush=True)
self.ge_client = GraphEmbeddingsClient(
pulsar_host=pulsar_host,
pulsar_api_key=pulsar_api_key,
subscriber=module + "-ge",
input_queue=ge_request_queue,
output_queue=ge_response_queue,
)
self.triples_client = TriplesQueryClient(
pulsar_host=pulsar_host,
pulsar_api_key=pulsar_api_key,
subscriber=module + "-tpl",
input_queue=tpl_request_queue,
output_queue=tpl_response_queue
)
self.embeddings = EmbeddingsClient(
pulsar_host=pulsar_host,
pulsar_api_key=pulsar_api_key,
input_queue=emb_request_queue,
output_queue=emb_response_queue,
subscriber=module + "-emb",
)
self.label_cache = {}
self.prompt = PromptClient(
pulsar_host=pulsar_host,
pulsar_api_key=pulsar_api_key,
input_queue=pr_request_queue,
output_queue=pr_response_queue,
subscriber=module + "-prompt",
)
if self.verbose:
print("Initialised", flush=True)
def query(
self, query, user="trustgraph", collection="default",
entity_limit=50, triple_limit=30, max_subgraph_size=1000,
max_path_length=2,
):
if self.verbose:
print("Construct prompt...", flush=True)
q = Query(
rag=self, user=user, collection=collection, verbose=self.verbose,
entity_limit=entity_limit, triple_limit=triple_limit,
max_subgraph_size=max_subgraph_size,
max_path_length=max_path_length,
)
kg = q.get_labelgraph(query)
if self.verbose:
print("Invoke LLM...", flush=True)
print(kg)
print(query)
resp = self.prompt.request_kg_prompt(query, kg)
if self.verbose:
print("Done", flush=True)
return resp

2
trustgraph-flow/trustgraph/librarian/service.py
View file

@ -35,7 +35,7 @@ from .. exceptions import RequestError
from . librarian import Librarian
module = ".".join(__name__.split(".")[1:-1])
module = "librarian"
default_input_queue = librarian_request_queue
default_output_queue = librarian_response_queue

3
trustgraph-flow/trustgraph/metering/counter.py
View file

@ -10,12 +10,11 @@ from .. schema import text_completion_response_queue
from .. log_level import LogLevel
from .. base import Consumer
module = ".".join(__name__.split(".")[1:-1])
module = "metering"
default_input_queue = text_completion_response_queue
default_subscriber = module
class Processor(Consumer):
def __init__(self, **params):

2
trustgraph-flow/trustgraph/model/prompt/generic/service.py
View file

@ -27,7 +27,7 @@ from .... clients.llm_client import LlmClient
from . prompts import to_definitions, to_relationships, to_topics
from . prompts import to_kg_query, to_document_query, to_rows
module = ".".join(__name__.split(".")[1:-1])
module = "prompt"
default_input_queue = prompt_request_queue
default_output_queue = prompt_response_queue

17
trustgraph-flow/trustgraph/model/prompt/template/prompt_manager.py
View file

@ -4,8 +4,6 @@ import json
from jsonschema import validate
import re
from trustgraph.clients.llm_client import LlmClient
class PromptConfiguration:
def __init__(self, system_template, global_terms={}, prompts={}):
self.system_template = system_template
@ -21,8 +19,7 @@ class Prompt:
class PromptManager:
def __init__(self, llm, config):
self.llm = llm
def __init__(self, config):
self.config = config
self.terms = config.global_terms
@ -54,7 +51,9 @@ class PromptManager:
return json.loads(json_str)
def invoke(self, id, input):
async def invoke(self, id, input, llm):
print("Invoke...", flush=True)
if id not in self.prompts:
raise RuntimeError("ID invalid")
@ -68,9 +67,7 @@ class PromptManager:
"prompt": self.templates[id].render(terms)
}
resp = self.llm.request(**prompt)
print(resp, flush=True)
resp = await llm(**prompt)
if resp_type == "text":
return resp
@ -81,13 +78,13 @@ class PromptManager:
try:
obj = self.parse_json(resp)
except:
print("Parse fail:", resp, flush=True)
raise RuntimeError("JSON parse fail")
print(obj, flush=True)
if self.prompts[id].schema:
try:
print(self.prompts[id].schema)
validate(instance=obj, schema=self.prompts[id].schema)
print("Validated", flush=True)
except Exception as e:
raise RuntimeError(f"Schema validation fail: {e}")

133
trustgraph-flow/trustgraph/model/prompt/template/service.py
View file

@ -3,6 +3,7 @@
Language service abstracts prompt engineering from LLM.
"""
import asyncio
import json
import re
@ -10,74 +11,59 @@ from .... schema import Definition, Relationship, Triple
from .... schema import Topic
from .... schema import PromptRequest, PromptResponse, Error
from .... schema import TextCompletionRequest, TextCompletionResponse
from .... schema import text_completion_request_queue
from .... schema import text_completion_response_queue
from .... schema import prompt_request_queue, prompt_response_queue
from .... base import ConsumerProducer
from .... clients.llm_client import LlmClient
from .... base import FlowProcessor
from .... base import ProducerSpec, ConsumerSpec, TextCompletionClientSpec
from . prompt_manager import PromptConfiguration, Prompt, PromptManager
module = ".".join(__name__.split(".")[1:-1])
default_ident = "prompt"
default_input_queue = prompt_request_queue
default_output_queue = prompt_response_queue
default_subscriber = module
class Processor(ConsumerProducer):
class Processor(FlowProcessor):
def __init__(self, **params):
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
)
id = params.get("id")
# Config key for prompts
self.config_key = params.get("config_type", "prompt")
super(Processor, self).__init__(
**params | {
"input_queue": input_queue,
"output_queue": output_queue,
"subscriber": subscriber,
"input_schema": PromptRequest,
"output_schema": PromptResponse,
"text_completion_request_queue": tc_request_queue,
"text_completion_response_queue": tc_response_queue,
"id": id,
}
)
self.llm = LlmClient(
subscriber=subscriber,
input_queue=tc_request_queue,
output_queue=tc_response_queue,
pulsar_host = self.pulsar_host,
pulsar_api_key=self.pulsar_api_key,
self.register_specification(
ConsumerSpec(
name = "request",
schema = PromptRequest,
handler = self.on_request
)
)
# System prompt hack
class Llm:
def __init__(self, llm):
self.llm = llm
def request(self, system, prompt):
print(system)
print(prompt, flush=True)
return self.llm.request(system, prompt)
self.register_specification(
TextCompletionClientSpec(
request_name = "text-completion-request",
response_name = "text-completion-response",
)
)
self.llm = Llm(self.llm)
self.register_specification(
ProducerSpec(
name = "response",
schema = PromptResponse
)
)
self.register_config_handler(self.on_prompt_config)
# Null configuration, should reload quickly
self.manager = PromptManager(
llm = self.llm,
config = PromptConfiguration("", {}, {})
)
async def on_config(self, version, config):
async def on_prompt_config(self, config, version):
print("Loading configuration version", version)
@ -111,7 +97,6 @@ class Processor(ConsumerProducer):
)
self.manager = PromptManager(
self.llm,
PromptConfiguration(
system,
{},
@ -126,7 +111,7 @@ class Processor(ConsumerProducer):
print("Exception:", e, flush=True)
print("Configuration reload failed", flush=True)
async def handle(self, msg):
async def on_request(self, msg, consumer, flow):
v = msg.value()
@ -138,7 +123,7 @@ class Processor(ConsumerProducer):
try:
print(v.terms)
print(v.terms, flush=True)
input = {
k: json.loads(v)
@ -146,14 +131,33 @@ class Processor(ConsumerProducer):
}
print(f"Handling kind {kind}...", flush=True)
print(input, flush=True)
resp = self.manager.invoke(kind, input)
async def llm(system, prompt):
print(system, flush=True)
print(prompt, flush=True)
resp = await flow("text-completion-request").text_completion(
system = system, prompt = prompt,
)
try:
return resp
except Exception as e:
print("LLM Exception:", e, flush=True)
return None
try:
resp = await self.manager.invoke(kind, input, llm)
except Exception as e:
print("Invocation exception:", e, flush=True)
raise e
print(resp, flush=True)
if isinstance(resp, str):
print("Send text response...", flush=True)
print(resp, flush=True)
r = PromptResponse(
text=resp,
@ -161,7 +165,7 @@ class Processor(ConsumerProducer):
error=None,
)
await self.send(r, properties={"id": id})
await flow("response").send(r, properties={"id": id})
return
@ -176,13 +180,13 @@ class Processor(ConsumerProducer):
error=None,
)
await self.send(r, properties={"id": id})
await flow("response").send(r, properties={"id": id})
return
except Exception as e:
print(f"Exception: {e}")
print(f"Exception: {e}", flush=True)
print("Send error response...", flush=True)
@ -194,11 +198,11 @@ class Processor(ConsumerProducer):
response=None,
)
await self.send(r, properties={"id": id})
await flow("response").send(r, properties={"id": id})
except Exception as e:
print(f"Exception: {e}")
print(f"Exception: {e}", flush=True)
print("Send error response...", flush=True)
@ -215,22 +219,7 @@ class Processor(ConsumerProducer):
@staticmethod
def add_args(parser):
ConsumerProducer.add_args(
parser, default_input_queue, default_subscriber,
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})',
)
FlowProcessor.add_args(parser)
parser.add_argument(
'--config-type',
@ -240,5 +229,5 @@ class Processor(ConsumerProducer):
def run():
Processor.launch(module, __doc__)
Processor.launch(default_ident, __doc__)

2
trustgraph-flow/trustgraph/model/text_completion/azure/llm.py
View file

@ -16,7 +16,7 @@ from .... log_level import LogLevel
from .... base import ConsumerProducer
from .... exceptions import TooManyRequests
module = ".".join(__name__.split(".")[1:-1])
module = "text-completion"
default_input_queue = text_completion_request_queue
default_output_queue = text_completion_response_queue

2
trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py
View file

@ -16,7 +16,7 @@ from .... log_level import LogLevel
from .... base import ConsumerProducer
from .... exceptions import TooManyRequests
module = ".".join(__name__.split(".")[1:-1])
module = "text-completion"
default_input_queue = text_completion_request_queue
default_output_queue = text_completion_response_queue

2
trustgraph-flow/trustgraph/model/text_completion/claude/llm.py
View file

@ -15,7 +15,7 @@ from .... log_level import LogLevel
from .... base import ConsumerProducer
from .... exceptions import TooManyRequests
module = ".".join(__name__.split(".")[1:-1])
module = "text-completion"
default_input_queue = text_completion_request_queue
default_output_queue = text_completion_response_queue

2
trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py
View file

@ -15,7 +15,7 @@ from .... log_level import LogLevel
from .... base import ConsumerProducer
from .... exceptions import TooManyRequests
module = ".".join(__name__.split(".")[1:-1])
module = "text-completion"
default_input_queue = text_completion_request_queue
default_output_queue = text_completion_response_queue

2
trustgraph-flow/trustgraph/model/text_completion/googleaistudio/llm.py
View file

@ -17,7 +17,7 @@ from .... log_level import LogLevel
from .... base import ConsumerProducer
from .... exceptions import TooManyRequests
module = ".".join(__name__.split(".")[1:-1])
module = "text-completion"
default_input_queue = text_completion_request_queue
default_output_queue = text_completion_response_queue

2
trustgraph-flow/trustgraph/model/text_completion/llamafile/llm.py
View file

@ -14,7 +14,7 @@ from .... log_level import LogLevel
from .... base import ConsumerProducer
from .... exceptions import TooManyRequests
module = ".".join(__name__.split(".")[1:-1])
module = "text-completion"
default_input_queue = text_completion_request_queue
default_output_queue = text_completion_response_queue

2
trustgraph-flow/trustgraph/model/text_completion/lmstudio/llm.py
View file

@ -15,7 +15,7 @@ from .... log_level import LogLevel
from .... base import ConsumerProducer
from .... exceptions import TooManyRequests
module = ".".join(__name__.split(".")[1:-1])
module = "text-completion"
default_input_queue = text_completion_request_queue
default_output_queue = text_completion_response_queue

2
trustgraph-flow/trustgraph/model/text_completion/mistral/llm.py
View file

@ -15,7 +15,7 @@ from .... log_level import LogLevel
from .... base import ConsumerProducer
from .... exceptions import TooManyRequests
module = ".".join(__name__.split(".")[1:-1])
module = "text-completion"
default_input_queue = text_completion_request_queue
default_output_queue = text_completion_response_queue

2
trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py
View file

@ -15,7 +15,7 @@ from .... log_level import LogLevel
from .... base import ConsumerProducer
from .... exceptions import TooManyRequests
module = ".".join(__name__.split(".")[1:-1])
module = "text-completion"
default_input_queue = text_completion_request_queue
default_output_queue = text_completion_response_queue

2
trustgraph-flow/trustgraph/model/text_completion/openai/llm.py
View file

@ -15,7 +15,7 @@ from .... log_level import LogLevel
from .... base import ConsumerProducer
from .... exceptions import TooManyRequests
module = ".".join(__name__.split(".")[1:-1])
module = "text-completion"
default_input_queue = text_completion_request_queue
default_output_queue = text_completion_response_queue

2
trustgraph-flow/trustgraph/query/doc_embeddings/milvus/service.py
View file

@ -11,7 +11,7 @@ from .... schema import document_embeddings_request_queue
from .... schema import document_embeddings_response_queue
from .... base import ConsumerProducer
module = ".".join(__name__.split(".")[1:-1])
module = "de-query"
default_input_queue = document_embeddings_request_queue
default_output_queue = document_embeddings_response_queue

2
trustgraph-flow/trustgraph/query/doc_embeddings/pinecone/service.py
View file

@ -16,7 +16,7 @@ from .... schema import document_embeddings_request_queue
from .... schema import document_embeddings_response_queue
from .... base import ConsumerProducer
module = ".".join(__name__.split(".")[1:-1])
module = "de-query"
default_input_queue = document_embeddings_request_queue
default_output_queue = document_embeddings_response_queue

70
trustgraph-flow/trustgraph/query/doc_embeddings/qdrant/service.py
View file

@ -7,71 +7,51 @@ of chunks
from qdrant_client import QdrantClient
from qdrant_client.models import PointStruct
from qdrant_client.models import Distance, VectorParams
import uuid
from .... schema import DocumentEmbeddingsRequest, DocumentEmbeddingsResponse
from .... schema import DocumentEmbeddingsResponse
from .... schema import Error, Value
from .... schema import document_embeddings_request_queue
from .... schema import document_embeddings_response_queue
from .... base import ConsumerProducer
from .... base import DocumentEmbeddingsQueryService
module = ".".join(__name__.split(".")[1:-1])
default_ident = "de-query"
default_input_queue = document_embeddings_request_queue
default_output_queue = document_embeddings_response_queue
default_subscriber = module
default_store_uri = 'http://localhost:6333'
class Processor(ConsumerProducer):
class Processor(DocumentEmbeddingsQueryService):
def __init__(self, **params):
input_queue = params.get("input_queue", default_input_queue)
output_queue = params.get("output_queue", default_output_queue)
subscriber = params.get("subscriber", default_subscriber)
store_uri = params.get("store_uri", default_store_uri)
#optional api key
api_key = params.get("api_key", None)
super(Processor, self).__init__(
**params | {
"input_queue": input_queue,
"output_queue": output_queue,
"subscriber": subscriber,
"input_schema": DocumentEmbeddingsRequest,
"output_schema": DocumentEmbeddingsResponse,
"store_uri": store_uri,
"api_key": api_key,
}
)
self.client = QdrantClient(url=store_uri, api_key=api_key)
self.qdrant = QdrantClient(url=store_uri, api_key=api_key)
async def handle(self, msg):
async def query_document_embeddings(self, msg):
try:
v = msg.value()
# Sender-produced ID
id = msg.properties()["id"]
print(f"Handling input {id}...", flush=True)
chunks = []
for vec in v.vectors:
for vec in msg.vectors:
dim = len(vec)
collection = (
"d_" + v.user + "_" + v.collection + "_" +
"d_" + msg.user + "_" + msg.collection + "_" +
str(dim)
)
search_result = self.client.query_points(
search_result = self.qdrant.query_points(
collection_name=collection,
query=vec,
limit=v.limit,
limit=msg.limit,
with_payload=True,
).points
@ -79,37 +59,17 @@ class Processor(ConsumerProducer):
ent = r.payload["doc"]
chunks.append(ent)
print("Send response...", flush=True)
r = DocumentEmbeddingsResponse(documents=chunks, error=None)
await self.send(r, properties={"id": id})
print("Done.", flush=True)
return chunks
except Exception as e:
print(f"Exception: {e}")
print("Send error response...", flush=True)
r = DocumentEmbeddingsResponse(
error=Error(
type = "llm-error",
message = str(e),
),
documents=None,
)
await self.send(r, properties={"id": id})
self.consumer.acknowledge(msg)
raise e
@staticmethod
def add_args(parser):
ConsumerProducer.add_args(
parser, default_input_queue, default_subscriber,
default_output_queue,
)
DocumentEmbeddingsQueryService.add_args(parser)
parser.add_argument(
'-t', '--store-uri',
@ -125,5 +85,5 @@ class Processor(ConsumerProducer):
def run():
Processor.launch(module, __doc__)
Processor.launch(default_ident, __doc__)

2
trustgraph-flow/trustgraph/query/graph_embeddings/milvus/service.py
View file

@ -11,7 +11,7 @@ from .... schema import graph_embeddings_request_queue
from .... schema import graph_embeddings_response_queue
from .... base import ConsumerProducer
module = ".".join(__name__.split(".")[1:-1])
module = "ge-query"
default_input_queue = graph_embeddings_request_queue
default_output_queue = graph_embeddings_response_queue

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

@ -16,7 +16,7 @@ from .... schema import graph_embeddings_request_queue
from .... schema import graph_embeddings_response_queue
from .... base import ConsumerProducer
module = ".".join(__name__.split(".")[1:-1])
module = "ge-query"
default_input_queue = graph_embeddings_request_queue
default_output_queue = graph_embeddings_response_queue

72
trustgraph-flow/trustgraph/query/graph_embeddings/qdrant/service.py
View file

@ -7,44 +7,32 @@ entities
from qdrant_client import QdrantClient
from qdrant_client.models import PointStruct
from qdrant_client.models import Distance, VectorParams
import uuid
from .... schema import GraphEmbeddingsRequest, GraphEmbeddingsResponse
from .... schema import GraphEmbeddingsResponse
from .... schema import Error, Value
from .... schema import graph_embeddings_request_queue
from .... schema import graph_embeddings_response_queue
from .... base import ConsumerProducer
from .... base import GraphEmbeddingsQueryService
module = ".".join(__name__.split(".")[1:-1])
default_ident = "ge-query"
default_input_queue = graph_embeddings_request_queue
default_output_queue = graph_embeddings_response_queue
default_subscriber = module
default_store_uri = 'http://localhost:6333'
class Processor(ConsumerProducer):
class Processor(GraphEmbeddingsQueryService):
def __init__(self, **params):
input_queue = params.get("input_queue", default_input_queue)
output_queue = params.get("output_queue", default_output_queue)
subscriber = params.get("subscriber", default_subscriber)
store_uri = params.get("store_uri", default_store_uri)
#optional api key
api_key = params.get("api_key", None)
super(Processor, self).__init__(
**params | {
"input_queue": input_queue,
"output_queue": output_queue,
"subscriber": subscriber,
"input_schema": GraphEmbeddingsRequest,
"output_schema": GraphEmbeddingsResponse,
"store_uri": store_uri,
"api_key": api_key,
}
)
self.client = QdrantClient(url=store_uri, api_key=api_key)
self.qdrant = QdrantClient(url=store_uri, api_key=api_key)
def create_value(self, ent):
if ent.startswith("http://") or ent.startswith("https://"):
@ -52,34 +40,27 @@ class Processor(ConsumerProducer):
else:
return Value(value=ent, is_uri=False)
async def handle(self, msg):
async def query_graph_embeddings(self, msg):
try:
v = msg.value()
# Sender-produced ID
id = msg.properties()["id"]
print(f"Handling input {id}...", flush=True)
entity_set = set()
entities = []
for vec in v.vectors:
for vec in msg.vectors:
dim = len(vec)
collection = (
"t_" + v.user + "_" + v.collection + "_" +
"t_" + msg.user + "_" + msg.collection + "_" +
str(dim)
)
# Heuristic hack, get (2*limit), so that we have more chance
# of getting (limit) entities
search_result = self.client.query_points(
search_result = self.qdrant.query_points(
collection_name=collection,
query=vec,
limit=v.limit * 2,
limit=msg.limit * 2,
with_payload=True,
).points
@ -92,10 +73,10 @@ class Processor(ConsumerProducer):
entities.append(ent)
# Keep adding entities until limit
if len(entity_set) >= v.limit: break
if len(entity_set) >= msg.limit: break
# Keep adding entities until limit
if len(entity_set) >= v.limit: break
if len(entity_set) >= msg.limit: break
ents2 = []
@ -105,36 +86,19 @@ class Processor(ConsumerProducer):
entities = ents2
print("Send response...", flush=True)
r = GraphEmbeddingsResponse(entities=entities, error=None)
await self.send(r, properties={"id": id})
return entities
print("Done.", flush=True)
except Exception as e:
print(f"Exception: {e}")
print("Send error response...", flush=True)
r = GraphEmbeddingsResponse(
error=Error(
type = "llm-error",
message = str(e),
),
entities=None,
)
await self.send(r, properties={"id": id})
self.consumer.acknowledge(msg)
raise e
@staticmethod
def add_args(parser):
ConsumerProducer.add_args(
parser, default_input_queue, default_subscriber,
default_output_queue,
)
GraphEmbeddingsQueryService.add_args(parser)
parser.add_argument(
'-t', '--store-uri',
@ -150,5 +114,5 @@ class Processor(ConsumerProducer):
def run():
Processor.launch(module, __doc__)
Processor.launch(default_ident, __doc__)

121
trustgraph-flow/trustgraph/query/triples/cassandra/service.py
View file

@ -7,38 +7,24 @@ null. Output is a list of triples.
from .... direct.cassandra import TrustGraph
from .... schema import TriplesQueryRequest, TriplesQueryResponse, Error
from .... schema import Value, Triple
from .... schema import triples_request_queue
from .... schema import triples_response_queue
from .... base import ConsumerProducer
from .... base import TriplesQueryService
module = ".".join(__name__.split(".")[1:-1])
default_ident = "triples-query"
default_input_queue = triples_request_queue
default_output_queue = triples_response_queue
default_subscriber = module
default_graph_host='localhost'
class Processor(ConsumerProducer):
class Processor(TriplesQueryService):
def __init__(self, **params):
input_queue = params.get("input_queue", default_input_queue)
output_queue = params.get("output_queue", default_output_queue)
subscriber = params.get("subscriber", default_subscriber)
graph_host = params.get("graph_host", default_graph_host)
graph_username = params.get("graph_username", None)
graph_password = params.get("graph_password", None)
super(Processor, self).__init__(
**params | {
"input_queue": input_queue,
"output_queue": output_queue,
"subscriber": subscriber,
"input_schema": TriplesQueryRequest,
"output_schema": TriplesQueryResponse,
"graph_host": graph_host,
"graph_username": graph_username,
"graph_password": graph_password,
}
)
@ -53,92 +39,85 @@ class Processor(ConsumerProducer):
else:
return Value(value=ent, is_uri=False)
async def handle(self, msg):
async def query_triples(self, query):
try:
v = msg.value()
table = (v.user, v.collection)
table = (query.user, query.collection)
if table != self.table:
if self.username and self.password:
self.tg = TrustGraph(
hosts=self.graph_host,
keyspace=v.user, table=v.collection,
keyspace=query.user, table=query.collection,
username=self.username, password=self.password
)
else:
self.tg = TrustGraph(
hosts=self.graph_host,
keyspace=v.user, table=v.collection,
keyspace=query.user, table=query.collection,
)
self.table = table
# Sender-produced ID
id = msg.properties()["id"]
print(f"Handling input {id}...", flush=True)
triples = []
if v.s is not None:
if v.p is not None:
if v.o is not None:
if query.s is not None:
if query.p is not None:
if query.o is not None:
resp = self.tg.get_spo(
v.s.value, v.p.value, v.o.value,
limit=v.limit
query.s.value, query.p.value, query.o.value,
limit=query.limit
)
triples.append((v.s.value, v.p.value, v.o.value))
triples.append((query.s.value, query.p.value, query.o.value))
else:
resp = self.tg.get_sp(
v.s.value, v.p.value,
limit=v.limit
query.s.value, query.p.value,
limit=query.limit
)
for t in resp:
triples.append((v.s.value, v.p.value, t.o))
triples.append((query.s.value, query.p.value, t.o))
else:
if v.o is not None:
if query.o is not None:
resp = self.tg.get_os(
v.o.value, v.s.value,
limit=v.limit
query.o.value, query.s.value,
limit=query.limit
)
for t in resp:
triples.append((v.s.value, t.p, v.o.value))
triples.append((query.s.value, t.p, query.o.value))
else:
resp = self.tg.get_s(
v.s.value,
limit=v.limit
query.s.value,
limit=query.limit
)
for t in resp:
triples.append((v.s.value, t.p, t.o))
triples.append((query.s.value, t.p, t.o))
else:
if v.p is not None:
if v.o is not None:
if query.p is not None:
if query.o is not None:
resp = self.tg.get_po(
v.p.value, v.o.value,
limit=v.limit
query.p.value, query.o.value,
limit=query.limit
)
for t in resp:
triples.append((t.s, v.p.value, v.o.value))
triples.append((t.s, query.p.value, query.o.value))
else:
resp = self.tg.get_p(
v.p.value,
limit=v.limit
query.p.value,
limit=query.limit
)
for t in resp:
triples.append((t.s, v.p.value, t.o))
triples.append((t.s, query.p.value, t.o))
else:
if v.o is not None:
if query.o is not None:
resp = self.tg.get_o(
v.o.value,
limit=v.limit
query.o.value,
limit=query.limit
)
for t in resp:
triples.append((t.s, t.p, v.o.value))
triples.append((t.s, t.p, query.o.value))
else:
resp = self.tg.get_all(
limit=v.limit
limit=query.limit
)
for t in resp:
triples.append((t.s, t.p, t.o))
@ -152,37 +131,17 @@ class Processor(ConsumerProducer):
for t in triples
]
print("Send response...", flush=True)
r = TriplesQueryResponse(triples=triples, error=None)
await self.send(r, properties={"id": id})
print("Done.", flush=True)
return triples
except Exception as e:
print(f"Exception: {e}")
print("Send error response...", flush=True)
r = TriplesQueryResponse(
error=Error(
type = "llm-error",
message = str(e),
),
response=None,
)
await self.send(r, properties={"id": id})
self.consumer.acknowledge(msg)
raise e
@staticmethod
def add_args(parser):
ConsumerProducer.add_args(
parser, default_input_queue, default_subscriber,
default_output_queue,
)
TriplesQueryService.add_args(parser)
parser.add_argument(
'-g', '--graph-host',
@ -205,5 +164,5 @@ class Processor(ConsumerProducer):
def run():
Processor.launch(module, __doc__)
Processor.launch(default_ident, __doc__)

2
trustgraph-flow/trustgraph/query/triples/falkordb/service.py
View file

@ -13,7 +13,7 @@ from .... schema import triples_request_queue
from .... schema import triples_response_queue
from .... base import ConsumerProducer
module = ".".join(__name__.split(".")[1:-1])
module = "triples-query"
default_input_queue = triples_request_queue
default_output_queue = triples_response_queue

2
trustgraph-flow/trustgraph/query/triples/memgraph/service.py
View file

@ -13,7 +13,7 @@ from .... schema import triples_request_queue
from .... schema import triples_response_queue
from .... base import ConsumerProducer
module = ".".join(__name__.split(".")[1:-1])
module = "triples-query"
default_input_queue = triples_request_queue
default_output_queue = triples_response_queue

2
trustgraph-flow/trustgraph/query/triples/neo4j/service.py
View file

@ -13,7 +13,7 @@ from .... schema import triples_request_queue
from .... schema import triples_response_queue
from .... base import ConsumerProducer
module = ".".join(__name__.split(".")[1:-1])
module = "triples-query"
default_input_queue = triples_request_queue
default_output_queue = triples_response_queue

94
trustgraph-flow/trustgraph/retrieval/document_rag/document_rag.py Normal file
View file

@ -0,0 +1,94 @@
import asyncio
LABEL="http://www.w3.org/2000/01/rdf-schema#label"
class Query:
def __init__(
self, rag, user, collection, verbose,
doc_limit=20
):
self.rag = rag
self.user = user
self.collection = collection
self.verbose = verbose
self.doc_limit = doc_limit
async def get_vector(self, query):
if self.verbose:
print("Compute embeddings...", flush=True)
qembeds = await self.rag.embeddings_client.embed(query)
if self.verbose:
print("Done.", flush=True)
return qembeds
async def get_docs(self, query):
vectors = await self.get_vector(query)
if self.verbose:
print("Get docs...", flush=True)
docs = await self.rag.doc_embeddings_client.query(
vectors, limit=self.doc_limit,
user=self.user, collection=self.collection,
)
if self.verbose:
print("Docs:", flush=True)
for doc in docs:
print(doc, flush=True)
return docs
class DocumentRag:
def __init__(
self, prompt_client, embeddings_client, doc_embeddings_client,
verbose=False,
):
self.verbose = verbose
self.prompt_client = prompt_client
self.embeddings_client = embeddings_client
self.doc_embeddings_client = doc_embeddings_client
if self.verbose:
print("Initialised", flush=True)
async def query(
self, query, user="trustgraph", collection="default",
doc_limit=20,
):
if self.verbose:
print("Construct prompt...", flush=True)
q = Query(
rag=self, user=user, collection=collection, verbose=self.verbose,
doc_limit=doc_limit
)
docs = await q.get_docs(query)
if self.verbose:
print("Invoke LLM...", flush=True)
print(docs)
print(query)
resp = await self.prompt_client.document_prompt(
query = query,
documents = docs
)
if self.verbose:
print("Done", flush=True)
return resp

191
trustgraph-flow/trustgraph/retrieval/document_rag/rag.py
View file

@ -5,88 +5,77 @@ Input is query, output is response.
"""
from ... schema import DocumentRagQuery, DocumentRagResponse, Error
from ... schema import document_rag_request_queue, document_rag_response_queue
from ... schema import prompt_request_queue
from ... schema import prompt_response_queue
from ... schema import embeddings_request_queue
from ... schema import embeddings_response_queue
from ... schema import document_embeddings_request_queue
from ... schema import document_embeddings_response_queue
from ... log_level import LogLevel
from ... document_rag import DocumentRag
from ... base import ConsumerProducer
from . document_rag import DocumentRag
from ... base import FlowProcessor, ConsumerSpec, ProducerSpec
from ... base import PromptClientSpec, EmbeddingsClientSpec
from ... base import DocumentEmbeddingsClientSpec
module = ".".join(__name__.split(".")[1:-1])
default_ident = "document-rag"
default_input_queue = document_rag_request_queue
default_output_queue = document_rag_response_queue
default_subscriber = module
class Processor(ConsumerProducer):
class Processor(FlowProcessor):
def __init__(self, **params):
input_queue = params.get("input_queue", default_input_queue)
output_queue = params.get("output_queue", default_output_queue)
subscriber = params.get("subscriber", default_subscriber)
pr_request_queue = params.get(
"prompt_request_queue", prompt_request_queue
)
pr_response_queue = params.get(
"prompt_response_queue", prompt_response_queue
)
emb_request_queue = params.get(
"embeddings_request_queue", embeddings_request_queue
)
emb_response_queue = params.get(
"embeddings_response_queue", embeddings_response_queue
)
de_request_queue = params.get(
"document_embeddings_request_queue",
document_embeddings_request_queue
)
de_response_queue = params.get(
"document_embeddings_response_queue",
document_embeddings_response_queue
)
id = params.get("id", default_ident)
doc_limit = params.get("doc_limit", 10)
doc_limit = params.get("doc_limit", 5)
super(Processor, self).__init__(
**params | {
"input_queue": input_queue,
"output_queue": output_queue,
"subscriber": subscriber,
"input_schema": DocumentRagQuery,
"output_schema": DocumentRagResponse,
"prompt_request_queue": pr_request_queue,
"prompt_response_queue": pr_response_queue,
"embeddings_request_queue": emb_request_queue,
"embeddings_response_queue": emb_response_queue,
"document_embeddings_request_queue": de_request_queue,
"document_embeddings_response_queue": de_response_queue,
"id": id,
"doc_limit": doc_limit,
}
)
self.rag = DocumentRag(
pulsar_host=self.pulsar_host,
pulsar_api_key=self.pulsar_api_key,
pr_request_queue=pr_request_queue,
pr_response_queue=pr_response_queue,
emb_request_queue=emb_request_queue,
emb_response_queue=emb_response_queue,
de_request_queue=de_request_queue,
de_response_queue=de_response_queue,
verbose=True,
module=module,
)
self.doc_limit = doc_limit
async def handle(self, msg):
self.register_specification(
ConsumerSpec(
name = "request",
schema = DocumentRagQuery,
handler = self.on_request,
)
)
self.register_specification(
EmbeddingsClientSpec(
request_name = "embeddings-request",
response_name = "embeddings-response",
)
)
self.register_specification(
DocumentEmbeddingsClientSpec(
request_name = "document-embeddings-request",
response_name = "document-embeddings-response",
)
)
self.register_specification(
PromptClientSpec(
request_name = "prompt-request",
response_name = "prompt-response",
)
)
self.register_specification(
ProducerSpec(
name = "response",
schema = DocumentRagResponse,
)
)
async def on_request(self, msg, consumer, flow):
try:
self.rag = DocumentRag(
embeddings_client = flow("embeddings-request"),
doc_embeddings_client = flow("document-embeddings-request"),
prompt_client = flow("prompt-request"),
verbose=True,
)
v = msg.value()
# Sender-produced ID
@ -99,11 +88,15 @@ class Processor(ConsumerProducer):
else:
doc_limit = self.doc_limit
response = self.rag.query(v.query, doc_limit=doc_limit)
response = await self.rag.query(v.query, doc_limit=doc_limit)
print("Send response...", flush=True)
r = DocumentRagResponse(response = response, error=None)
await self.send(r, properties={"id": id})
await flow("response").send(
DocumentRagResponse(
response = response,
error = None
),
properties = {"id": id}
)
print("Done.", flush=True)
@ -113,25 +106,21 @@ class Processor(ConsumerProducer):
print("Send error response...", flush=True)
r = DocumentRagResponse(
error=Error(
type = "llm-error",
message = str(e),
await flow("response").send(
DocumentRagResponse(
response = None,
error = Error(
type = "document-rag-error",
message = str(e),
),
),
response=None,
properties = {"id": id}
)
await self.send(r, properties={"id": id})
self.consumer.acknowledge(msg)
@staticmethod
def add_args(parser):
ConsumerProducer.add_args(
parser, default_input_queue, default_subscriber,
default_output_queue,
)
FlowProcessor.add_args(parser)
parser.add_argument(
'-d', '--doc-limit',
@ -140,43 +129,7 @@ class Processor(ConsumerProducer):
help=f'Default document fetch limit (default: 10)'
)
parser.add_argument(
'--prompt-request-queue',
default=prompt_request_queue,
help=f'Prompt request queue (default: {prompt_request_queue})',
)
parser.add_argument(
'--prompt-response-queue',
default=prompt_response_queue,
help=f'Prompt response queue (default: {prompt_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 response queue (default: {embeddings_response_queue})',
)
parser.add_argument(
'--document-embeddings-request-queue',
default=document_embeddings_request_queue,
help=f'Document embeddings request queue (default: {document_embeddings_request_queue})',
)
parser.add_argument(
'--document-embeddings-response-queue',
default=document_embeddings_response_queue,
help=f'Document embeddings response queue (default: {document_embeddings_response_queue})',
)
def run():
Processor.launch(module, __doc__)
Processor.launch(default_ident, __doc__)

218
trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py Normal file
View file

@ -0,0 +1,218 @@
import asyncio
LABEL="http://www.w3.org/2000/01/rdf-schema#label"
class Query:
def __init__(
self, rag, user, collection, verbose,
entity_limit=50, triple_limit=30, max_subgraph_size=1000,
max_path_length=2,
):
self.rag = rag
self.user = user
self.collection = collection
self.verbose = verbose
self.entity_limit = entity_limit
self.triple_limit = triple_limit
self.max_subgraph_size = max_subgraph_size
self.max_path_length = max_path_length
async def get_vector(self, query):
if self.verbose:
print("Compute embeddings...", flush=True)
qembeds = await self.rag.embeddings_client.embed(query)
if self.verbose:
print("Done.", flush=True)
return qembeds
async def get_entities(self, query):
vectors = await self.get_vector(query)
if self.verbose:
print("Get entities...", flush=True)
entities = await self.rag.graph_embeddings_client.query(
vectors=vectors, limit=self.entity_limit,
user=self.user, collection=self.collection,
)
entities = [
str(e)
for e in entities
]
if self.verbose:
print("Entities:", flush=True)
for ent in entities:
print(" ", ent, flush=True)
return entities
async def maybe_label(self, e):
if e in self.rag.label_cache:
return self.rag.label_cache[e]
res = await self.rag.triples_client.query(
s=e, p=LABEL, o=None, limit=1,
user=self.user, collection=self.collection,
)
if len(res) == 0:
self.rag.label_cache[e] = e
return e
self.rag.label_cache[e] = str(res[0].o)
return self.rag.label_cache[e]
async def follow_edges(self, ent, subgraph, path_length):
# Not needed?
if path_length <= 0:
return
# Stop spanning around if the subgraph is already maxed out
if len(subgraph) >= self.max_subgraph_size:
return
res = await self.rag.triples_client.query(
s=ent, p=None, o=None,
limit=self.triple_limit,
user=self.user, collection=self.collection,
)
for triple in res:
subgraph.add(
(str(triple.s), str(triple.p), str(triple.o))
)
if path_length > 1:
await self.follow_edges(str(triple.o), subgraph, path_length-1)
res = await self.rag.triples_client.query(
s=None, p=ent, o=None,
limit=self.triple_limit,
user=self.user, collection=self.collection,
)
for triple in res:
subgraph.add(
(str(triple.s), str(triple.p), str(triple.o))
)
res = await self.rag.triples_client.query(
s=None, p=None, o=ent,
limit=self.triple_limit,
user=self.user, collection=self.collection,
)
for triple in res:
subgraph.add(
(str(triple.s), str(triple.p), str(triple.o))
)
if path_length > 1:
await self.follow_edges(
str(triple.s), subgraph, path_length-1
)
async def get_subgraph(self, query):
entities = await self.get_entities(query)
if self.verbose:
print("Get subgraph...", flush=True)
subgraph = set()
for ent in entities:
await self.follow_edges(ent, subgraph, self.max_path_length)
subgraph = list(subgraph)
return subgraph
async def get_labelgraph(self, query):
subgraph = await self.get_subgraph(query)
sg2 = []
for edge in subgraph:
if edge[1] == LABEL:
continue
s = await self.maybe_label(edge[0])
p = await self.maybe_label(edge[1])
o = await self.maybe_label(edge[2])
sg2.append((s, p, o))
sg2 = sg2[0:self.max_subgraph_size]
if self.verbose:
print("Subgraph:", flush=True)
for edge in sg2:
print(" ", str(edge), flush=True)
if self.verbose:
print("Done.", flush=True)
return sg2
class GraphRag:
def __init__(
self, prompt_client, embeddings_client, graph_embeddings_client,
triples_client, verbose=False,
):
self.verbose = verbose
self.prompt_client = prompt_client
self.embeddings_client = embeddings_client
self.graph_embeddings_client = graph_embeddings_client
self.triples_client = triples_client
self.label_cache = {}
if self.verbose:
print("Initialised", flush=True)
async def query(
self, query, user = "trustgraph", collection = "default",
entity_limit = 50, triple_limit = 30, max_subgraph_size = 1000,
max_path_length = 2,
):
if self.verbose:
print("Construct prompt...", flush=True)
q = Query(
rag = self, user = user, collection = collection,
verbose = self.verbose, entity_limit = entity_limit,
triple_limit = triple_limit,
max_subgraph_size = max_subgraph_size,
max_path_length = max_path_length,
)
kg = await q.get_labelgraph(query)
if self.verbose:
print("Invoke LLM...", flush=True)
print(kg)
print(query)
resp = await self.prompt_client.kg_prompt(query, kg)
if self.verbose:
print("Done", flush=True)
return resp

228
trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py
View file

@ -5,57 +5,18 @@ Input is query, output is response.
"""
from ... schema import GraphRagQuery, GraphRagResponse, Error
from ... schema import graph_rag_request_queue, graph_rag_response_queue
from ... schema import prompt_request_queue
from ... schema import prompt_response_queue
from ... schema import embeddings_request_queue
from ... schema import embeddings_response_queue
from ... schema import graph_embeddings_request_queue
from ... schema import graph_embeddings_response_queue
from ... schema import triples_request_queue
from ... schema import triples_response_queue
from ... log_level import LogLevel
from ... graph_rag import GraphRag
from ... base import ConsumerProducer
from . graph_rag import GraphRag
from ... base import FlowProcessor, ConsumerSpec, ProducerSpec
from ... base import PromptClientSpec, EmbeddingsClientSpec
from ... base import GraphEmbeddingsClientSpec, TriplesClientSpec
module = ".".join(__name__.split(".")[1:-1])
default_ident = "graph-rag"
default_input_queue = graph_rag_request_queue
default_output_queue = graph_rag_response_queue
default_subscriber = module
class Processor(ConsumerProducer):
class Processor(FlowProcessor):
def __init__(self, **params):
input_queue = params.get("input_queue", default_input_queue)
output_queue = params.get("output_queue", default_output_queue)
subscriber = params.get("subscriber", default_subscriber)
pr_request_queue = params.get(
"prompt_request_queue", prompt_request_queue
)
pr_response_queue = params.get(
"prompt_response_queue", prompt_response_queue
)
emb_request_queue = params.get(
"embeddings_request_queue", embeddings_request_queue
)
emb_response_queue = params.get(
"embeddings_response_queue", embeddings_response_queue
)
ge_request_queue = params.get(
"graph_embeddings_request_queue", graph_embeddings_request_queue
)
ge_response_queue = params.get(
"graph_embeddings_response_queue", graph_embeddings_response_queue
)
tpl_request_queue = params.get(
"triples_request_queue", triples_request_queue
)
tpl_response_queue = params.get(
"triples_response_queue", triples_response_queue
)
id = params.get("id", default_ident)
entity_limit = params.get("entity_limit", 50)
triple_limit = params.get("triple_limit", 30)
@ -64,49 +25,74 @@ class Processor(ConsumerProducer):
super(Processor, self).__init__(
**params | {
"input_queue": input_queue,
"output_queue": output_queue,
"subscriber": subscriber,
"input_schema": GraphRagQuery,
"output_schema": GraphRagResponse,
"id": id,
"entity_limit": entity_limit,
"triple_limit": triple_limit,
"max_subgraph_size": max_subgraph_size,
"prompt_request_queue": pr_request_queue,
"prompt_response_queue": pr_response_queue,
"embeddings_request_queue": emb_request_queue,
"embeddings_response_queue": emb_response_queue,
"graph_embeddings_request_queue": ge_request_queue,
"graph_embeddings_response_queue": ge_response_queue,
"triples_request_queue": triples_request_queue,
"triples_response_queue": triples_response_queue,
"max_path_length": max_path_length,
}
)
self.rag = GraphRag(
pulsar_host=self.pulsar_host,
pulsar_api_key=self.pulsar_api_key,
pr_request_queue=pr_request_queue,
pr_response_queue=pr_response_queue,
emb_request_queue=emb_request_queue,
emb_response_queue=emb_response_queue,
ge_request_queue=ge_request_queue,
ge_response_queue=ge_response_queue,
tpl_request_queue=triples_request_queue,
tpl_response_queue=triples_response_queue,
verbose=True,
module=module,
)
self.default_entity_limit = entity_limit
self.default_triple_limit = triple_limit
self.default_max_subgraph_size = max_subgraph_size
self.default_max_path_length = max_path_length
async def handle(self, msg):
self.register_specification(
ConsumerSpec(
name = "request",
schema = GraphRagQuery,
handler = self.on_request,
)
)
self.register_specification(
EmbeddingsClientSpec(
request_name = "embeddings-request",
response_name = "embeddings-response",
)
)
self.register_specification(
GraphEmbeddingsClientSpec(
request_name = "graph-embeddings-request",
response_name = "graph-embeddings-response",
)
)
self.register_specification(
TriplesClientSpec(
request_name = "triples-request",
response_name = "triples-response",
)
)
self.register_specification(
PromptClientSpec(
request_name = "prompt-request",
response_name = "prompt-response",
)
)
self.register_specification(
ProducerSpec(
name = "response",
schema = GraphRagResponse,
)
)
async def on_request(self, msg, consumer, flow):
try:
self.rag = GraphRag(
embeddings_client = flow("embeddings-request"),
graph_embeddings_client = flow("graph-embeddings-request"),
triples_client = flow("triples-request"),
prompt_client = flow("prompt-request"),
verbose=True,
)
v = msg.value()
# Sender-produced ID
@ -134,16 +120,20 @@ class Processor(ConsumerProducer):
else:
max_path_length = self.default_max_path_length
response = self.rag.query(
query=v.query, user=v.user, collection=v.collection,
entity_limit=entity_limit, triple_limit=triple_limit,
max_subgraph_size=max_subgraph_size,
max_path_length=max_path_length,
response = await self.rag.query(
query = v.query, user = v.user, collection = v.collection,
entity_limit = entity_limit, triple_limit = triple_limit,
max_subgraph_size = max_subgraph_size,
max_path_length = max_path_length,
)
print("Send response...", flush=True)
r = GraphRagResponse(response=response, error=None)
await self.send(r, properties={"id": id})
await flow("response").send(
GraphRagResponse(
response = response,
error = None
),
properties = {"id": id}
)
print("Done.", flush=True)
@ -153,25 +143,21 @@ class Processor(ConsumerProducer):
print("Send error response...", flush=True)
r = GraphRagResponse(
error=Error(
type = "llm-error",
message = str(e),
await flow("response").send(
GraphRagResponse(
response = None,
error = Error(
type = "graph-rag-error",
message = str(e),
),
),
response=None,
properties = {"id": id}
)
await self.send(r, properties={"id": id})
self.consumer.acknowledge(msg)
@staticmethod
def add_args(parser):
ConsumerProducer.add_args(
parser, default_input_queue, default_subscriber,
default_output_queue,
)
FlowProcessor.add_args(parser)
parser.add_argument(
'-e', '--entity-limit',
@ -201,55 +187,7 @@ class Processor(ConsumerProducer):
help=f'Default max path length (default: 2)'
)
parser.add_argument(
'--prompt-request-queue',
default=prompt_request_queue,
help=f'Prompt request queue (default: {prompt_request_queue})',
)
parser.add_argument(
'--prompt-response-queue',
default=prompt_response_queue,
help=f'Prompt response queue (default: {prompt_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 response queue (default: {embeddings_response_queue})',
)
parser.add_argument(
'--graph-embeddings-request-queue',
default=graph_embeddings_request_queue,
help=f'Graph embeddings request queue (default: {graph_embeddings_request_queue})',
)
parser.add_argument(
'--graph-embeddings-response-queue',
default=graph_embeddings_response_queue,
help=f'Graph embeddings response queue (default: {graph_embeddings_response_queue})',
)
parser.add_argument(
'--triples-request-queue',
default=triples_request_queue,
help=f'Triples request queue (default: {triples_request_queue})',
)
parser.add_argument(
'--triples-response-queue',
default=triples_response_queue,
help=f'Triples response queue (default: {triples_response_queue})',
)
def run():
Processor.launch(module, __doc__)
Processor.launch(default_ident, __doc__)

2
trustgraph-flow/trustgraph/storage/doc_embeddings/milvus/write.py
View file

@ -10,7 +10,7 @@ from .... schema import document_embeddings_store_queue
from .... log_level import LogLevel
from .... base import Consumer
module = ".".join(__name__.split(".")[1:-1])
module = "de-write"
default_input_queue = document_embeddings_store_queue
default_subscriber = module

2
trustgraph-flow/trustgraph/storage/doc_embeddings/pinecone/write.py
View file

@ -16,7 +16,7 @@ from .... schema import document_embeddings_store_queue
from .... log_level import LogLevel
from .... base import Consumer
module = ".".join(__name__.split(".")[1:-1])
module = "de-write"
default_input_queue = document_embeddings_store_queue
default_subscriber = module

39
trustgraph-flow/trustgraph/storage/doc_embeddings/qdrant/write.py
View file

@ -8,31 +8,21 @@ from qdrant_client.models import PointStruct
from qdrant_client.models import Distance, VectorParams
import uuid
from .... schema import DocumentEmbeddings
from .... schema import document_embeddings_store_queue
from .... log_level import LogLevel
from .... base import Consumer
from .... base import DocumentEmbeddingsStoreService
module = ".".join(__name__.split(".")[1:-1])
default_ident = "de-write"
default_input_queue = document_embeddings_store_queue
default_subscriber = module
default_store_uri = 'http://localhost:6333'
class Processor(Consumer):
class Processor(DocumentEmbeddingsStoreService):
def __init__(self, **params):
input_queue = params.get("input_queue", default_input_queue)
subscriber = params.get("subscriber", default_subscriber)
store_uri = params.get("store_uri", default_store_uri)
api_key = params.get("api_key", None)
super(Processor, self).__init__(
**params | {
"input_queue": input_queue,
"subscriber": subscriber,
"input_schema": DocumentEmbeddings,
"store_uri": store_uri,
"api_key": api_key,
}
@ -40,13 +30,11 @@ class Processor(Consumer):
self.last_collection = None
self.client = QdrantClient(url=store_uri)
self.qdrant = QdrantClient(url=store_uri, api_key=api_key)
async def handle(self, msg):
async def store_document_embeddings(self, message):
v = msg.value()
for emb in v.chunks:
for emb in message.chunks:
chunk = emb.chunk.decode("utf-8")
if chunk == "": return
@ -55,16 +43,17 @@ class Processor(Consumer):
dim = len(vec)
collection = (
"d_" + v.metadata.user + "_" + v.metadata.collection + "_" +
"d_" + message.metadata.user + "_" +
message.metadata.collection + "_" +
str(dim)
)
if collection != self.last_collection:
if not self.client.collection_exists(collection):
if not self.qdrant.collection_exists(collection):
try:
self.client.create_collection(
self.qdrant.create_collection(
collection_name=collection,
vectors_config=VectorParams(
size=dim, distance=Distance.COSINE
@ -76,7 +65,7 @@ class Processor(Consumer):
self.last_collection = collection
self.client.upsert(
self.qdrant.upsert(
collection_name=collection,
points=[
PointStruct(
@ -92,9 +81,7 @@ class Processor(Consumer):
@staticmethod
def add_args(parser):
Consumer.add_args(
parser, default_input_queue, default_subscriber,
)
DocumentEmbeddingsStoreService.add_args(parser)
parser.add_argument(
'-t', '--store-uri',
@ -110,5 +97,5 @@ class Processor(Consumer):
def run():
Processor.launch(module, __doc__)
Processor.launch(default_ident, __doc__)

2
trustgraph-flow/trustgraph/storage/graph_embeddings/milvus/write.py
View file

@ -9,7 +9,7 @@ from .... log_level import LogLevel
from .... direct.milvus_graph_embeddings import EntityVectors
from .... base import Consumer
module = ".".join(__name__.split(".")[1:-1])
module = "ge-write"
default_input_queue = graph_embeddings_store_queue
default_subscriber = module

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

@ -15,7 +15,7 @@ from .... schema import graph_embeddings_store_queue
from .... log_level import LogLevel
from .... base import Consumer
module = ".".join(__name__.split(".")[1:-1])
module = "ge-write"
default_input_queue = graph_embeddings_store_queue
default_subscriber = module

38
trustgraph-flow/trustgraph/storage/graph_embeddings/qdrant/write.py
View file

@ -8,31 +8,21 @@ from qdrant_client.models import PointStruct
from qdrant_client.models import Distance, VectorParams
import uuid
from .... schema import GraphEmbeddings
from .... schema import graph_embeddings_store_queue
from .... log_level import LogLevel
from .... base import Consumer
from .... base import GraphEmbeddingsStoreService
module = ".".join(__name__.split(".")[1:-1])
default_ident = "ge-write"
default_input_queue = graph_embeddings_store_queue
default_subscriber = module
default_store_uri = 'http://localhost:6333'
class Processor(Consumer):
class Processor(GraphEmbeddingsStoreService):
def __init__(self, **params):
input_queue = params.get("input_queue", default_input_queue)
subscriber = params.get("subscriber", default_subscriber)
store_uri = params.get("store_uri", default_store_uri)
api_key = params.get("api_key", None)
super(Processor, self).__init__(
**params | {
"input_queue": input_queue,
"subscriber": subscriber,
"input_schema": GraphEmbeddings,
"store_uri": store_uri,
"api_key": api_key,
}
@ -40,7 +30,7 @@ class Processor(Consumer):
self.last_collection = None
self.client = QdrantClient(url=store_uri, api_key=api_key)
self.qdrant = QdrantClient(url=store_uri, api_key=api_key)
def get_collection(self, dim, user, collection):
@ -50,10 +40,10 @@ class Processor(Consumer):
if cname != self.last_collection:
if not self.client.collection_exists(cname):
if not self.qdrant.collection_exists(cname):
try:
self.client.create_collection(
self.qdrant.create_collection(
collection_name=cname,
vectors_config=VectorParams(
size=dim, distance=Distance.COSINE
@ -67,11 +57,9 @@ class Processor(Consumer):
return cname
async def handle(self, msg):
async def store_graph_embeddings(self, message):
v = msg.value()
for entity in v.entities:
for entity in message.entities:
if entity.entity.value == "" or entity.entity.value is None: return
@ -80,10 +68,10 @@ class Processor(Consumer):
dim = len(vec)
collection = self.get_collection(
dim, v.metadata.user, v.metadata.collection
dim, message.metadata.user, message.metadata.collection
)
self.client.upsert(
self.qdrant.upsert(
collection_name=collection,
points=[
PointStruct(
@ -99,9 +87,7 @@ class Processor(Consumer):
@staticmethod
def add_args(parser):
Consumer.add_args(
parser, default_input_queue, default_subscriber,
)
GraphEmbeddingsStoreService.add_args(parser)
parser.add_argument(
'-t', '--store-uri',
@ -117,5 +103,5 @@ class Processor(Consumer):
def run():
Processor.launch(module, __doc__)
Processor.launch(default_ident, __doc__)

2
trustgraph-flow/trustgraph/storage/object_embeddings/milvus/write.py
View file

@ -9,7 +9,7 @@ from .... log_level import LogLevel
from .... direct.milvus_object_embeddings import ObjectVectors
from .... base import Consumer
module = ".".join(__name__.split(".")[1:-1])
module = "oe-write"
default_input_queue = object_embeddings_store_queue
default_subscriber = module

2
trustgraph-flow/trustgraph/storage/rows/cassandra/write.py
View file

@ -17,7 +17,7 @@ from .... schema import rows_store_queue
from .... log_level import LogLevel
from .... base import Consumer
module = ".".join(__name__.split(".")[1:-1])
module = "rows-write"
ssl_context = SSLContext(PROTOCOL_TLSv1_2)
default_input_queue = rows_store_queue

41
trustgraph-flow/trustgraph/storage/triples/cassandra/write.py
View file

@ -10,35 +10,26 @@ import argparse
import time
from .... direct.cassandra import TrustGraph
from .... schema import Triples
from .... schema import triples_store_queue
from .... log_level import LogLevel
from .... base import Consumer
from .... base import TriplesStoreService
module = ".".join(__name__.split(".")[1:-1])
default_ident = "triples-write"
default_input_queue = triples_store_queue
default_subscriber = module
default_graph_host='localhost'
class Processor(Consumer):
class Processor(TriplesStoreService):
def __init__(self, **params):
input_queue = params.get("input_queue", default_input_queue)
subscriber = params.get("subscriber", default_subscriber)
id = params.get("id", default_ident)
graph_host = params.get("graph_host", default_graph_host)
graph_username = params.get("graph_username", None)
graph_password = params.get("graph_password", None)
super(Processor, self).__init__(
**params | {
"input_queue": input_queue,
"subscriber": subscriber,
"input_schema": Triples,
"graph_host": graph_host,
"graph_username": graph_username,
"graph_password": graph_password,
"graph_username": graph_username
}
)
@ -47,11 +38,9 @@ class Processor(Consumer):
self.password = graph_password
self.table = None
async def handle(self, msg):
async def store_triples(self, message):
v = msg.value()
table = (v.metadata.user, v.metadata.collection)
table = (message.metadata.user, message.metadata.collection)
if self.table is None or self.table != table:
@ -61,13 +50,15 @@ class Processor(Consumer):
if self.username and self.password:
self.tg = TrustGraph(
hosts=self.graph_host,
keyspace=v.metadata.user, table=v.metadata.collection,
keyspace=message.metadata.user,
table=message.metadata.collection,
username=self.username, password=self.password
)
else:
self.tg = TrustGraph(
hosts=self.graph_host,
keyspace=v.metadata.user, table=v.metadata.collection,
keyspace=message.metadata.user,
table=message.metadata.collection,
)
except Exception as e:
print("Exception", e, flush=True)
@ -76,7 +67,7 @@ class Processor(Consumer):
self.table = table
for t in v.triples:
for t in message.triples:
self.tg.insert(
t.s.value,
t.p.value,
@ -86,9 +77,7 @@ class Processor(Consumer):
@staticmethod
def add_args(parser):
Consumer.add_args(
parser, default_input_queue, default_subscriber,
)
TriplesStoreService.add_args(parser)
parser.add_argument(
'-g', '--graph-host',
@ -110,5 +99,5 @@ class Processor(Consumer):
def run():
Processor.launch(module, __doc__)
Processor.launch(default_ident, __doc__)

2
trustgraph-flow/trustgraph/storage/triples/falkordb/write.py
View file

@ -16,7 +16,7 @@ from .... schema import triples_store_queue
from .... log_level import LogLevel
from .... base import Consumer
module = ".".join(__name__.split(".")[1:-1])
module = "triples-write"
default_input_queue = triples_store_queue
default_subscriber = module

2
trustgraph-flow/trustgraph/storage/triples/memgraph/write.py
View file

@ -16,7 +16,7 @@ from .... schema import triples_store_queue
from .... log_level import LogLevel
from .... base import Consumer
module = ".".join(__name__.split(".")[1:-1])
module = "triples-write"
default_input_queue = triples_store_queue
default_subscriber = module

2
trustgraph-flow/trustgraph/storage/triples/neo4j/write.py
View file

@ -16,7 +16,7 @@ from .... schema import triples_store_queue
from .... log_level import LogLevel
from .... base import Consumer
module = ".".join(__name__.split(".")[1:-1])
module = "triples-write"
default_input_queue = triples_store_queue
default_subscriber = module