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

Working for LLM + prompt

Browse source
This commit is contained in:
Cyber MacGeddon 2024-11-20 17:44:00 +00:00
parent a8e37037d2
commit 04a6059f33
2 changed files with 87 additions and 50 deletions
Download patch file Download diff file Expand all files Collapse all files

19
test-api
View file

@ -19,14 +19,20 @@ resp = requests.post(
json=input,
)
print(resp.json()["response"])
sys.exit(0)
resp = resp.json()
if "error" in resp:
print(f"Error: {resp['error']}")
sys.exit(1)
print(resp["response"])
############################################################################
input = {
"id": "question",
"variables": {
"question": "Add 12 and 14, and then make a poem about llamas which incorporates that number. Then write a joke about llamas."
"question": "Write a joke about llamas."
}
}
@ -35,8 +41,13 @@ resp = requests.post(
json=input,
)
print(resp.json()["text"])
if "error" in resp:
print(f"Error: {resp['error']}")
sys.exit(1)
print(resp["response"])
sys.exit(0)
############################################################################
input = {

118
trustgraph-flow/scripts/api-gateway
View file

@ -19,33 +19,34 @@ from trustgraph.schema import TextCompletionRequest, TextCompletionResponse
from trustgraph.schema import text_completion_request_queue
from trustgraph.schema import text_completion_response_queue
from trustgraph.schema import PromptRequest, PromptResponse
from trustgraph.schema import prompt_request_queue
from trustgraph.schema import prompt_response_queue
logger = logging.getLogger("api")
logger.setLevel(logging.INFO)
pulsar_host = "pulsar://localhost:6650"
TIME_OUT = 600
class Publisher:
def __init__(self, pulsar_host, topic, schema=None, max_size=10):
self.pulsar_host = pulsar_host
self.topic = topic
self.schema = schema
self.q = asyncio.Queue(maxsize=max_size)
async def run(self):
async with aiopulsar.connect(self.pulsar_host) as client:
async with client.create_producer(
topic=self.topic,
schema=self.schema,
) as producer:
while True:
id, item = await self.q.get()
await producer.send(item, { "id": id })
print("message out")
async def send(self, id, msg):
await self.q.put((id, msg))
@ -54,33 +55,26 @@ class Subscriber:
def __init__(self, pulsar_host, topic, subscription, consumer_name,
schema=None, max_size=10):
self.pulsar_host = pulsar_host
self.topic = topic
self.subscription = subscription
self.consumer_name = consumer_name
self.schema = schema
self.q = {}
async def run(self):
async with aiopulsar.connect(pulsar_host) as client:
async with client.subscribe(
topic=self.topic,
subscription_name=self.subscription,
consumer_name=self.consumer_name,
schema=self.schema,
) as consumer:
while True:
msg = await consumer.receive()
print("message in")
id = msg.properties()["id"]
value = msg.value()
if id in self.q:
await self.q[id].put(value)
@ -91,7 +85,6 @@ class Subscriber:
async def unsubscribe(self, id):
if id in self.q:
# self.q[id].shutdown()
del self.q[id]
class Api:
@ -101,15 +94,26 @@ class Api:
self.port = int(config.get("port", "8088"))
self.app = web.Application(middlewares=[])
self.llm_out = Publisher(
pulsar_host, text_completion_request_queue,
schema=JsonSchema(TextCompletionRequest)
)
self.llm_in = Subscriber(
pulsar_host, text_completion_response_queue,
"api-gateway", "api-gateway",
JsonSchema(TextCompletionResponse)
)
self.llm_out = Publisher(
pulsar_host, text_completion_request_queue,
schema=JsonSchema(TextCompletionRequest)
self.prompt_out = Publisher(
pulsar_host, prompt_request_queue,
schema=JsonSchema(PromptRequest)
)
self.prompt_in = Subscriber(
pulsar_host, prompt_response_queue,
"api-gateway", "api-gateway",
JsonSchema(PromptResponse)
)
self.app.add_routes([
@ -119,8 +123,6 @@ class Api:
async def llm(self, request):
logger.info("LLM...")
id = str(uuid.uuid4())
try:
@ -137,66 +139,90 @@ class Api:
)
)
resp = await q.get()
try:
resp = await asyncio.wait_for(q.get(), TIME_OUT)
except:
raise RuntimeError("Timeout waiting for response")
await self.llm_in.unsubscribe(id)
if resp.error:
return web.json_response(
{ "error": resp.error.message }
)
return web.json_response(
{
"response": resp.response,
},
content_type="application/json",
{ "response": resp.response }
)
except Exception as e:
logging.error(f"Exception: {e}")
return web.HTTPInternalServerError()
return web.json_response(
{ "error": str(e) }
)
finally:
await self.llm_in.unsubscribe(id)
async def prompt(self, request):
logger.info("LLM...")
id = str(uuid.uuid4())
try:
data = await request.json()
try:
id = data["id"]
variables = data["variables"]
except:
return web.HTTPBadRequest()
q = await self.prompt_in.subscribe(id)
resp = self.prompt_client.request(
id=id, variables=variables
terms = {
k: json.dumps(v)
for k, v in data["variables"].items()
}
await self.prompt_out.send(
id,
PromptRequest(
id=data["id"],
terms=terms
)
)
if isinstance(resp, str):
try:
resp = await asyncio.wait_for(q.get(), TIME_OUT)
except:
raise RuntimeError("Timeout waiting for response")
if resp.error:
return web.json_response(
{
"text": resp,
},
content_type="application/json",
{ "error": resp.error.message }
)
else:
if resp.object:
return web.json_response(
{
"object": resp,
},
content_type="application/json",
{ "object": resp.object }
)
return web.json_response(
{ "text": resp.text }
)
except Exception as e:
logging.error(f"Exception: {e}")
return web.HTTPInternalServerError()
return web.json_response(
{ "error": str(e) }
)
finally:
await self.llm_in.unsubscribe(id)
async def app_factory(self):
self.llm_pub_task = asyncio.create_task(self.llm_in.run())
self.llm_sub_task = asyncio.create_task(self.llm_out.run())
self.prompt_pub_task = asyncio.create_task(self.prompt_in.run())
self.prompt_sub_task = asyncio.create_task(self.prompt_out.run())
return self.app
def run(self):