diff --git a/test-api b/test-api new file mode 100755 index 00000000..c68bbb7b --- /dev/null +++ b/test-api @@ -0,0 +1,57 @@ +#!/usr/bin/env python3 + +import requests +import json +import sys + +url = "http://localhost:8088/api/v1/" + +############################################################################ + +input = { + "system": "Respond in French. Use long word, form of numbers, no digits", +# "prompt": "Add 2 and 12" + "prompt": "Add 12 and 14, and then make a poem about llamas which incorporates that number. Then write a joke about llamas" +} + +resp = requests.post( + f"{url}text-completion", + json=input, +) + +print(resp.json()["response"]) +sys.exit(0) +############################################################################ + +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." + } +} + +resp = requests.post( + f"{url}prompt", + json=input, +) + +print(resp.json()["text"]) + +############################################################################ + +input = { + "id": "extract-definitions", + "variables": { + "text": "A cat is a large mammal." + } +} + +resp = requests.post( + f"{url}prompt", + json=input, +) + +print(json.dumps(resp.json()["object"], indent=4)) + +############################################################################ + diff --git a/trustgraph-flow/scripts/api-gateway b/trustgraph-flow/scripts/api-gateway new file mode 100755 index 00000000..d449e21d --- /dev/null +++ b/trustgraph-flow/scripts/api-gateway @@ -0,0 +1,207 @@ +#!/usr/bin/env python3 + +import asyncio +from aiohttp import web +import json +import logging +import uuid + +import pulsar +from pulsar.asyncio import Client +from pulsar.schema import JsonSchema +import _pulsar +import aiopulsar + +from trustgraph.clients.llm_client import LlmClient +from trustgraph.clients.prompt_client import PromptClient + +from trustgraph.schema import TextCompletionRequest, TextCompletionResponse +from trustgraph.schema import text_completion_request_queue +from trustgraph.schema import text_completion_response_queue + +logger = logging.getLogger("api") +logger.setLevel(logging.INFO) + +pulsar_host = "pulsar://localhost:6650" + +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 }) + + async def send(self, id, msg): + await self.q.put((id, msg)) + +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() + + id = msg.properties()["id"] + value = msg.value() + + if id in self.q: + await self.q[id].put(value) + + async def subscribe(self, id): + q = asyncio.Queue() + self.q[id] = q + return q + + async def unsubscribe(self, id): + if id in self.q: +# self.q[id].shutdown() + del self.q[id] + +class Api: + + def __init__(self, **config): + + self.port = int(config.get("port", "8088")) + self.app = web.Application(middlewares=[]) + + 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.app.add_routes([ + web.post("/api/v1/text-completion", self.llm), + web.post("/api/v1/prompt", self.prompt), + ]) + + async def llm(self, request): + + logger.info("LLM...") + + id = str(uuid.uuid4()) + + try: + + data = await request.json() + + q = await self.llm_in.subscribe(id) + + await self.llm_out.send( + id, + TextCompletionRequest( + system=data["system"], + prompt=data["prompt"] + ) + ) + + resp = await q.get() + + await self.llm_in.unsubscribe(id) + + return web.json_response( + { + "response": resp.response, + }, + content_type="application/json", + ) + + except Exception as e: + logging.error(f"Exception: {e}") + return web.HTTPInternalServerError() + + async def prompt(self, request): + + logger.info("LLM...") + + try: + + data = await request.json() + + try: + id = data["id"] + variables = data["variables"] + except: + return web.HTTPBadRequest() + + resp = self.prompt_client.request( + id=id, variables=variables + ) + + if isinstance(resp, str): + + return web.json_response( + { + "text": resp, + }, + content_type="application/json", + ) + + else: + + return web.json_response( + { + "object": resp, + }, + content_type="application/json", + ) + + except Exception as e: + logging.error(f"Exception: {e}") + return web.HTTPInternalServerError() + + 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()) + + return self.app + + def run(self): + web.run_app(self.app_factory(), port=self.port) + +a = Api() +a.run() + diff --git a/trustgraph-flow/setup.py b/trustgraph-flow/setup.py index 8b46b2d2..fb929759 100644 --- a/trustgraph-flow/setup.py +++ b/trustgraph-flow/setup.py @@ -58,6 +58,7 @@ setuptools.setup( "google-generativeai", "ibis", "jsonschema", + "aiohttp", ], scripts=[ "scripts/agent-manager-react",