Atualizar app/app.py
This commit is contained in:
+137
@@ -1,3 +1,140 @@
|
||||
import os, json, uuid
|
||||
import requests
|
||||
from flask import Flask, Response, request, jsonify, session, redirect, url_for
|
||||
from authlib.integrations.flask_client import OAuth
|
||||
from functools import wraps
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
from markdown import markdown
|
||||
|
||||
from config import Config
|
||||
from models import Base, Conversation, Message, Attachment
|
||||
|
||||
app = Flask(__name__)
|
||||
app.config.from_object(Config)
|
||||
app.secret_key = app.config["SECRET_KEY"]
|
||||
|
||||
engine = create_engine(os.getenv("DATABASE_URL"))
|
||||
SessionLocal = sessionmaker(bind=engine)
|
||||
Base.metadata.create_all(engine)
|
||||
|
||||
oauth = OAuth(app)
|
||||
google = oauth.register(
|
||||
name="google",
|
||||
client_id=app.config["GOOGLE_CLIENT_ID"],
|
||||
client_secret=app.config["GOOGLE_CLIENT_SECRET"],
|
||||
access_token_url="https://oauth2.googleapis.com/token",
|
||||
authorize_url="https://accounts.google.com/o/oauth2/v2/auth",
|
||||
api_base_url="https://www.googleapis.com/oauth2/v2/",
|
||||
client_kwargs={"scope": "openid email profile"},
|
||||
)
|
||||
|
||||
def login_required(f):
|
||||
@wraps(f)
|
||||
def wrapper(*args, **kwargs):
|
||||
if "user" not in session:
|
||||
return redirect(url_for("login"))
|
||||
return f(*args, **kwargs)
|
||||
return wrapper
|
||||
|
||||
@app.route("/login")
|
||||
def login():
|
||||
return google.authorize_redirect(app.config["GOOGLE_REDIRECT_URI"])
|
||||
|
||||
@app.route("/oauth/callback")
|
||||
def oauth_callback():
|
||||
token = google.authorize_access_token()
|
||||
user_info = google.get("userinfo").json()
|
||||
session["user"] = user_info
|
||||
return redirect("/")
|
||||
|
||||
@app.route("/api/new_conversation", methods=["POST"])
|
||||
@login_required
|
||||
def new_conversation():
|
||||
db = SessionLocal()
|
||||
cid = str(uuid.uuid4())
|
||||
title = request.json.get("title", "New conversation")
|
||||
conv = Conversation(id=cid, title=title)
|
||||
db.add(conv)
|
||||
db.commit()
|
||||
db.close()
|
||||
return jsonify({"conversation_id": cid, "title": title})
|
||||
|
||||
@app.route("/api/conversations")
|
||||
@login_required
|
||||
def list_conversations():
|
||||
db = SessionLocal()
|
||||
convs = db.query(Conversation).all()
|
||||
data = [{"id": c.id, "title": c.title} for c in convs]
|
||||
db.close()
|
||||
return jsonify(data)
|
||||
|
||||
@app.route("/api/history/<cid>")
|
||||
@login_required
|
||||
def history(cid):
|
||||
db = SessionLocal()
|
||||
msgs = db.query(Message).filter_by(conversation_id=cid).order_by(Message.created_at).all()
|
||||
data = [{"role": m.role, "content": markdown(m.content)} for m in msgs]
|
||||
db.close()
|
||||
return jsonify(data)
|
||||
|
||||
@app.route("/api/attachment/<cid>", methods=["POST"])
|
||||
@login_required
|
||||
def upload_attachment(cid):
|
||||
db = SessionLocal()
|
||||
file = request.files.get("file")
|
||||
if not file:
|
||||
return jsonify({"error": "No file"}), 400
|
||||
content = file.read().decode("latin1")
|
||||
att = Attachment(
|
||||
id=str(uuid.uuid4()),
|
||||
conversation_id=cid,
|
||||
filename=file.filename,
|
||||
mime_type=file.mimetype,
|
||||
data=content,
|
||||
)
|
||||
db.add(att)
|
||||
db.commit()
|
||||
db.close()
|
||||
return jsonify({"status": "ok"})
|
||||
|
||||
@app.route("/api/stream")
|
||||
@login_required
|
||||
def stream():
|
||||
cid = request.args.get("conversation_id")
|
||||
model = request.args.get("model", "llama3.2")
|
||||
prompt = request.args.get("prompt")
|
||||
|
||||
db = SessionLocal()
|
||||
db.add(Message(id=str(uuid.uuid4()), conversation_id=cid, role="user", content=prompt))
|
||||
db.commit()
|
||||
|
||||
def generate():
|
||||
url = f"{app.config['OLLAMA_BASE_URL']}/api/generate"
|
||||
payload = {"model": model, "prompt": prompt, "stream": True}
|
||||
|
||||
with requests.post(url, json=payload, stream=True) as r:
|
||||
buffer = ""
|
||||
for line in r.iter_lines():
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
data = json.loads(line.decode())
|
||||
token = data.get("response", "")
|
||||
buffer += token
|
||||
# streaming animation: send partial buffer
|
||||
yield f"data: {json.dumps({'token': token})}\n\n"
|
||||
except:
|
||||
continue
|
||||
db.add(Message(id=str(uuid.uuid4()), conversation_id=cid, role="assistant", content=buffer))
|
||||
db.commit()
|
||||
db.close()
|
||||
|
||||
return Response(generate(), mimetype="text/event-stream")
|
||||
|
||||
|
||||
# ***
|
||||
|
||||
import json
|
||||
import requests
|
||||
from flask import Flask, Response, request, jsonify, session, redirect, url_for
|
||||
|
||||
Reference in New Issue
Block a user