File size: 4,594 Bytes
8174ba7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 |
from flask import (
Flask,
jsonify,
request,
render_template_string,
abort,
)
from flask_cors import CORS
import unicodedata
import markdown
import time
import os
import gc
import base64
from io import BytesIO
from random import randint
import hashlib
from colorama import Fore, Style, init as colorama_init
import chromadb
import posthog
from chromadb.config import Settings
from sentence_transformers import SentenceTransformer
from werkzeug.middleware.proxy_fix import ProxyFix
colorama_init()
port = 7860
host = "0.0.0.0"
embedding_model = 'sentence-transformers/all-mpnet-base-v2'
print("Initializing ChromaDB")
# 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 get_real_ip():
return request.remote_addr
@app.route("/", methods=["GET"])
def index():
with open("./README.md", "r", encoding="utf8") as f:
content = f.read()
return render_template_string(markdown.markdown(content, extensions=["tables"]))
@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
)
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/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"]
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
)
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.run(host=host, port=port) |