service

#!/usr/bin/env python3

import _jsonnet as j
from aiohttp import web
import aiohttp
import asyncio
import logging
import os
import json
import yaml

from trustgraph.clients.embeddings_client import EmbeddingsClient
from trustgraph.clients.graph_rag_client import GraphRagClient
from trustgraph.clients.prompt_client import PromptClient
from trustgraph.clients.triples_query_client import TriplesQueryClient
from trustgraph.clients.llm_client import LlmClient

logger = logging.getLogger("service")
logger.setLevel(logging.INFO)

class Api:
    def __init__(self, config=None):

        if not config:
            self.config = {}

        self.config = config

        pulsar_host = config.get("pulsar-host", "pulsar://localhost:6650")

        self.port = int(self.config.get("port", "8081"))
        self.app = web.Application(middlewares=[])

        self.app.router.add_route('GET', '/ws', self.websocket_handler)
        self.app.add_routes([web.get("/{tail:.*}", self.everything_else)])

        self.embed = EmbeddingsClient(pulsar_host=pulsar_host)
        self.llm = LlmClient(pulsar_host=pulsar_host)

        print("Initialised")

    async def websocket_handler(self, request):

        print('Websocket connection starting')
        ws = aiohttp.web.WebSocketResponse()
        await ws.prepare(request)
        print('Websocket connection ready')

        handlers = {
            "message": self.handle_message
        }

        async for msg in ws:

            print(msg)

            if msg.type == aiohttp.WSMsgType.TEXT:

                try:
                    obj = json.loads(msg.data)
                except:
                    await ws.send_json({
                        "type": "error",
                        "body": {
                            "message": f"JSON parse fail",
                        }
                    })
                    continue

                try:

                    if "type" not in obj:
                        raise RuntimeError("Message has no type field")

                    if obj["type"] == "close":
                        await ws.close()
                        break

                    if "body" not in obj:
                        raise RuntimeError("Message has no body field")

                    if obj["type"] not in handlers:
                        raise RuntimeError("Don't understand that type")

                    handler = handlers.get(obj["type"])
                    resp = await handler(obj["body"])

                    # Code above must either set 'resp' or raise an exception

                    await ws.send_json(resp)
                    continue

                except Exception as e:
                    await ws.send_json({
                        "type": "error",
                        "body": {
                            "message": str(e),
                        }
                    })
                    continue

        print('Websocket connection closed')
        return ws

    async def handle_message(self, obj):

        if "text" not in obj:
            raise RuntimeError("Message has no message field")

        q = obj["text"]
        print(">", q)
        resp = self.llm.request(q)

        print("<", resp)

        return {
            "type": "message",
            "body": {
                "role": "ai",
                "text": resp,
            },
        }

    async def everything_else(self, request):

        print(">", request.path)
        if ".." in request.path:
            return web.HTTPNotFound()

        if request.path == "/":
            with open("dist/index.html", "r") as f:
                return web.Response(
                    text=f.read(), content_type="text/html"
                )

        if request.path == "/api/patterns":
            return web.Response(
                text=self.patterns, content_type="application/json"
            )

        if request.path.endswith(".css"):
            with open("dist" + request.path, "r") as f:
                data = f.read()
            return web.Response(
                text=data, content_type="text/css"
            )

        if request.path.endswith(".js"):
            with open("dist" + request.path, "r") as f:
                data = f.read()
            return web.Response(
                text=data, content_type="text/javascript"
            )

        if request.path.endswith(".html"):
            with open("dist" + request.path, "r") as f:
                data = f.read()
            return web.Response(
                text=data, content_type="text/html"
            )

        return web.HTTPNotFound()

    async def generate(self, request):

        print("Generate...")
        config = await request.text()

        print(config)

        config = config.encode("utf-8")

        gen = Generator(config)

        with open("./templates/config-loader.jsonnet", "r") as f:
            wrapper = f.read()
        processed = gen.process(wrapper)

        return web.Response(
            text=yaml.dump(processed), content_type = "text/plain"
        )

    def run(self):

        web.run_app(self.app, port=self.port)

config = {
    "pulsar-host": "pulsar://localhost:6650"
}

a = Api(config=config)
a.run()