Fix API gateway integration, added to templates (#165)

This commit is contained in:
cybermaggedon 2024-11-20 20:56:23 +00:00 committed by GitHub
parent 92b84441eb
commit ba6d6c13af
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 81 additions and 24 deletions

View file

@ -5,9 +5,51 @@ local prompt = import "prompt-template.jsonnet";
{ {
"api-gateway-port":: 8088,
"api-gateway-timeout":: 600,
"chunk-size":: 250, "chunk-size":: 250,
"chunk-overlap":: 15, "chunk-overlap":: 15,
"api-gateway" +: {
create:: function(engine)
local port = $["api-gateway-port"];
local container =
engine.container("api-gateway")
.with_image(images.trustgraph)
.with_command([
"api-gateway",
"-p",
url.pulsar,
"--timeout",
std.toString($["api-gateway-timeout"]),
"--port",
std.toString(port),
])
.with_limits("0.5", "256M")
.with_reservations("0.1", "256M")
.with_port(8000, 8000, "metrics")
.with_port(port, port, "api");
local containerSet = engine.containers(
"api-gateway", [ container ]
);
local service =
engine.internalService(containerSet)
.with_port(8000, 8000, "metrics")
.with_port(port, port, "api");
engine.resources([
containerSet,
service,
])
},
"chunker" +: { "chunker" +: {
create:: function(engine) create:: function(engine)

View file

@ -1,10 +1,17 @@
#!/usr/bin/env python3 #!/usr/bin/env python3
# FIXME: Subscribes to Pulsar unnecessarily, should only do it when there
# are active listeners
# FIXME: Connection errors in publishers / subscribers cause those threads
# to fail and are not failed or retried
import asyncio import asyncio
from aiohttp import web from aiohttp import web
import json import json
import logging import logging
import uuid import uuid
import os
import pulsar import pulsar
from pulsar.asyncio import Client from pulsar.asyncio import Client
@ -42,7 +49,7 @@ from trustgraph.schema import embeddings_response_queue
logger = logging.getLogger("api") logger = logging.getLogger("api")
logger.setLevel(logging.INFO) logger.setLevel(logging.INFO)
pulsar_host = "pulsar://localhost:6650" pulsar_host = os.getenv("PULSAR_HOST", "pulsar://pulsar:6650")
TIME_OUT = 600 TIME_OUT = 600
class Publisher: class Publisher:
@ -54,15 +61,18 @@ class Publisher:
self.q = asyncio.Queue(maxsize=max_size) self.q = asyncio.Queue(maxsize=max_size)
async def run(self): async def run(self):
async with aiopulsar.connect(self.pulsar_host) as client: try:
async with client.create_producer( async with aiopulsar.connect(self.pulsar_host) as client:
topic=self.topic, async with client.create_producer(
schema=self.schema, topic=self.topic,
) as producer: schema=self.schema,
while True: ) as producer:
id, item = await self.q.get() while True:
await producer.send(item, { "id": id }) id, item = await self.q.get()
# print("message out") await producer.send(item, { "id": id })
# print("message out")
except Exception as e:
print("Exception:", e, flush=True)
async def send(self, id, msg): async def send(self, id, msg):
await self.q.put((id, msg)) await self.q.put((id, msg))
@ -79,20 +89,23 @@ class Subscriber:
self.q = {} self.q = {}
async def run(self): async def run(self):
async with aiopulsar.connect(pulsar_host) as client: try:
async with client.subscribe( async with aiopulsar.connect(pulsar_host) as client:
topic=self.topic, async with client.subscribe(
subscription_name=self.subscription, topic=self.topic,
consumer_name=self.consumer_name, subscription_name=self.subscription,
schema=self.schema, consumer_name=self.consumer_name,
) as consumer: schema=self.schema,
while True: ) as consumer:
msg = await consumer.receive() while True:
# print("message in", self.topic) msg = await consumer.receive()
id = msg.properties()["id"] # print("message in", self.topic)
value = msg.value() id = msg.properties()["id"]
if id in self.q: value = msg.value()
await self.q[id].put(value) if id in self.q:
await self.q[id].put(value)
except Exception as e:
print("Exception:", e, flush=True)
async def subscribe(self, id): async def subscribe(self, id):
q = asyncio.Queue() q = asyncio.Queue()

View file

@ -59,8 +59,10 @@ setuptools.setup(
"ibis", "ibis",
"jsonschema", "jsonschema",
"aiohttp", "aiohttp",
"aiopulsar-py",
], ],
scripts=[ scripts=[
"scripts/api-gateway",
"scripts/agent-manager-react", "scripts/agent-manager-react",
"scripts/chunker-recursive", "scripts/chunker-recursive",
"scripts/chunker-token", "scripts/chunker-token",