app.py

import inspect
import os
import json
from io import BytesIO
from typing import List, Type

from flask import Flask, jsonify, render_template, request, send_file
from flask_restx import Resource, Api, fields
from funcs import emb_wiki, emb_arxiv, WikiKnowledgeBase, ArXivKnowledgeBase

app = Flask(__name__)
api = Api(
    app,
    version="0.1",
    terms_url="https://myscale.com/terms/",
    contact_email="support@myscale.com",
    title="MyScale Open Knowledge Base",
    description="An API to get relevant page from MyScale Open Knowledge Base",
)

query_result = api.model(
    "QueryResult",
    {
        "documents": fields.String,
        "num_retrieved": fields.Integer,
    },
)

kb_list = {
    "wiki": lambda: WikiKnowledgeBase(embedding=emb_wiki),
    "arxiv": lambda: ArXivKnowledgeBase(embedding=emb_arxiv),
}

query_parser = api.parser()
query_parser.add_argument(
    "subject",
    required=True,
    type=str,
    help="a sentence or phrase describes the subject you want to query.",
)
query_parser.add_argument(
    "where_str", required=True, type=str, help="a sql-like where string to build filter"
)
query_parser.add_argument(
    "limit", required=False, type=int, default=4, help="desired number of retrieved documents"
)


@api.route(
    "/get_related_arxiv",
    doc={
        "description": (
            "Get some related papers.\nYou should use schema here:\n\n"
            "CREATE TABLE ArXiv (\n"
            "    `id` String,\n"
            "    `abstract` String, \n"
            "    `pubdate` DateTime, \n"
            "    `title` String, \n"
            "    `categories` Array(String), -- arxiv category\n"
            "    `authors` Array(String),\n"
            "    `comment` String,\n"
            "ORDER BY id\n\n"
        ),
    },
)
class get_related_arxiv(Resource):
    @api.expect(query_parser)
    @api.marshal_with(query_result)
    @api.doc(id='get_related_arxiv')
    def get(self):
        args = query_parser.parse_args()
        kb = kb_list['arxiv']()
        docs, num_docs = kb(args.subject, args.where_str, args.limit)
        return {"documents": docs, "num_retrieved": num_docs}


@api.route(
    "/get_related_wiki",
    doc={
        "description": (
            "Get some related wiki pages.\nYou should use schema here:\n\n"
            "CREATE TABLE Wikipedia (\n"
            "    `id` String,\n"
            "    `text` String,\n"
            "    `title` String,\n"
            "    `view` Float32,\n"
            "    `url` String, -- URL to this wiki page\n"
            "ORDER BY id\n\n"
            "You should avoid using LIKE on long text columns."
        ),
    },
)
class get_related_wiki(Resource):
    @api.expect(query_parser)
    @api.marshal_with(query_result)
    @api.doc(id='get_related_wiki')
    def get(self):
        args = query_parser.parse_args()
        kb = kb_list['wiki']()
        docs, num_docs = kb(args.subject, args.where_str, args.limit)
        return {"documents": docs, "num_retrieved": num_docs}


if __name__ == "__main__":
    # print(json.dumps(api.__schema__))
    app.run(host="0.0.0.0", port=7860)