ohmu/main/reranker-adapter/adapter.py

"""OWUI (Cohere-stiil) -> TEI rerank adapter.

OWUI saadab:  {"model":..,"query":..,"documents":[...],"top_n":N}
              ootab: {"results":[{"index":i,"relevance_score":s}, ...]}
TEI ootab:    {"query":..,"texts":[...]}
              tagastab: [{"index":i,"score":s}, ...]
"""
import json
import urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

TEI = "http://reranker:80/rerank"


class Handler(BaseHTTPRequestHandler):
    server_version = "rerank-adapter/1.0"

    def _send(self, code, obj):
        body = json.dumps(obj).encode()
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self):
        self._send(200, {"ok": True, "tei": TEI})

    def do_POST(self):
        try:
            n = int(self.headers.get("Content-Length") or 0)
            body = json.loads(self.rfile.read(n) or b"{}")
        except Exception as e:
            return self._send(400, {"error": "bad json: %s" % e})

        docs = body.get("documents") or body.get("texts") or []
        if isinstance(docs, str):
            docs = [docs]
        docs = [d if isinstance(d, str) else json.dumps(d) for d in docs]
        if not docs:
            return self._send(200, {"results": []})

        req = urllib.request.Request(
            TEI,
            data=json.dumps({"query": body.get("query", ""), "texts": docs}).encode(),
            headers={"Content-Type": "application/json"},
        )
        try:
            out = json.load(urllib.request.urlopen(req, timeout=600))
        except Exception as e:
            return self._send(500, {"error": "tei: %s" % e})

        if isinstance(out, dict):
            out = out.get("results") or out.get("data") or []
        results = []
        for o in out:
            try:
                results.append({"index": int(o["index"]), "relevance_score": float(o["score"])})
            except Exception:
                continue
        self._send(200, {"results": results})

    def log_message(self, *a):
        pass


if __name__ == "__main__":
    ThreadingHTTPServer(("0.0.0.0", 80), Handler).serve_forever()