add api

README.md

@@ -1,5 +1,5 @@
 ---
-title: GPTs Myscale Backend
+title: GPTs Myscale Backend RestAPI
 emoji: 📚
 colorFrom: gray
 colorTo: gray

app.py

@@ -0,0 +1,91 @@
+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"
+)
+
+
+@api.route(
+    "/get_related_docs/<string:knowledge_base>",
+    doc={
+        "description": (
+            "Get some related papers.\nYou should use schema here:\n\n"
+            "CREATE TABLE ArXiv (\n"
+            "    `id` String,\n"
+            "    `abstract` String, -- abstract of the paper. avoid using this column to do LIKE match\n"
+            "    `pubdate` DateTime, \n"
+            "    `title` String, -- title of the paper\n"
+            "    `categories` Array(String), -- arxiv category of the paper\n"
+            "    `authors` Array(String), -- authors of the paper\n"
+            "    `comment` String, -- extra comments of the paper\n"
+            "ORDER BY id\n\n"
+            "CREATE TABLE Wikipedia (\n"
+            "    `id` String,\n"
+            "    `text` String, -- abstract of the wiki page. avoid using this column to do LIKE match\n"
+            "    `title` String, -- title of the paper\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."
+        ),
+    },
+)
+@api.param("knowledge_base", "Knowledge base used to query. Must be one of ['wiki', 'arxiv']")
+class get_related_docs(Resource):
+    @api.expect(query_parser)
+    @api.marshal_with(query_result)
+    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, debug=True)

funcs.py

@@ -0,0 +1,72 @@
+
+from typing import List, Tuple
+import clickhouse_connect
+from sentence_transformers import SentenceTransformer
+from InstructorEmbedding import INSTRUCTOR
+
+
+emb_wiki = SentenceTransformer("sentence-transformers/paraphrase-multilingual-mpnet-base-v2")
+emb_arxiv = INSTRUCTOR('hkunlp/instructor-xl')
+
+class ArXivKnowledgeBase:
+    def __init__(self, embedding: SentenceTransformer) -> None:    
+        self.db = clickhouse_connect.get_client(
+            host='msc-4a9e710a.us-east-1.aws.staging.myscale.cloud',
+            port=443,
+            username='chatdata',
+            password='myscale_rocks'
+        )
+        self.embedding: SentenceTransformer = embedding
+        self.table: str = 'default.ChatArXiv'
+        self.embedding_col = "vector"
+        self.must_have_cols: List[str] = ['id', 'abstract', 'authors', 'categories', 'comment', 'title', 'pubdate']
+    
+
+    def __call__(self, subject: str, where_str: str = None, limit: int = 5) -> Tuple[str, int]:
+        q_emb = self.embedding.encode(subject).tolist()
+        q_emb_str = ",".join(map(str, q_emb))
+        if where_str:
+            where_str = f"WHERE {where_str}"
+        else:
+            where_str = ""
+
+        q_str = f"""
+            SELECT dist, {','.join(self.must_have_cols)}
+            FROM {self.table}
+            {where_str}
+            ORDER BY distance({self.embedding_col}, [{q_emb_str}]) 
+                AS dist ASC
+            LIMIT {limit}
+            """
+
+        docs = [r for r in self.db.query(q_str).named_results()]
+        return '\n'.join([str(d) for d in docs]), len(docs)
+    
+class WikiKnowledgeBase(ArXivKnowledgeBase):
+    def __init__(self, embedding: SentenceTransformer) -> None:    
+        super().__init__(embedding)
+        self.table: str = 'wiki.Wikipedia'
+        self.embedding_col = "emb"
+        self.must_have_cols: List[str] = ['text', 'title', 'views', 'url']
+
+
+if __name__ == '__main__':
+    # kb = ArXivKnowledgeBase(embedding=emb_arxiv)
+    kb = WikiKnowledgeBase(embedding=emb_wiki)
+    d = kb("When did Steven Jobs die?", "", 5)
+    print(d)
+
+
+d = {"components": {
+    "schemas": {
+    "type": "object",
+    "properties": {
+        "todos":{
+          "type": "array",
+          "items":{"type": "string"},
+          "description": "The list of todos.",
+         }
+    }
+   }
+  }
+  }

requirements.txt

@@ -0,0 +1,3 @@
+clickhouse_connect
+flask
+flask-restx