Spaces:
Running
Running
import gradio as gr | |
import inspect | |
import os | |
import requests | |
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" | |
) | |
class get_related_docs(Resource): | |
def get(self, knowledge_base): | |
args = query_parser.parse_args() | |
kb = kb_list[knowledge_base]() | |
print(kb) | |
print(args.subject, args.where_str, args.limit) | |
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) | |