mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-17 09:11:03 +02:00
Prompt template working
This commit is contained in:
Cyber MacGeddon
parent
44b1d7e508
commit
b110af41fb
15 changed files with 380 additions and 526 deletions
4
Makefile
4
Makefile
|
|
@ -65,8 +65,8 @@ some-containers:
|
|||
-t ${CONTAINER_BASE}/trustgraph-base:${VERSION} .
|
||||
${DOCKER} build -f containers/Containerfile.flow \
|
||||
-t ${CONTAINER_BASE}/trustgraph-flow:${VERSION} .
|
||||
${DOCKER} build -f containers/Containerfile.vertexai \
|
||||
-t ${CONTAINER_BASE}/trustgraph-vertexai:${VERSION} .
|
||||
# ${DOCKER} build -f containers/Containerfile.vertexai \
|
||||
# -t ${CONTAINER_BASE}/trustgraph-vertexai:${VERSION} .
|
||||
|
||||
basic-containers: update-package-versions
|
||||
${DOCKER} build -f containers/Containerfile.base \
|
||||
|
|
|
|||
|
|
@ -11,11 +11,9 @@ llm = LlmClient(
|
|||
)
|
||||
|
||||
system = "You are a lovely assistant."
|
||||
prompt="Write a funny limerick about a llama"
|
||||
prompt="what is 2 + 2 == 5"
|
||||
|
||||
resp = llm.request(system, prompt)
|
||||
|
||||
print(resp)
|
||||
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -3,13 +3,18 @@
|
|||
import pulsar
|
||||
from trustgraph.clients.prompt_client import PromptClient
|
||||
|
||||
p = PromptClient(pulsar_host="pulsar://localhost:6650")
|
||||
p = PromptClient(
|
||||
pulsar_host="pulsar://localhost:6650",
|
||||
input_queue="non-persistent://tg/request/prompt:default",
|
||||
output_queue="non-persistent://tg/response/prompt:default",
|
||||
subscriber="test1",
|
||||
)
|
||||
|
||||
question = """What is the square root of 16?"""
|
||||
|
||||
resp = p.request(
|
||||
id="question",
|
||||
terms = {
|
||||
variables = {
|
||||
"question": question
|
||||
}
|
||||
)
|
||||
|
|
|
|||
|
|
@ -6,7 +6,8 @@ from . producer import Producer
|
|||
from . publisher import Publisher
|
||||
from . subscriber import Subscriber
|
||||
from . metrics import ProcessorMetrics, ConsumerMetrics, ProducerMetrics
|
||||
from . flow_processor import FlowProcessor
|
||||
from . flow_processor import FlowProcessor, SubscriberSpec, ProducerSpec
|
||||
from . flow_processor import ConsumerSpec, SettingSpec
|
||||
from . request_response import RequestResponseService
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -30,6 +30,8 @@ class Consumer:
|
|||
|
||||
self.metrics = metrics
|
||||
|
||||
self.consumer = None
|
||||
|
||||
def __del__(self):
|
||||
self.running = False
|
||||
|
||||
|
|
@ -39,30 +41,63 @@ class Consumer:
|
|||
async def stop(self):
|
||||
|
||||
self.running = False
|
||||
await self.task
|
||||
|
||||
async def start(self):
|
||||
|
||||
self.running = True
|
||||
|
||||
self.consumer = self.client.subscribe(
|
||||
self.queue, self.subscriber, self.schema,
|
||||
start_of_messages = self.start_of_messages,
|
||||
)
|
||||
|
||||
print("Subscribed.", flush=True)
|
||||
|
||||
# Puts it in the stopped state, the run thread should set running
|
||||
if self.metrics:
|
||||
self.metrics.state("stopped")
|
||||
|
||||
self.task = self.taskgroup.create_task(self.run())
|
||||
|
||||
print("Subscriber started", flush=True)
|
||||
|
||||
async def run(self):
|
||||
|
||||
if self.metrics:
|
||||
self.metrics.state("running")
|
||||
self.metrics.state("stopped")
|
||||
|
||||
while self.running:
|
||||
|
||||
try:
|
||||
|
||||
print(self.queue, "subscribing...", flush=True)
|
||||
|
||||
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
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
|
||||
print("Exception:", e, flush=True)
|
||||
await asyncio.sleep(2)
|
||||
continue
|
||||
|
||||
print(self.queue, "subscribed", flush=True)
|
||||
|
||||
if self.metrics:
|
||||
self.metrics.state("running")
|
||||
|
||||
try:
|
||||
|
||||
await self.consume()
|
||||
|
||||
if self.metrics:
|
||||
self.metrics.state("stopped")
|
||||
|
||||
except Exception as e:
|
||||
|
||||
print("Exception:", e, flush=True)
|
||||
self.consumer.close()
|
||||
self.consumer = None
|
||||
await asyncio.sleep(2)
|
||||
continue
|
||||
|
||||
async def consume(self):
|
||||
|
||||
while self.running:
|
||||
|
||||
|
|
@ -79,7 +114,7 @@ class Consumer:
|
|||
expiry = time.time() + self.rate_limit_timeout
|
||||
|
||||
# This loop is for retry on rate-limit / resource limits
|
||||
while True:
|
||||
while self.running:
|
||||
|
||||
if time.time() > expiry:
|
||||
|
||||
|
|
@ -122,7 +157,6 @@ class Consumer:
|
|||
|
||||
print("TooManyRequests: will retry...", flush=True)
|
||||
|
||||
# FIXME
|
||||
if self.metrics:
|
||||
self.metrics.rate_limit()
|
||||
|
||||
|
|
@ -145,7 +179,3 @@ class Consumer:
|
|||
|
||||
# Break out of retry loop, processes next message
|
||||
break
|
||||
|
||||
if self.metrics:
|
||||
self.metrics.state("stopped")
|
||||
|
||||
|
|
|
|||
|
|
@ -1,187 +0,0 @@
|
|||
|
||||
import asyncio
|
||||
from pulsar.schema import JsonSchema
|
||||
import pulsar
|
||||
from prometheus_client import Histogram, Info, Counter, Enum
|
||||
import time
|
||||
|
||||
from . base_processor import BaseProcessor
|
||||
from .. exceptions import TooManyRequests
|
||||
|
||||
default_rate_limit_retry = 10
|
||||
default_rate_limit_timeout = 7200
|
||||
|
||||
class Consumer(BaseProcessor):
|
||||
|
||||
def __init__(self, **params):
|
||||
|
||||
if not hasattr(__class__, "state_metric"):
|
||||
__class__.state_metric = Enum(
|
||||
'processor_state', 'Processor state',
|
||||
states=['starting', 'running', 'stopped']
|
||||
)
|
||||
__class__.state_metric.state('starting')
|
||||
|
||||
__class__.state_metric.state('starting')
|
||||
|
||||
super(Consumer, self).__init__(**params)
|
||||
|
||||
self.subscriber = params.get("subscriber")
|
||||
self.input_schema = params.get("input_schema")
|
||||
|
||||
self.rate_limit_retry = params.get(
|
||||
"rate_limit_retry", default_rate_limit_retry
|
||||
)
|
||||
self.rate_limit_timeout = params.get(
|
||||
"rate_limit_timeout", default_rate_limit_timeout
|
||||
)
|
||||
|
||||
if self.input_schema == None:
|
||||
raise RuntimeError("input_schema must be specified")
|
||||
|
||||
if not hasattr(__class__, "request_metric"):
|
||||
__class__.request_metric = Histogram(
|
||||
'request_latency', 'Request latency (seconds)',
|
||||
["flow"]
|
||||
)
|
||||
|
||||
if not hasattr(__class__, "pubsub_metric"):
|
||||
__class__.pubsub_metric = Info(
|
||||
'pubsub', 'Pub/sub configuration',
|
||||
["flow"]
|
||||
)
|
||||
|
||||
if not hasattr(__class__, "processing_metric"):
|
||||
__class__.processing_metric = Counter(
|
||||
'processing_count', 'Processing count',
|
||||
["flow", "status"]
|
||||
)
|
||||
|
||||
if not hasattr(__class__, "rate_limit_metric"):
|
||||
__class__.rate_limit_metric = Counter(
|
||||
'rate_limit_count', 'Rate limit event count',
|
||||
["flow"]
|
||||
)
|
||||
|
||||
# __class__.pubsub_metric.info({
|
||||
# "input_queue": self.input_queue,
|
||||
# "subscriber": self.subscriber,
|
||||
# "input_schema": self.input_schema.__name__,
|
||||
# "rate_limit_retry": str(self.rate_limit_retry),
|
||||
# "rate_limit_timeout": str(self.rate_limit_timeout),
|
||||
# })
|
||||
|
||||
# self.consumer = self.client.subscribe(
|
||||
# self.input_queue, self.subscriber,
|
||||
# consumer_type=pulsar.ConsumerType.Shared,
|
||||
# schema=JsonSchema(self.input_schema),
|
||||
# )
|
||||
|
||||
self.config_queues = {}
|
||||
self.config_handlers = [self.on_config_pubsub]
|
||||
|
||||
print("Initialised consumer.", flush=True)
|
||||
|
||||
async def on_config_pubsub(self, version, config):
|
||||
|
||||
print("Configuring using config", version)
|
||||
|
||||
flows = config.get("flows", {}).get(self.ident, {})
|
||||
|
||||
print(flows)
|
||||
|
||||
print("Configuration completed OK")
|
||||
|
||||
async def run(self):
|
||||
|
||||
__class__.state_metric.state('running')
|
||||
|
||||
while True:
|
||||
|
||||
await asyncio.sleep(1)
|
||||
|
||||
continue
|
||||
|
||||
msg = await asyncio.to_thread(self.consumer.receive)
|
||||
|
||||
expiry = time.time() + self.rate_limit_timeout
|
||||
|
||||
# This loop is for retry on rate-limit / resource limits
|
||||
while True:
|
||||
|
||||
if time.time() > expiry:
|
||||
|
||||
print("Gave up waiting for rate-limit retry", flush=True)
|
||||
|
||||
# Message failed to be processed, this causes it to
|
||||
# be retried
|
||||
self.consumer.negative_acknowledge(msg)
|
||||
|
||||
__class__.processing_metric.labels(status="error").inc()
|
||||
|
||||
# Break out of retry loop, processes next message
|
||||
break
|
||||
|
||||
try:
|
||||
|
||||
with __class__.request_metric.time():
|
||||
await self.handle(msg)
|
||||
|
||||
# Acknowledge successful processing of the message
|
||||
self.consumer.acknowledge(msg)
|
||||
|
||||
__class__.processing_metric.labels(status="success").inc()
|
||||
|
||||
# Break out of retry loop
|
||||
break
|
||||
|
||||
except TooManyRequests:
|
||||
|
||||
print("TooManyRequests: will retry...", flush=True)
|
||||
|
||||
__class__.rate_limit_metric.inc()
|
||||
|
||||
# Sleep
|
||||
time.sleep(self.rate_limit_retry)
|
||||
|
||||
# Contine from retry loop, just causes a reprocessing
|
||||
continue
|
||||
|
||||
except Exception as e:
|
||||
|
||||
print("Exception:", e, flush=True)
|
||||
|
||||
# Message failed to be processed, this causes it to
|
||||
# be retried
|
||||
self.consumer.negative_acknowledge(msg)
|
||||
|
||||
__class__.processing_metric.labels(status="error").inc()
|
||||
|
||||
# Break out of retry loop, processes next message
|
||||
break
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser, default_input_queue, default_subscriber):
|
||||
|
||||
BaseProcessor.add_args(parser)
|
||||
|
||||
parser.add_argument(
|
||||
'-s', '--subscriber',
|
||||
default=default_subscriber,
|
||||
help=f'Queue subscriber name (default: {default_subscriber})'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--rate-limit-retry',
|
||||
type=int,
|
||||
default=default_rate_limit_retry,
|
||||
help=f'Rate limit retry (default: {default_rate_limit_retry})'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--rate-limit-timeout',
|
||||
type=int,
|
||||
default=default_rate_limit_timeout,
|
||||
help=f'Rate limit timeout (default: {default_rate_limit_timeout})'
|
||||
)
|
||||
|
||||
|
|
@ -1,62 +0,0 @@
|
|||
|
||||
from pulsar.schema import JsonSchema
|
||||
import pulsar
|
||||
from prometheus_client import Histogram, Info, Counter, Enum
|
||||
import time
|
||||
|
||||
from . consumer import Consumer
|
||||
from .. exceptions import TooManyRequests
|
||||
|
||||
class ConsumerProducer(Consumer):
|
||||
|
||||
def __init__(self, **params):
|
||||
|
||||
super(ConsumerProducer, self).__init__(**params)
|
||||
|
||||
self.output_queue = params.get("output_queue")
|
||||
self.output_schema = params.get("output_schema")
|
||||
|
||||
if not hasattr(__class__, "output_metric"):
|
||||
__class__.output_metric = Counter(
|
||||
'output_count', 'Output items created'
|
||||
)
|
||||
|
||||
__class__.pubsub_metric.info({
|
||||
"input_queue": self.input_queue,
|
||||
"output_queue": self.output_queue,
|
||||
"subscriber": self.subscriber,
|
||||
"input_schema": self.input_schema.__name__,
|
||||
"output_schema": self.output_schema.__name__,
|
||||
"rate_limit_retry": str(self.rate_limit_retry),
|
||||
"rate_limit_timeout": str(self.rate_limit_timeout),
|
||||
})
|
||||
|
||||
if self.output_schema == None:
|
||||
raise RuntimeError("output_schema must be specified")
|
||||
|
||||
self.producer = self.client.create_producer(
|
||||
topic=self.output_queue,
|
||||
schema=JsonSchema(self.output_schema),
|
||||
chunking_enabled=True,
|
||||
)
|
||||
|
||||
print("Initialised consumer/producer.")
|
||||
|
||||
async def send(self, msg, properties={}):
|
||||
self.producer.send(msg, properties)
|
||||
__class__.output_metric.inc()
|
||||
|
||||
@staticmethod
|
||||
def add_args(
|
||||
parser, default_input_queue, default_subscriber,
|
||||
default_output_queue,
|
||||
):
|
||||
|
||||
Consumer.add_args(parser, default_input_queue, default_subscriber)
|
||||
|
||||
parser.add_argument(
|
||||
'-o', '--output-queue',
|
||||
default=default_output_queue,
|
||||
help=f'Output queue (default: {default_output_queue})'
|
||||
)
|
||||
|
||||
|
|
@ -11,70 +11,115 @@ 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 AsyncProcessor, Consumer, Producer, Subscriber
|
||||
from . metrics import ConsumerMetrics, ProducerMetrics
|
||||
|
||||
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.producer = {}
|
||||
self.consumer = {}
|
||||
self.config = {}
|
||||
self.id = id
|
||||
self.name = flow
|
||||
|
||||
for spec in processor.config_spec:
|
||||
name = spec
|
||||
self.config[name] = defn[name]
|
||||
self.producer = {}
|
||||
self.consumer = {}
|
||||
self.setting = {}
|
||||
|
||||
if not hasattr(self, name):
|
||||
setattr(self, name, defn[name])
|
||||
|
||||
for spec in processor.producer_spec:
|
||||
|
||||
name, schema = spec
|
||||
|
||||
producer_metrics = ProducerMetrics(
|
||||
id, f"{flow}-{name}"
|
||||
)
|
||||
|
||||
producer = processor.publish(
|
||||
queue = defn[name],
|
||||
schema = schema,
|
||||
metrics = producer_metrics,
|
||||
)
|
||||
|
||||
self.producer[name] = producer
|
||||
|
||||
if not hasattr(self, name):
|
||||
setattr(self, name, producer)
|
||||
|
||||
for spec in processor.consumer_spec:
|
||||
|
||||
name, schema, handler = spec
|
||||
|
||||
consumer_metrics = ConsumerMetrics(
|
||||
id, f"{flow}-{name}"
|
||||
)
|
||||
|
||||
consumer = processor.subscribe(
|
||||
flow = self,
|
||||
queue = defn[name],
|
||||
subscriber = id,
|
||||
schema = schema,
|
||||
handler = handler,
|
||||
metrics = consumer_metrics,
|
||||
)
|
||||
|
||||
# Consumer handle gets access to producers and other
|
||||
# metadata
|
||||
consumer.id = id
|
||||
consumer.name = name
|
||||
consumer.flow = self
|
||||
|
||||
self.consumer[name] = consumer
|
||||
|
||||
if not hasattr(self, name):
|
||||
setattr(self, name, consumer)
|
||||
for spec in processor.specifications:
|
||||
spec.add(self, processor, defn)
|
||||
|
||||
async def start(self):
|
||||
for c in self.consumer.values():
|
||||
|
|
@ -101,29 +146,26 @@ class FlowProcessor(AsyncProcessor):
|
|||
|
||||
# These can be overriden by a derived class:
|
||||
|
||||
# Consumer specification, array of ("name", SchemaType, handler)
|
||||
self.consumer_spec = []
|
||||
|
||||
# Producer specification, array of ("name", SchemaType)
|
||||
self.producer_spec = []
|
||||
|
||||
# Configuration specification, collects some flow variables from
|
||||
# config, array of "name"
|
||||
self.config_spec = []
|
||||
# Array of specifications: ConsumerSpec, ProducerSpec, SettingSpec
|
||||
self.specifications = []
|
||||
|
||||
print("Service initialised.")
|
||||
|
||||
# Register a new consumer name
|
||||
def register_consumer(self, name, schema, handler):
|
||||
self.consumer_spec.append((name, schema, handler))
|
||||
self.specifications.append(ConsumerSpec(name, schema, handler))
|
||||
|
||||
# Register a producer name
|
||||
def register_producer(self, name, schema):
|
||||
self.producer_spec.append((name, schema))
|
||||
self.specifications.append(ProducerSpec(name, schema))
|
||||
|
||||
# Register a configuration variable
|
||||
def register_config(self, name):
|
||||
self.config_spec.append(name)
|
||||
self.specifications.append(SettingSpec(name))
|
||||
|
||||
# Register a configuration variable
|
||||
def register_specification(self, spec):
|
||||
self.specifications.append(spec)
|
||||
|
||||
# Start processing for a new flow
|
||||
async def start_flow(self, flow, defn):
|
||||
|
|
|
|||
|
|
@ -1,4 +1,6 @@
|
|||
|
||||
import asyncio
|
||||
|
||||
class Producer:
|
||||
|
||||
def __init__(self, client, queue, schema, metrics=None):
|
||||
|
|
@ -6,19 +8,58 @@ class Producer:
|
|||
self.queue = queue
|
||||
self.schema = schema
|
||||
|
||||
self.producer = self.client.publish(self.queue, self.schema)
|
||||
|
||||
self.metrics = metrics
|
||||
|
||||
self.running = True
|
||||
self.producer = None
|
||||
|
||||
def __del__(self):
|
||||
|
||||
self.running = False
|
||||
|
||||
if hasattr(self, "producer"):
|
||||
self.producer.close()
|
||||
if self.producer:
|
||||
self.producer.close()
|
||||
|
||||
async def start(self):
|
||||
self.running = True
|
||||
|
||||
async def stop(self):
|
||||
self.running = False
|
||||
|
||||
async def send(self, msg, properties={}):
|
||||
|
||||
if self.metrics:
|
||||
self.metrics.inc()
|
||||
if not self.running: return
|
||||
|
||||
self.producer.send(msg, properties)
|
||||
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)
|
||||
except Exception as e:
|
||||
print("Exception:", e, flush=True)
|
||||
await asyncio.sleep(2)
|
||||
|
||||
if not self.running: break
|
||||
|
||||
while self.running:
|
||||
|
||||
try:
|
||||
|
||||
await asyncio.to_thread(
|
||||
self.producer.send,
|
||||
msg, properties
|
||||
)
|
||||
|
||||
if self.metrics:
|
||||
self.metrics.inc()
|
||||
|
||||
# Delivery success, break out of loop
|
||||
break
|
||||
|
||||
except Exception as e:
|
||||
print("Exception:", e, flush=True)
|
||||
self.producer.close()
|
||||
self.producer = None
|
||||
|
||||
|
|
|
|||
|
|
@ -1,57 +0,0 @@
|
|||
|
||||
from pulsar.schema import JsonSchema
|
||||
from prometheus_client import Info, Counter
|
||||
|
||||
from . base_processor import BaseProcessor
|
||||
|
||||
class Producer(BaseProcessor):
|
||||
|
||||
def __init__(self, **params):
|
||||
|
||||
output_queue = params.get("output_queue")
|
||||
output_schema = params.get("output_schema")
|
||||
|
||||
if not hasattr(__class__, "output_metric"):
|
||||
__class__.output_metric = Counter(
|
||||
'output_count', 'Output items created'
|
||||
)
|
||||
|
||||
if not hasattr(__class__, "pubsub_metric"):
|
||||
__class__.pubsub_metric = Info(
|
||||
'pubsub', 'Pub/sub configuration'
|
||||
)
|
||||
|
||||
__class__.pubsub_metric.info({
|
||||
"output_queue": output_queue,
|
||||
"output_schema": output_schema.__name__,
|
||||
})
|
||||
|
||||
super(Producer, self).__init__(**params)
|
||||
|
||||
if output_schema == None:
|
||||
raise RuntimeError("output_schema must be specified")
|
||||
|
||||
self.producer = self.client.create_producer(
|
||||
topic=output_queue,
|
||||
schema=JsonSchema(output_schema),
|
||||
chunking_enabled=True,
|
||||
)
|
||||
|
||||
async def send(self, msg, properties={}):
|
||||
self.producer.send(msg, properties)
|
||||
__class__.output_metric.inc()
|
||||
|
||||
@staticmethod
|
||||
def add_args(
|
||||
parser, default_input_queue, default_subscriber,
|
||||
default_output_queue,
|
||||
):
|
||||
|
||||
BaseProcessor.add_args(parser)
|
||||
|
||||
parser.add_argument(
|
||||
'-o', '--output-queue',
|
||||
default=default_output_queue,
|
||||
help=f'Output queue (default: {default_output_queue})'
|
||||
)
|
||||
|
||||
|
|
@ -1,8 +1,7 @@
|
|||
|
||||
import queue
|
||||
import asyncio
|
||||
import time
|
||||
import pulsar
|
||||
import threading
|
||||
|
||||
class Publisher:
|
||||
|
||||
|
|
@ -11,22 +10,22 @@ class Publisher:
|
|||
self.client = pulsar_client
|
||||
self.topic = topic
|
||||
self.schema = schema
|
||||
self.q = queue.Queue(maxsize=max_size)
|
||||
self.q = asyncio.Queue(maxsize=max_size)
|
||||
self.chunking_enabled = chunking_enabled
|
||||
self.running = True
|
||||
|
||||
def start(self):
|
||||
self.task = threading.Thread(target=self.run)
|
||||
self.task.start()
|
||||
async def start(self):
|
||||
self.task = asyncio.create_task(self.run())
|
||||
await self.task.start()
|
||||
|
||||
def stop(self):
|
||||
async def stop(self):
|
||||
self.running = False
|
||||
|
||||
def join(self):
|
||||
self.stop()
|
||||
self.task.join()
|
||||
async def join(self):
|
||||
await self.stop()
|
||||
await self.task
|
||||
|
||||
def run(self):
|
||||
async def run(self):
|
||||
|
||||
while self.running:
|
||||
|
||||
|
|
@ -40,8 +39,13 @@ class Publisher:
|
|||
while self.running:
|
||||
|
||||
try:
|
||||
id, item = self.q.get(timeout=0.5)
|
||||
except queue.Empty:
|
||||
id, item = asyncio.wait_for(
|
||||
self.q.get(),
|
||||
timeout=0.5
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
continue
|
||||
except asyncio.QueueEmpty:
|
||||
continue
|
||||
|
||||
if id:
|
||||
|
|
@ -55,7 +59,6 @@ class Publisher:
|
|||
# If handler drops out, sleep a retry
|
||||
time.sleep(2)
|
||||
|
||||
def send(self, id, msg):
|
||||
self.q.put((id, msg))
|
||||
async def send(self, id, item):
|
||||
await self.q.put((id, item))
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
|
||||
import queue
|
||||
import pulsar
|
||||
import threading
|
||||
from pulsar.schema import JsonSchema
|
||||
import asyncio
|
||||
import _pulsar
|
||||
import time
|
||||
|
||||
class Subscriber:
|
||||
|
|
@ -16,41 +16,50 @@ class Subscriber:
|
|||
self.q = {}
|
||||
self.full = {}
|
||||
self.max_size = max_size
|
||||
self.lock = threading.Lock()
|
||||
self.lock = asyncio.Lock()
|
||||
self.running = True
|
||||
|
||||
def start(self):
|
||||
self.task = threading.Thread(target=self.run)
|
||||
self.task.start()
|
||||
|
||||
def stop(self):
|
||||
async def __del__(self):
|
||||
self.running = False
|
||||
|
||||
def join(self):
|
||||
self.task.join()
|
||||
async def start(self):
|
||||
self.task = asyncio.create_task(self.run())
|
||||
|
||||
def run(self):
|
||||
async def stop(self):
|
||||
self.running = False
|
||||
|
||||
print("SUBSCRIBER LOOP KICKS IN", flush=True)
|
||||
async def join(self):
|
||||
await self.stop()
|
||||
await self.task
|
||||
|
||||
async def run(self):
|
||||
|
||||
while self.running:
|
||||
|
||||
try:
|
||||
|
||||
consumer = self.client.subscribe(
|
||||
topic=self.topic,
|
||||
subscription_name=self.subscription,
|
||||
consumer_name=self.consumer_name,
|
||||
schema=self.schema,
|
||||
topic = self.topic,
|
||||
subscription_name = self.subscription,
|
||||
consumer_name = self.consumer_name,
|
||||
schema = JsonSchema(self.schema),
|
||||
)
|
||||
|
||||
print("SUBSCRIBER RUNNING...", flush=True)
|
||||
print("Subscriber running...", flush=True)
|
||||
|
||||
while self.running:
|
||||
|
||||
msg = consumer.receive()
|
||||
|
||||
print("GOT MESSAGE...", flush=True)
|
||||
try:
|
||||
msg = await asyncio.to_thread(
|
||||
consumer.receive,
|
||||
timeout_millis=2000
|
||||
)
|
||||
except _pulsar.Timeout:
|
||||
continue
|
||||
except Exception as e:
|
||||
print("Exception:", e, flush=True)
|
||||
print(type(e))
|
||||
raise e
|
||||
|
||||
# Acknowledge successful reception of the message
|
||||
consumer.acknowledge(msg)
|
||||
|
|
@ -62,57 +71,68 @@ class Subscriber:
|
|||
|
||||
value = msg.value()
|
||||
|
||||
with self.lock:
|
||||
async with self.lock:
|
||||
|
||||
# FIXME: Hard-coded timeouts
|
||||
|
||||
if id in self.q:
|
||||
|
||||
try:
|
||||
# FIXME: Timeout means data goes missing
|
||||
self.q[id].put(value, timeout=0.5)
|
||||
except:
|
||||
pass
|
||||
await asyncio.wait_for(
|
||||
self.q[id].put(value),
|
||||
timeout=2
|
||||
)
|
||||
except Exception as e:
|
||||
print("Q Put:", e, flush=True)
|
||||
|
||||
for q in self.full.values():
|
||||
try:
|
||||
# FIXME: Timeout means data goes missing
|
||||
q.put(value, timeout=0.5)
|
||||
except:
|
||||
pass
|
||||
await asyncio.wait_for(
|
||||
q.put(value),
|
||||
timeout=2
|
||||
)
|
||||
except Exception as e:
|
||||
print("Q Put:", e, flush=True)
|
||||
|
||||
except Exception as e:
|
||||
print("Subscriber exception:", e, flush=True)
|
||||
|
||||
consumer.close()
|
||||
|
||||
# If handler drops out, sleep a retry
|
||||
time.sleep(2)
|
||||
|
||||
def subscribe(self, id):
|
||||
async def subscribe(self, id):
|
||||
|
||||
with self.lock:
|
||||
async with self.lock:
|
||||
|
||||
q = queue.Queue(maxsize=self.max_size)
|
||||
q = asyncio.Queue(maxsize=self.max_size)
|
||||
self.q[id] = q
|
||||
|
||||
return q
|
||||
|
||||
def unsubscribe(self, id):
|
||||
async def unsubscribe(self, id):
|
||||
|
||||
with self.lock:
|
||||
async with self.lock:
|
||||
|
||||
if id in self.q:
|
||||
# self.q[id].shutdown(immediate=True)
|
||||
del self.q[id]
|
||||
|
||||
def subscribe_all(self, id):
|
||||
async def subscribe_all(self, id):
|
||||
|
||||
with self.lock:
|
||||
async with self.lock:
|
||||
|
||||
q = queue.Queue(maxsize=self.max_size)
|
||||
q = asyncio.Queue(maxsize=self.max_size)
|
||||
self.full[id] = q
|
||||
|
||||
return q
|
||||
|
||||
def unsubscribe_all(self, id):
|
||||
async def unsubscribe_all(self, id):
|
||||
|
||||
with self.lock:
|
||||
async with self.lock:
|
||||
|
||||
if id in self.full:
|
||||
# self.full[id].shutdown(immediate=True)
|
||||
|
|
|
|||
|
|
@ -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,7 @@ class PromptManager:
|
|||
|
||||
return json.loads(json_str)
|
||||
|
||||
def invoke(self, id, input):
|
||||
async def invoke(self, id, input, llm):
|
||||
|
||||
if id not in self.prompts:
|
||||
raise RuntimeError("ID invalid")
|
||||
|
|
@ -68,9 +65,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
|
||||
|
|
@ -83,10 +78,9 @@ class PromptManager:
|
|||
except:
|
||||
raise RuntimeError("JSON parse fail")
|
||||
|
||||
print(obj, flush=True)
|
||||
if self.prompts[id].schema:
|
||||
try:
|
||||
print(self.prompts[id].schema)
|
||||
print(self.prompts[id].schema, flush=True)
|
||||
validate(instance=obj, schema=self.prompts[id].schema)
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Schema validation fail: {e}")
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@
|
|||
Language service abstracts prompt engineering from LLM.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import re
|
||||
|
||||
|
|
@ -10,63 +11,73 @@ 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 .... clients.llm_client import LlmClient
|
||||
from .... base import RequestResponseService
|
||||
from .... base import FlowProcessor, SubscriberSpec, ConsumerSpec
|
||||
from .... base import ProducerSpec
|
||||
|
||||
from . prompt_manager import PromptConfiguration, Prompt, PromptManager
|
||||
|
||||
module = "prompt"
|
||||
default_subscriber = module
|
||||
default_ident = "prompt"
|
||||
|
||||
class Processor(RequestResponseService):
|
||||
class Processor(FlowProcessor):
|
||||
|
||||
def __init__(self, **params):
|
||||
|
||||
id = params.get("id")
|
||||
subscriber = params.get("subscriber", default_subscriber)
|
||||
|
||||
# Config key for prompts
|
||||
self.config_key = params.get("config_type", "prompt")
|
||||
|
||||
super(Processor, self).__init__(
|
||||
**params | {
|
||||
"id": id,
|
||||
"subscriber": subscriber,
|
||||
"request_schema": PromptRequest,
|
||||
"response_schema": PromptResponse,
|
||||
}
|
||||
)
|
||||
|
||||
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.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(
|
||||
ProducerSpec(
|
||||
name = "text-completion-request",
|
||||
schema = TextCompletionRequest
|
||||
)
|
||||
)
|
||||
|
||||
self.llm = Llm(self.llm)
|
||||
self.register_specification(
|
||||
SubscriberSpec(
|
||||
name = "text-completion-response",
|
||||
schema = TextCompletionResponse,
|
||||
)
|
||||
)
|
||||
|
||||
self.config_handlers.append(self.on_config)
|
||||
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)
|
||||
|
||||
|
|
@ -100,7 +111,6 @@ class Processor(RequestResponseService):
|
|||
)
|
||||
|
||||
self.manager = PromptManager(
|
||||
self.llm,
|
||||
PromptConfiguration(
|
||||
system,
|
||||
{},
|
||||
|
|
@ -115,7 +125,7 @@ class Processor(RequestResponseService):
|
|||
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()
|
||||
|
||||
|
|
@ -127,7 +137,7 @@ class Processor(RequestResponseService):
|
|||
|
||||
try:
|
||||
|
||||
print(v.terms)
|
||||
print(v.terms, flush=True)
|
||||
|
||||
input = {
|
||||
k: json.loads(v)
|
||||
|
|
@ -135,9 +145,40 @@ class Processor(RequestResponseService):
|
|||
}
|
||||
|
||||
print(f"Handling kind {kind}...", flush=True)
|
||||
print(input, flush=True)
|
||||
|
||||
resp = self.manager.invoke(kind, input)
|
||||
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(
|
||||
TextCompletionRequest(
|
||||
system=system, prompt=prompt
|
||||
),
|
||||
properties={"id": id}
|
||||
)
|
||||
|
||||
# FIXME: hard-coded?
|
||||
resp = await asyncio.wait_for(
|
||||
q.get(),
|
||||
timeout=600
|
||||
)
|
||||
|
||||
try:
|
||||
return resp.response
|
||||
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
|
||||
finally:
|
||||
await flow.consumer["text-completion-response"].unsubscribe(id)
|
||||
|
||||
if isinstance(resp, str):
|
||||
|
||||
|
|
@ -150,7 +191,7 @@ class Processor(RequestResponseService):
|
|||
error=None,
|
||||
)
|
||||
|
||||
await self.send(r, properties={"id": id})
|
||||
await flow.response.send(r, properties={"id": id})
|
||||
|
||||
return
|
||||
|
||||
|
|
@ -165,13 +206,13 @@ class Processor(RequestResponseService):
|
|||
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)
|
||||
|
||||
|
|
@ -183,11 +224,11 @@ class Processor(RequestResponseService):
|
|||
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)
|
||||
|
||||
|
|
@ -204,22 +245,7 @@ class Processor(RequestResponseService):
|
|||
@staticmethod
|
||||
def add_args(parser):
|
||||
|
||||
RequestResponseService.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',
|
||||
|
|
@ -229,5 +255,5 @@ class Processor(RequestResponseService):
|
|||
|
||||
def run():
|
||||
|
||||
Processor.launch(module, __doc__)
|
||||
Processor.launch(default_ident, __doc__)
|
||||
|
||||
|
|
|
|||
|
|
@ -181,7 +181,7 @@ class Processor(RequestResponseService):
|
|||
@staticmethod
|
||||
def add_args(parser):
|
||||
|
||||
RequestResponseService.add_args(parser, default_subscriber)
|
||||
RequestResponseService.add_args(parser)
|
||||
|
||||
parser.add_argument(
|
||||
'-m', '--model',
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue