Update server.py

server.py

@@ -1,101 +1,39 @@
-from functools import wraps
 from flask import (
     Flask,
     jsonify,
     request,
     render_template_string,
     abort,
-    send_from_directory,
-    send_file,
 )
 from flask_cors import CORS
-import markdown
-import argparse
-from transformers import AutoTokenizer, AutoProcessor, pipeline
-from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM
-from transformers import BlipForConditionalGeneration, GPT2Tokenizer
 import unicodedata
-import torch
+import markdown
 import time
 import os
 import gc
-from PIL import Image
 import base64
 from io import BytesIO
 from random import randint
-import webuiapi
 import hashlib
+import chromadb
+import posthog
+from chromadb.config import Settings
+from sentence_transformers import SentenceTransformer
+from werkzeug.middleware.proxy_fix import ProxyFix
+import argparse
+from transformers import AutoTokenizer, AutoProcessor, pipeline
+from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM
+from transformers import BlipForConditionalGeneration, GPT2Tokenizer
+from PIL import Image
+import webuiapi
 from constants import *
 from colorama import Fore, Style, init as colorama_init
 
 colorama_init()
 
-
-class SplitArgs(argparse.Action):
-    def __call__(self, parser, namespace, values, option_string=None):
-        setattr(
-            namespace, self.dest, values.replace('"', "").replace("'", "").split(",")
-        )
-
-
-# Script arguments
-parser = argparse.ArgumentParser(
-    prog="TavernAI Extras", description="Web API for transformers models"
-)
-parser.add_argument(
-    "--port", type=int, help="Specify the port on which the application is hosted"
-)
-parser.add_argument(
-    "--listen", action="store_true", help="Host the app on the local network"
-)
-parser.add_argument(
-    "--share", action="store_true", help="Share the app on CloudFlare tunnel"
-)
-parser.add_argument("--cpu", action="store_true", help="Run the models on the CPU")
-parser.add_argument("--summarization-model", help="Load a custom summarization model")
 parser.add_argument(
     "--classification-model", help="Load a custom text classification model"
 )
-parser.add_argument("--captioning-model", help="Load a custom captioning model")
-parser.add_argument(
-    "--keyphrase-model", help="Load a custom keyphrase extraction model"
-)
-parser.add_argument("--prompt-model", help="Load a custom prompt generation model")
-parser.add_argument("--embedding-model", help="Load a custom text embedding model")
-
-sd_group = parser.add_mutually_exclusive_group()
-
-local_sd = sd_group.add_argument_group("sd-local")
-local_sd.add_argument("--sd-model", help="Load a custom SD image generation model")
-local_sd.add_argument("--sd-cpu", help="Force the SD pipeline to run on the CPU")
-
-remote_sd = sd_group.add_argument_group("sd-remote")
-remote_sd.add_argument(
-    "--sd-remote", action="store_true", help="Use a remote backend for SD"
-)
-remote_sd.add_argument(
-    "--sd-remote-host", type=str, help="Specify the host of the remote SD backend"
-)
-remote_sd.add_argument(
-    "--sd-remote-port", type=int, help="Specify the port of the remote SD backend"
-)
-remote_sd.add_argument(
-    "--sd-remote-ssl", action="store_true", help="Use SSL for the remote SD backend"
-)
-remote_sd.add_argument(
-    "--sd-remote-auth",
-    type=str,
-    help="Specify the username:password for the remote SD backend (if required)",
-)
-
-parser.add_argument(
-    "--enable-modules",
-    action=SplitArgs,
-    default=[],
-    help="Override a list of enabled modules",
-)
-
-args = parser.parse_args()
 
 port = 7860
 host = "0.0.0.0"
@@ -111,31 +49,17 @@ classification_model = (
     else DEFAULT_CLASSIFICATION_MODEL
 )
 
-modules = (
-    args.enable_modules if args.enable_modules and len(args.enable_modules) > 0 else []
-)
 
-if len(modules) == 0:
-    print(
-        f"{Fore.RED}{Style.BRIGHT}You did not select any modules to run! Choose them by adding an --enable-modules option"
-    )
-    print(f"Example: --enable-modules=caption,summarize{Style.RESET_ALL}")
+embedding_model = 'sentence-transformers/all-mpnet-base-v2'
 
-# Models init
-device_string = "cuda:0" if torch.cuda.is_available() and not args.cpu else "cpu"
-device = torch.device(device_string)
-torch_dtype = torch.float32 if device_string == "cpu" else torch.float16
+print("Initializing a text summarization model...")
 
-if "summarize" in modules:
-    print("Initializing a text summarization model...")
-    summarization_tokenizer = AutoTokenizer.from_pretrained(summarization_model)
-    summarization_transformer = AutoModelForSeq2SeqLM.from_pretrained(
-        summarization_model, torch_dtype=torch_dtype
-    ).to(device)
+summarization_tokenizer = AutoTokenizer.from_pretrained(summarization_model)
+summarization_transformer = AutoModelForSeq2SeqLM.from_pretrained(
+summarization_model, torch_dtype=torch_dtype).to(device)
 
-if "classify" in modules:
-    print("Initializing a sentiment classification pipeline...")
-    classification_pipe = pipeline(
+print("Initializing a sentiment classification pipeline...")
+classification_pipe = pipeline(
         "text-classification",
         model=classification_model,
         top_k=None,
@@ -143,40 +67,34 @@ if "classify" in modules:
         torch_dtype=torch_dtype,
     )
 
-if "chromadb" in modules:
-    print("Initializing ChromaDB")
-    import chromadb
-    import posthog
-    from chromadb.config import Settings
-    from sentence_transformers import SentenceTransformer
 
-    # disable chromadb telemetry
-    posthog.capture = lambda *args, **kwargs: None
-    chromadb_client = chromadb.Client(Settings(anonymized_telemetry=False))
-    chromadb_embedder = SentenceTransformer(embedding_model)
-    chromadb_embed_fn = chromadb_embedder.encode
 
+print("Initializing ChromaDB")
+
+device_string = "cpu"
+device = torch.device(device_string)
+torch_dtype = torch.float32 if device_string == "cpu" else torch.float16
+
+
+
+# disable chromadb telemetry
+posthog.capture = lambda *args, **kwargs: None
+chromadb_client = chromadb.Client(Settings(anonymized_telemetry=False))
+chromadb_embedder = SentenceTransformer(embedding_model)
+chromadb_embed_fn = chromadb_embedder.encode
 
 # Flask init
 app = Flask(__name__)
 CORS(app)  # allow cross-domain requests
 app.config["MAX_CONTENT_LENGTH"] = 100 * 1024 * 1024
 
+app.wsgi_app = ProxyFix(
+    app.wsgi_app, x_for=2, x_proto=1, x_host=1, x_prefix=1
+)
 
-def require_module(name):
-    def wrapper(fn):
-        @wraps(fn)
-        def decorated_view(*args, **kwargs):
-            if name not in modules:
-                abort(403, "Module is disabled by config")
-            return fn(*args, **kwargs)
-
-        return decorated_view
-
-    return wrapper
-
+def get_real_ip():
+    return request.remote_addr
 
-# AI stuff
 def classify_text(text: str) -> list:
     output = classification_pipe(
         text,
@@ -243,7 +161,6 @@ def after_request(response):
     response.headers["X-Request-Duration"] = str(duration)
     return response
 
-
 @app.route("/", methods=["GET"])
 def index():
     with open("./README.md", "r", encoding="utf8") as f:
@@ -251,42 +168,104 @@ def index():
     return render_template_string(markdown.markdown(content, extensions=["tables"]))
 
 
-@app.route("/api/extensions", methods=["GET"])
-def get_extensions():
-    extensions = dict(
-        {
-            "extensions": [
-                {
-                    "name": "not-supported",
-                    "metadata": {
-                        "display_name": """<span style="white-space:break-spaces;">Extensions serving using Extensions API is no longer supported. Please update the mod from: <a href="https://github.com/Cohee1207/SillyTavern">https://github.com/Cohee1207/SillyTavern</a></span>""",
-                        "requires": [],
-                        "assets": [],
-                    },
-                }
-            ]
-        }
+@app.route("/api/modules", methods=["GET"])
+def get_modules():
+    return jsonify({"modules": ['chromadb']})
+
+@app.route("/api/chromadb", methods=["POST"])
+def chromadb_add_messages():
+    data = request.get_json()
+    if "chat_id" not in data or not isinstance(data["chat_id"], str):
+        abort(400, '"chat_id" is required')
+    if "messages" not in data or not isinstance(data["messages"], list):
+        abort(400, '"messages" is required')
+
+    ip = get_real_ip()
+    chat_id_md5 = hashlib.md5(f'{ip}-{data["chat_id"]}'.encode()).hexdigest()
+    collection = chromadb_client.get_or_create_collection(
+        name=f"chat-{chat_id_md5}", embedding_function=chromadb_embed_fn
     )
-    return jsonify(extensions)
+
+    documents = [m["content"] for m in data["messages"]]
+    ids = [m["id"] for m in data["messages"]]
+    metadatas = [
+        {"role": m["role"], "date": m["date"], "meta": m.get("meta", "")}
+        for m in data["messages"]
+    ]
+
+    if len(ids) > 0:
+        collection.upsert(
+            ids=ids,
+            documents=documents,
+            metadatas=metadatas,
+        )
+
+    return jsonify({"count": len(ids)})
 
 
-@app.route("/api/caption", methods=["POST"])
-@require_module("caption")
-def api_caption():
+@app.route("/api/chromadb/query", methods=["POST"])
+def chromadb_query():
     data = request.get_json()
+    if "chat_id" not in data or not isinstance(data["chat_id"], str):
+        abort(400, '"chat_id" is required')
+    if "query" not in data or not isinstance(data["query"], str):
+        abort(400, '"query" is required')
+
+    if "n_results" not in data or not isinstance(data["n_results"], int):
+        n_results = 1
+    else:
+        n_results = data["n_results"]
 
-    if "image" not in data or not isinstance(data["image"], str):
-        abort(400, '"image" is required')
+    ip = get_real_ip()
+    chat_id_md5 = hashlib.md5(f'{ip}-{data["chat_id"]}'.encode()).hexdigest()
+    collection = chromadb_client.get_or_create_collection(
+        name=f"chat-{chat_id_md5}", embedding_function=chromadb_embed_fn
+    )
 
-    image = Image.open(BytesIO(base64.b64decode(data["image"])))
-    image = image.convert("RGB")
-    image.thumbnail((512, 512))
-    caption = caption_image(image)
-    thumbnail = image_to_base64(image)
-    print("Caption:", caption, sep="\n")
-    gc.collect()
-    return jsonify({"caption": caption, "thumbnail": thumbnail})
+    n_results = min(collection.count(), n_results)
 
+    messages = []
+    if n_results > 0:
+        query_result = collection.query(
+            query_texts=[data["query"]],
+            n_results=n_results,
+        )
+    
+        documents = query_result["documents"][0]
+        ids = query_result["ids"][0]
+        metadatas = query_result["metadatas"][0]
+        distances = query_result["distances"][0]
+    
+        messages = [
+            {
+                "id": ids[i],
+                "date": metadatas[i]["date"],
+                "role": metadatas[i]["role"],
+                "meta": metadatas[i]["meta"],
+                "content": documents[i],
+                "distance": distances[i],
+            }
+            for i in range(len(ids))
+        ]
+
+    return jsonify(messages)
+    
+@app.route("/api/chromadb/purge", methods=["POST"])
+def chromadb_purge():
+    data = request.get_json()
+    if "chat_id" not in data or not isinstance(data["chat_id"], str):
+        abort(400, '"chat_id" is required')
+
+    ip = get_real_ip()
+    chat_id_md5 = hashlib.md5(f'{ip}-{data["chat_id"]}'.encode()).hexdigest()
+    collection = chromadb_client.get_or_create_collection(
+        name=f"chat-{chat_id_md5}", embedding_function=chromadb_embed_fn
+    )
+
+    deleted = collection.delete()
+    print("ChromaDB embeddings deleted", len(deleted))
+
+    return 'Ok', 200
 
 @app.route("/api/summarize", methods=["POST"])
 @require_module("summarize")
@@ -308,8 +287,8 @@ def api_summarize():
     return jsonify({"summary": summary})
 
 
+
 @app.route("/api/classify", methods=["POST"])
-@require_module("classify")
 def api_classify():
     data = request.get_json()
 
@@ -324,138 +303,10 @@ def api_classify():
 
 
 @app.route("/api/classify/labels", methods=["GET"])
-@require_module("classify")
 def api_classify_labels():
     classification = classify_text("")
     labels = [x["label"] for x in classification]
     return jsonify({"labels": labels})
 
 
-@app.route("/api/keywords", methods=["POST"])
-@require_module("keywords")
-def api_keywords():
-    data = request.get_json()
-
-    if "text" not in data or not isinstance(data["text"], str):
-        abort(400, '"text" is required')
-
-    print("Keywords input:", data["text"], sep="\n")
-    keywords = extract_keywords(data["text"])
-    print("Keywords output:", keywords, sep="\n")
-    return jsonify({"keywords": keywords})
-
-
-@app.route("/api/prompt", methods=["POST"])
-@require_module("prompt")
-def api_prompt():
-    data = request.get_json()
-
-    if "text" not in data or not isinstance(data["text"], str):
-        abort(400, '"text" is required')
-
-    keywords = extract_keywords(data["text"])
-
-    if "name" in data and isinstance(data["name"], str):
-        keywords.insert(0, data["name"])
-
-    print("Prompt input:", data["text"], sep="\n")
-    prompts = generate_prompt(keywords)
-    print("Prompt output:", prompts, sep="\n")
-    return jsonify({"prompts": prompts})
-
-@app.route("/api/modules", methods=["GET"])
-def get_modules():
-    return jsonify({"modules": modules})
-
-@app.route("/api/chromadb", methods=["POST"])
-@require_module("chromadb")
-def chromadb_add_messages():
-    data = request.get_json()
-    if "chat_id" not in data or not isinstance(data["chat_id"], str):
-        abort(400, '"chat_id" is required')
-    if "messages" not in data or not isinstance(data["messages"], list):
-        abort(400, '"messages" is required')
-
-    chat_id_md5 = hashlib.md5(data["chat_id"].encode()).hexdigest()
-    collection = chromadb_client.get_or_create_collection(
-        name=f"chat-{chat_id_md5}", embedding_function=chromadb_embed_fn
-    )
-
-    documents = [m["content"] for m in data["messages"]]
-    ids = [m["id"] for m in data["messages"]]
-    metadatas = [
-        {"role": m["role"], "date": m["date"], "meta": m.get("meta", "")}
-        for m in data["messages"]
-    ]
-
-    collection.upsert(
-        ids=ids,
-        documents=documents,
-        metadatas=metadatas,
-    )
-
-    return jsonify({"count": len(ids)})
-
-
-@app.route("/api/chromadb/purge", methods=["POST"])
-@require_module("chromadb")
-def chromadb_purge():
-    data = request.get_json()
-    if "chat_id" not in data or not isinstance(data["chat_id"], str):
-        abort(400, '"chat_id" is required')
-
-    chat_id_md5 = hashlib.md5(data["chat_id"].encode()).hexdigest()
-    collection = chromadb_client.get_or_create_collection(
-        name=f"chat-{chat_id_md5}", embedding_function=chromadb_embed_fn
-    )
-
-    deleted = collection.delete()
-    print("ChromaDB embeddings deleted", len(deleted))
-    return 'Ok', 200
-
-
-@app.route("/api/chromadb/query", methods=["POST"])
-@require_module("chromadb")
-def chromadb_query():
-    data = request.get_json()
-    if "chat_id" not in data or not isinstance(data["chat_id"], str):
-        abort(400, '"chat_id" is required')
-    if "query" not in data or not isinstance(data["query"], str):
-        abort(400, '"query" is required')
-
-    if "n_results" not in data or not isinstance(data["n_results"], int):
-        n_results = 1
-    else:
-        n_results = data["n_results"]
-
-    chat_id_md5 = hashlib.md5(data["chat_id"].encode()).hexdigest()
-    collection = chromadb_client.get_or_create_collection(
-        name=f"chat-{chat_id_md5}", embedding_function=chromadb_embed_fn
-    )
-
-    n_results = min(collection.count(), n_results)
-    query_result = collection.query(
-        query_texts=[data["query"]],
-        n_results=n_results,
-    )
-
-    documents = query_result["documents"][0]
-    ids = query_result["ids"][0]
-    metadatas = query_result["metadatas"][0]
-    distances = query_result["distances"][0]
-
-    messages = [
-        {
-            "id": ids[i],
-            "date": metadatas[i]["date"],
-            "role": metadatas[i]["role"],
-            "meta": metadatas[i]["meta"],
-            "content": documents[i],
-            "distance": distances[i],
-        }
-        for i in range(len(ids))
-    ]
-
-    return jsonify(messages)
-
-app.run(host=host, port=port)
+app.run(host=host, port=port)