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

Request/response spec working, implemented in extract-defs and vertexai

Browse source
This commit is contained in:
Cyber MacGeddon 2025-04-18 18:35:17 +01:00
parent 36013212e0
commit 1047cd2fa1
22 changed files with 378 additions and 370 deletions
Download patch file Download diff file Expand all files Collapse all files

8
trustgraph-base/trustgraph/base/__init__.py
View file

@ -7,7 +7,9 @@ from . publisher import Publisher
from . subscriber import Subscriber
from . metrics import ProcessorMetrics, ConsumerMetrics, ProducerMetrics
from . flow_processor import FlowProcessor
from . flow_processor import ConsumerSpec, SettingSpec
from . flow_processor import SubscriberSpec, ProducerSpec
from . request_response import RequestResponseService
from . consumer_spec import ConsumerSpec
from . setting_spec import SettingSpec
from . producer_spec import ProducerSpec
from . subscriber_spec import SubscriberSpec
from . request_response_spec import RequestResponseSpec

46
trustgraph-base/trustgraph/base/async_processor.py
View file

@ -30,7 +30,7 @@ class AsyncProcessor:
self.id = params.get("id")
# Register a pulsar client
self.client = PulsarClient(**params)
self.pulsar_client = PulsarClient(**params)
# Initialise metrics, records the parameters
ProcessorMetrics(id=self.id).info({
@ -58,10 +58,17 @@ class AsyncProcessor:
config_subscriber_id = str(uuid.uuid4())
# Subscribe to config queue
self.config_sub_task = self.subscribe(
self.config_sub_task = Consumer(
taskgroup = self.taskgroup,
client = self.client,
subscriber = config_subscriber_id,
flow = None,
queue = self.config_push_queue, subscriber = config_subscriber_id,
schema = ConfigPush, handler = self.on_config_change,
topic = self.config_push_queue,
schema = ConfigPush,
handler = self.on_config_change,
# This causes new subscriptions to view the entire history of
# configuration
@ -85,6 +92,10 @@ class AsyncProcessor:
@property
def pulsar_host(self): return self.client.pulsar_host
# Returns the pulsar client
@property
def client(self): return self.pulsar_client.client
# Register a new event handler for configuration change
def register_config_handler(self, handler):
self.config_handlers.append(handler)
@ -112,33 +123,6 @@ class AsyncProcessor:
while self.running:
await asyncio.sleep(2)
# Subscribe to a topic, returns a Consumer
def subscribe(
self, flow, queue, subscriber, schema, handler, metrics=None,
start_of_messages=False
):
return Consumer(
flow = flow,
taskgroup = self.taskgroup,
client = self.client,
queue = queue,
subscriber = subscriber,
schema = schema,
handler = handler,
metrics = metrics,
start_of_messages=start_of_messages,
)
# Open a mechanism to publish messages to a topic. Returns a subscriber
def publish(self, queue, schema, metrics=None):
return Producer(
client = self.client,
queue = queue,
schema = schema,
metrics = metrics,
)
# Startup fabric. This runs in 'async' mode, creates a taskgroup and
# runs the producer.
@classmethod

26
trustgraph-base/trustgraph/base/consumer.py
View file

@ -1,4 +1,6 @@
from pulsar.schema import JsonSchema
import pulsar
import _pulsar
import asyncio
import time
@ -8,19 +10,18 @@ from .. exceptions import TooManyRequests
class Consumer:
def __init__(
self, taskgroup, flow, client, queue, subscriber, schema,
self, taskgroup, flow, client, topic, subscriber, schema,
handler,
metrics = None,
start_of_messages=False,
rate_limit_retry_time = 10, rate_limit_timeout = 7200,
reconnect_time = 5,
):
self.taskgroup = taskgroup
self.flow = flow
self.client = client
self.queue = queue
self.topic = topic
self.subscriber = subscriber
self.schema = schema
self.handler = handler
@ -43,7 +44,8 @@ class Consumer:
self.running = False
if hasattr(self, "consumer"):
self.consumer.close()
if self.consumer:
self.consumer.close()
async def stop(self):
@ -69,13 +71,19 @@ class Consumer:
try:
print(self.queue, "subscribing...", flush=True)
print(self.topic, "subscribing...", flush=True)
if self.start_of_messages:
pos = pulsar.InitialPosition.Earliest
else:
pos = pulsar.InitialPosition.Latest
self.consumer = await asyncio.to_thread(
self.client.subscribe,
queue = self.queue, subscriber = self.subscriber,
schema = self.schema,
start_of_messages = self.start_of_messages
topic = self.topic,
subscription_name = self.subscriber,
schema = JsonSchema(self.schema),
initial_position = pos
)
except Exception as e:
@ -84,7 +92,7 @@ class Consumer:
await asyncio.sleep(self.reconnect_time)
continue
print(self.queue, "subscribed", flush=True)
print(self.topic, "subscribed", flush=True)
if self.metrics:
self.metrics.state("running")

36
trustgraph-base/trustgraph/base/consumer_spec.py Normal file
View file

@ -0,0 +1,36 @@
from . metrics import ConsumerMetrics
from . consumer import Consumer
from . spec import Spec
class ConsumerSpec(Spec):
def __init__(self, name, schema, handler):
self.name = name
self.schema = schema
self.handler = handler
def add(self, flow, processor, definition):
consumer_metrics = ConsumerMetrics(
flow.id, f"{flow.name}-{self.name}"
)
consumer = Consumer(
taskgroup = processor.taskgroup,
flow = flow,
client = processor.client,
topic = definition[self.name],
subscriber = flow.id,
schema = self.schema,
handler = self.handler,
metrics = consumer_metrics,
)
# Consumer handle gets access to producers and other
# metadata
consumer.id = flow.id
consumer.name = self.name
consumer.flow = flow
flow.consumer[self.name] = consumer

32
trustgraph-base/trustgraph/base/flow.py Normal file
View file

@ -0,0 +1,32 @@
import asyncio
class Flow:
def __init__(self, id, flow, processor, defn):
self.id = id
self.name = flow
self.producer = {}
# Consumers and publishers. Is this a bit untidy?
self.consumer = {}
self.setting = {}
for spec in processor.specifications:
spec.add(self, processor, defn)
async def start(self):
for c in self.consumer.values():
await c.start()
async def stop(self):
for c in self.consumer.values():
await c.stop()
def __call__(self, key):
if key in self.producer: return self.producer[key]
if key in self.consumer: return self.consumer[key]
if key in self.setting: return self.setting[key].value
return None

132
trustgraph-base/trustgraph/base/flow_processor.py
View file

@ -11,124 +11,10 @@ from .. schema import Error
from .. schema import config_request_queue, config_response_queue
from .. schema import config_push_queue
from .. log_level import LogLevel
from .. base import AsyncProcessor, Consumer, Producer, Subscriber
from . metrics import ConsumerMetrics, ProducerMetrics
from . async_processor import AsyncProcessor
from . subscriber import Subscriber
from . flow import Flow
class Spec:
pass
class SubscriberSpec(Spec):
def __init__(self, name, schema):
self.name = name
self.schema = schema
def add(self, flow, processor, definition):
subscriber_metrics = ConsumerMetrics(
flow.id, f"{flow.name}-{self.name}"
)
subscriber = Subscriber(
pulsar_client = processor.client.client,
topic = definition[self.name],
subscription = flow.id,
consumer_name = flow.id,
schema = self.schema,
)
# Put it in the consumer map, does that work?
# It means it gets start/stop call.
flow.consumer[self.name] = subscriber
if not hasattr(flow, self.name):
setattr(flow, self.name, subscriber)
class ConsumerSpec(Spec):
def __init__(self, name, schema, handler):
self.name = name
self.schema = schema
self.handler = handler
def add(self, flow, processor, definition):
consumer_metrics = ConsumerMetrics(
flow.id, f"{flow.name}-{self.name}"
)
consumer = processor.subscribe(
flow = flow,
queue = definition[self.name],
subscriber = flow.id,
schema = self.schema,
handler = self.handler,
metrics = consumer_metrics,
)
# Consumer handle gets access to producers and other
# metadata
consumer.id = flow.id
consumer.name = self.name
consumer.flow = flow
flow.consumer[self.name] = consumer
if not hasattr(flow, self.name):
setattr(flow, self.name, consumer)
class ProducerSpec(Spec):
def __init__(self, name, schema):
self.name = name
self.schema = schema
def add(self, flow, processor, definition):
producer_metrics = ProducerMetrics(
flow.id, f"{flow.name}-{self.name}"
)
producer = processor.publish(
queue = definition[self.name],
schema = self.schema,
metrics = producer_metrics,
)
flow.producer[self.name] = producer
if not hasattr(self, self.name):
setattr(flow, self.name, producer)
class SettingSpec(Spec):
def __init__(self, name):
self.name = name
def add(self, flow, processor, definition):
flow.config[self.name] = definition[self.name]
if not hasattr(flow, self.name):
setattr(flow, self.name, definition[self.name])
class Flow:
def __init__(self, id, flow, processor, defn):
self.id = id
self.name = flow
self.producer = {}
self.consumer = {}
self.setting = {}
for spec in processor.specifications:
spec.add(self, processor, defn)
async def start(self):
for c in self.consumer.values():
await c.start()
async def stop(self):
for c in self.consumer.values():
await c.stop()
# Parent class for configurable processors, configured with flows by
# the config service
class FlowProcessor(AsyncProcessor):
@ -151,18 +37,6 @@ class FlowProcessor(AsyncProcessor):
print("Service initialised.")
# Register a new consumer name
def register_consumer(self, name, schema, handler):
self.specifications.append(ConsumerSpec(name, schema, handler))
# Register a producer name
def register_producer(self, name, schema):
self.specifications.append(ProducerSpec(name, schema))
# Register a configuration variable
def register_config(self, name):
self.specifications.append(SettingSpec(name))
# Register a configuration variable
def register_specification(self, spec):
self.specifications.append(spec)

14
trustgraph-base/trustgraph/base/producer.py
View file

@ -1,11 +1,12 @@
from pulsar.schema import JsonSchema
import asyncio
class Producer:
def __init__(self, client, queue, schema, metrics=None):
def __init__(self, client, topic, schema, metrics=None):
self.client = client
self.queue = queue
self.topic = topic
self.schema = schema
self.metrics = metrics
@ -34,9 +35,12 @@ class Producer:
while self.running and self.producer is None:
try:
print("Connect publisher to", self.queue, "...", flush=True)
self.producer = self.client.publish(self.queue, self.schema)
print("Connected to", self.queue, flush=True)
print("Connect publisher to", self.topic, "...", flush=True)
self.producer = self.client.create_producer(
topic = self.topic,
schema = JsonSchema(self.schema)
)
print("Connected to", self.topic, flush=True)
except Exception as e:
print("Exception:", e, flush=True)
await asyncio.sleep(2)

25
trustgraph-base/trustgraph/base/producer_spec.py Normal file
View file

@ -0,0 +1,25 @@
from . producer import Producer
from . metrics import ProducerMetrics
from . spec import Spec
class ProducerSpec(Spec):
def __init__(self, name, schema):
self.name = name
self.schema = schema
def add(self, flow, processor, definition):
producer_metrics = ProducerMetrics(
flow.id, f"{flow.name}-{self.name}"
)
producer = Producer(
client = processor.client,
topic = definition[self.name],
schema = self.schema,
metrics = producer_metrics,
)
flow.producer[self.name] = producer

4
trustgraph-base/trustgraph/base/publisher.py
View file

@ -7,9 +7,9 @@ import pulsar
class Publisher:
def __init__(self, pulsar_client, topic, schema=None, max_size=10,
def __init__(self, client, topic, schema=None, max_size=10,
chunking_enabled=True):
self.client = pulsar_client
self.client = client
self.topic = topic
self.schema = schema
self.q = asyncio.Queue(maxsize=max_size)

24
trustgraph-base/trustgraph/base/pubsub.py
View file

@ -51,30 +51,6 @@ class PulsarClient:
if self.client:
self.client.close()
def subscribe(self, queue, subscriber, schema,
start_of_messages=False):
if start_of_messages:
pos = pulsar.InitialPosition.Earliest
else:
pos = pulsar.InitialPosition.Latest
return self.client.subscribe(
topic = queue,
subscription_name = subscriber,
consumer_type = pulsar.ConsumerType.Shared,
schema = JsonSchema(schema),
initial_position = pos,
)
def publish(self, queue, schema):
return self.client.create_producer(
topic=queue,
schema=JsonSchema(schema),
chunking_enabled=True,
)
@staticmethod
def add_args(parser):

73
trustgraph-base/trustgraph/base/request_response.py
View file

@ -1,73 +0,0 @@
import json
from pulsar.schema import JsonSchema
from .. schema import Error
from .. schema import config_request_queue, config_response_queue
from .. schema import config_push_queue
from .. log_level import LogLevel
from .. base import AsyncProcessor, Consumer, Producer
from .. base import ProcessorMetrics, ConsumerMetrics, ProducerMetrics
from . flow_processor import FlowProcessor
class RequestResponseService(FlowProcessor):
def __init__(self, **params):
super(RequestResponseService, self).__init__(**params)
self.response_schema = params.get("responsedrequest_schema")
# These can be overriden by a derived class
self.consumer_spec = [
("request", params.get("request_schema"), self.on_message)
]
self.producer_spec = [
("response", params.get("response_schema"))
]
print("Service initialised.")
async def on_message(self, message, consumer, flow):
v = message.value()
# Sender-produced ID
id = message.properties()["id"]
print(f"Handling input {id}...", flush=True)
try:
resp = await self.on_request(v, consumer, flow)
print("Send response...", flush=True)
await flow.producer["response"].send(resp, properties={"id": id})
return
except Exception as e:
print("Exception:", e, flush=True)
print("Send error response...", flush=True)
r = self.response_schema(
error=Error(
type="internal-error",
message = str(e)
)
)
await flow.producer["response"].send(r, properties={"id": id})
print("Done.", flush=True)
@staticmethod
def add_args(parser):
FlowProcessor.add_args(parser)
def run():
Processor.launch(module, __doc__)

104
trustgraph-base/trustgraph/base/request_response_spec.py Normal file
View file

@ -0,0 +1,104 @@
import uuid
import asyncio
from . subscriber import Subscriber
from . producer import Producer
from . spec import Spec
from . metrics import ConsumerMetrics, ProducerMetrics
class RequestResponse(Subscriber):
def __init__(
self, client, subscription, consumer_name,
request_topic, request_schema,
request_metrics,
response_topic, response_schema,
response_metrics,
):
super(RequestResponse, self).__init__(
client = client,
subscription = subscription,
consumer_name = consumer_name,
topic = response_topic,
schema = response_schema,
)
self.producer = Producer(
client = client,
topic = request_topic,
schema = request_schema,
metrics = request_metrics,
)
async def start(self):
await self.producer.start()
await super(RequestResponse, self).start()
async def stop(self):
await self.producer.stop()
await super(RequestResponse, self).stop()
async def request(self, req, timeout=300):
id = str(uuid.uuid4())
q = await self.subscribe(id)
try:
await self.producer.send(
req,
properties={"id": id}
)
resp = await asyncio.wait_for(
q.get(),
timeout=timeout
)
return resp
except Exception as e:
print("Exception:", e)
raise e
finally:
await self.unsubscribe(id)
# This deals with the request/response case. The caller needs to
# use another service in request/response mode. Uses two topics:
# - we send on the request topic as a producer
# - we receive on the response topic as a subscriber
class RequestResponseSpec(Spec):
def __init__(
self, request_name, request_schema, response_name, response_schema
):
self.request_name = request_name
self.request_schema = request_schema
self.response_name = response_name
self.response_schema = response_schema
def add(self, flow, processor, definition):
producer_metrics = ProducerMetrics(
flow.id, f"{flow.name}-{self.response_name}"
)
rr = RequestResponse(
client = processor.client,
subscription = flow.id,
consumer_name = flow.id,
request_topic = definition[self.request_name],
request_schema = self.request_schema,
request_metrics = producer_metrics,
response_topic = definition[self.response_name],
response_schema = self.response_schema,
response_metrics = None,
)
flow.consumer[self.request_name] = rr

19
trustgraph-base/trustgraph/base/setting_spec.py Normal file
View file

@ -0,0 +1,19 @@
from . spec import Spec
class Setting:
def __init__(self, value):
self.value = value
async def start():
pass
async def stop():
pass
class SettingSpec(Spec):
def __init__(self, name):
self.name = name
def add(self, flow, processor, definition):
flow.config[self.name] = Setting(definition[self.name])

4
trustgraph-base/trustgraph/base/spec.py Normal file
View file

@ -0,0 +1,4 @@
class Spec:
pass

4
trustgraph-base/trustgraph/base/subscriber.py
View file

@ -6,9 +6,9 @@ import time
class Subscriber:
def __init__(self, pulsar_client, topic, subscription, consumer_name,
def __init__(self, client, topic, subscription, consumer_name,
schema=None, max_size=100):
self.client = pulsar_client
self.client = client
self.topic = topic
self.subscription = subscription
self.consumer_name = consumer_name

30
trustgraph-base/trustgraph/base/subscriber_spec.py Normal file
View file

@ -0,0 +1,30 @@
from . metrics import ConsumerMetrics
from . subscriber import Subscriber
from . spec import Spec
class SubscriberSpec(Spec):
def __init__(self, name, schema):
self.name = name
self.schema = schema
def add(self, flow, processor, definition):
# FIXME: Metrics not used
subscriber_metrics = ConsumerMetrics(
flow.id, f"{flow.name}-{self.name}"
)
subscriber = Subscriber(
client = processor.client,
topic = definition[self.name],
subscription = flow.id,
consumer_name = flow.id,
schema = self.schema,
)
# Put it in the consumer map, does that work?
# It means it gets start/stop call.
flow.consumer[self.name] = subscriber

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

@ -10,7 +10,7 @@ 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 FlowProcessor
from ... base import FlowProcessor, ConsumerSpec, ProducerSpec
default_ident = "chunker"
@ -41,15 +41,19 @@ class Processor(FlowProcessor):
is_separator_regex=False,
)
self.register_consumer(
name = "input",
schema = TextDocument,
handler = self.on_message,
self.register_specification(
ConsumerSpec(
name = "input",
schema = TextDocument,
handler = self.on_message,
)
)
self.register_producer(
name = "output",
schema = Chunk,
self.register_specification(
ProducerSpec(
name = "output",
schema = Chunk,
)
)
print("Chunker initialised", flush=True)
@ -76,7 +80,7 @@ class Processor(FlowProcessor):
id=consumer.id, flow=consumer.flow
).observe(len(chunk.page_content))
await flow.producer["output"].send(r)
await flow("output").send(r)
print("Done.", flush=True)

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

@ -40,15 +40,19 @@ class Processor(FlowProcessor):
chunk_overlap=chunk_overlap,
)
self.register_consumer(
name = "input",
schema = TextDocument,
handler = self.on_message,
self.register_specification(
ConsumerSpec(
name = "input",
schema = TextDocument,
handler = self.on_message,
)
)
self.register_producer(
name = "output",
schema = Chunk,
self.register_specification(
ProducerSpec(
name = "output",
schema = Chunk,
)
)
print("Chunker initialised", flush=True)
@ -75,7 +79,7 @@ class Processor(FlowProcessor):
id=consumer.id, flow=consumer.flow
).observe(len(chunk.page_content))
await flow.producer["output"].send(r)
await flow("output").send(r)
print("Done.", flush=True)

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

@ -14,6 +14,7 @@ from trustgraph.base import AsyncProcessor, Consumer, Producer
from . config import Configuration
from ... base import ProcessorMetrics, ConsumerMetrics, ProducerMetrics
from ... base import Consumer, Producer
default_ident = "config-svc"
@ -42,29 +43,33 @@ class Processor(AsyncProcessor):
}
)
self.request_metrics = ConsumerMetrics(id + "-request")
self.response_metrics = ProducerMetrics(id + "-response")
self.push_metrics = ProducerMetrics(id + "-push")
request_metrics = ConsumerMetrics(id + "-request")
response_metrics = ProducerMetrics(id + "-response")
push_metrics = ProducerMetrics(id + "-push")
self.push_pub = self.publish(
queue = push_queue,
self.push_pub = Producer(
client = self.client,
topic = push_queue,
schema = ConfigPush,
metrics = self.push_metrics,
metrics = push_metrics,
)
self.response_pub = self.publish(
queue = response_queue,
self.response_pub = Producer(
client = self.client,
topic = response_queue,
schema = ConfigResponse,
metrics = self.response_metrics,
metrics = response_metrics,
)
self.subs = self.subscribe(
self.subs = Consumer(
taskgroup = self.taskgroup,
client = self.client,
flow = None,
queue = request_queue,
topic = request_queue,
subscriber = id,
schema = request_schema,
handler = self.on_message,
metrics = self.request_metrics,
metrics = request_metrics,
)
self.config = Configuration(self.push)

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

@ -11,7 +11,7 @@ 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 FlowProcessor
from ... base import FlowProcessor, ConsumerSpec, ProducerSpec
default_ident = "pdf-decoder"
@ -27,15 +27,19 @@ class Processor(FlowProcessor):
}
)
self.register_consumer(
name = "input",
schema = Document,
handler = self.on_message,
self.register_specification(
ConsumerSpec(
name = "input",
schema = Document,
handler = self.on_message,
)
)
self.register_producer(
name = "output",
schema = TextDocument,
self.register_specification(
ProducerSpec(
name = "output",
schema = TextDocument,
)
)
print("PDF inited", flush=True)
@ -67,7 +71,7 @@ class Processor(FlowProcessor):
text=page.page_content.encode("utf-8"),
)
await flow.producer["output"].send(r)
await flow("output").send(r)
print("Done.", flush=True)

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

@ -6,7 +6,6 @@ entity/context definitions for embedding.
"""
import json
import asyncio
import urllib.parse
from pulsar.schema import JsonSchema
import uuid
@ -18,7 +17,7 @@ from .... log_level import LogLevel
from .... clients.prompt_client import PromptClient
from .... rdf import TRUSTGRAPH_ENTITIES, DEFINITION, RDF_LABEL, SUBJECT_OF
from .... base import FlowProcessor, SubscriberSpec, ConsumerSpec
from .... base import FlowProcessor, RequestResponseSpec, ConsumerSpec
from .... base import ProducerSpec
DEFINITION_VALUE = Value(value=DEFINITION, is_uri=True)
@ -48,16 +47,11 @@ class Processor(FlowProcessor):
)
self.register_specification(
ProducerSpec(
name = "prompt-request",
schema = PromptRequest
)
)
self.register_specification(
SubscriberSpec(
name = "prompt-response",
schema = PromptResponse,
RequestResponseSpec(
request_name = "prompt-request",
request_schema = PromptRequest,
response_name = "prompt-response",
response_schema = PromptResponse,
)
)
@ -116,26 +110,16 @@ class Processor(FlowProcessor):
try:
q = await flow.consumer["prompt-response"].subscribe(id)
try:
await flow.producer["prompt-request"].send(
resp = await flow("prompt-request").request(
PromptRequest(
id="extract-definitions",
terms={
"text": json.dumps(chunk)
},
),
properties={"id": id}
)
)
# FIXME: hard-coded?
resp = await asyncio.wait_for(
q.get(),
timeout=600
)
print("Response", resp, flush=True)
if resp.error is not None:
@ -150,8 +134,6 @@ class Processor(FlowProcessor):
except Exception as e:
print("Prompt exception:", e, flush=True)
raise e
finally:
await flow.consumer["prompt-response"].unsubscribe(id)
triples = []
entities = []
@ -201,7 +183,7 @@ class Processor(FlowProcessor):
entities.append(ec)
await self.emit_triples(
flow.producer["triples"],
flow("triples"),
Metadata(
id=v.metadata.id,
metadata=[],
@ -212,7 +194,7 @@ class Processor(FlowProcessor):
)
await self.emit_ecs(
flow.producer["entity-contexts"],
flow("entity-contexts"),
Metadata(
id=v.metadata.id,
metadata=[],

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

@ -13,7 +13,7 @@ from .... schema import PromptRequest, PromptResponse, Error
from .... schema import TextCompletionRequest, TextCompletionResponse
from .... base import FlowProcessor
from .... base import ProducerSpec, SubscriberSpec, ConsumerSpec
from .... base import ProducerSpec, ConsumerSpec, RequestResponseSpec
from . prompt_manager import PromptConfiguration, Prompt, PromptManager
@ -43,16 +43,11 @@ class Processor(FlowProcessor):
)
self.register_specification(
ProducerSpec(
name = "text-completion-request",
schema = TextCompletionRequest
)
)
self.register_specification(
SubscriberSpec(
name = "text-completion-response",
schema = TextCompletionResponse,
RequestResponseSpec(
request_name = "text-completion-request",
request_schema = TextCompletionRequest,
response_name = "text-completion-response",
response_schema = TextCompletionResponse,
)
)
@ -139,24 +134,15 @@ class Processor(FlowProcessor):
print(f"Handling kind {kind}...", flush=True)
q = await flow.consumer["text-completion-response"].subscribe(id)
async def llm(system, prompt):
print(system, flush=True)
print(prompt, flush=True)
await flow.producer["text-completion-request"].send(
resp = await flow("text-completion-request").request(
TextCompletionRequest(
system=system, prompt=prompt
),
properties={"id": id}
)
# FIXME: hard-coded?
resp = await asyncio.wait_for(
q.get(),
timeout=600
)
try:
@ -170,8 +156,6 @@ class Processor(FlowProcessor):
except Exception as e:
print("Invocation exception:", e, flush=True)
raise e
finally:
await flow.consumer["text-completion-response"].unsubscribe(id)
print(resp, flush=True)
@ -185,7 +169,7 @@ class Processor(FlowProcessor):
error=None,
)
await flow.response.send(r, properties={"id": id})
await flow("response").send(r, properties={"id": id})
return
@ -200,7 +184,7 @@ class Processor(FlowProcessor):
error=None,
)
await flow.response.send(r, properties={"id": id})
await flow("response").send(r, properties={"id": id})
return
@ -218,7 +202,7 @@ class Processor(FlowProcessor):
response=None,
)
await flow.response.send(r, properties={"id": id})
await flow("response").send(r, properties={"id": id})
except Exception as e: