mirror of
https://github.com/wassname/stampy-chat.git
synced 2026-09-23 13:50:27 +08:00
Fraser
844ce744e6
# Default to querying embeddings from live website if pinecone url not # present in .env # # This helps people getting started developing or messing around with the # site, since setting up a vector DB with the embeddings is by far the # hardest part for those not already on the team.
91 lines
2.7 KiB
Python
91 lines
2.7 KiB
Python
from flask import Flask, jsonify, request, Response
|
|
from flask_cors import CORS, cross_origin
|
|
from get_blocks import get_top_k_blocks
|
|
from chat import talk_to_robot
|
|
import dataclasses
|
|
import os
|
|
import openai
|
|
import pinecone
|
|
from discord_webhook import DiscordWebhook
|
|
|
|
|
|
# ---------------------------------- env setup ---------------------------------
|
|
|
|
if os.path.exists('.env'):
|
|
from dotenv import load_dotenv
|
|
load_dotenv()
|
|
else:
|
|
print("'api/.env' not found. Rename the 'api/.env.example' file and fill in values.")
|
|
|
|
|
|
OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
|
|
PINECONE_API_KEY = os.environ.get('PINECONE_API_KEY')
|
|
LOGGING_URL = os.environ.get('LOGGING_URL')
|
|
PINECONE_INDEX = None
|
|
|
|
openai.api_key = OPENAI_API_KEY # non-optional
|
|
|
|
# Only init pinecone if we have an env value for it.
|
|
if PINECONE_API_KEY is not None and PINECONE_API_KEY != "":
|
|
|
|
pinecone.init(
|
|
api_key = PINECONE_API_KEY,
|
|
environment = "us-east1-gcp",
|
|
)
|
|
|
|
PINECONE_INDEX = pinecone.Index(index_name="alignment-search")
|
|
|
|
# log something only if the logging url is set
|
|
def log(*args, end="\n"):
|
|
message = " ".join([str(arg) for arg in args]) + end
|
|
# print(message)
|
|
if LOGGING_URL is not None and LOGGING_URL != "":
|
|
while len(message) > 2000 - 8:
|
|
m_section, message = message[:2000 - 8], message[2000 - 8:]
|
|
m_section = "```\n" + m_section + "\n```"
|
|
DiscordWebhook(url=LOGGING_URL, content=m_section).execute()
|
|
DiscordWebhook(url=LOGGING_URL, content="```\n" + message + "\n```").execute()
|
|
|
|
# ---------------------------------- web setup ---------------------------------
|
|
|
|
app = Flask(__name__)
|
|
cors = CORS(app)
|
|
app.config['CORS_HEADERS'] = 'Content-Type'
|
|
|
|
# ---------------------------------- sse stuff ---------------------------------
|
|
|
|
def stream(src):
|
|
yield from ('data: ' + '\ndata: '.join(message.splitlines()) + '\n\n' for message in src)
|
|
yield 'event: close\n\n'
|
|
|
|
# ------------------------------- semantic search ------------------------------
|
|
|
|
|
|
@app.route('/semantic', methods=['POST'])
|
|
@cross_origin()
|
|
def semantic():
|
|
query = request.json['query']
|
|
k = request.json['k'] if 'k' in request.json else 20
|
|
return jsonify([dataclasses.asdict(block) for block in get_top_k_blocks(PINECONE_INDEX, query, k)])
|
|
|
|
|
|
|
|
# ------------------------------------ chat ------------------------------------
|
|
|
|
|
|
@app.route('/chat', methods=['POST'])
|
|
@cross_origin()
|
|
def chat():
|
|
|
|
query = request.json['query']
|
|
history = request.json['history']
|
|
|
|
return Response(stream(talk_to_robot(PINECONE_INDEX, query, history, log = log)), mimetype='text/event-stream')
|
|
|
|
|
|
# ------------------------------------------------------------------------------
|
|
|
|
|
|
if __name__ == '__main__':
|
|
app.run(debug=True, port=3000)
|