Bare bones API gateway

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))
+
+############################################################################
+

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()
+

trustgraph-flow/setup.py

@@ -58,6 +58,7 @@ setuptools.setup(
         "google-generativeai",
         "ibis",
         "jsonschema",
+        "aiohttp",
     ],
     scripts=[
         "scripts/agent-manager-react",