diff --git a/app/app.py b/app/app.py index 089c2fc..3c820a7 100644 --- a/app/app.py +++ b/app/app.py @@ -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/") +@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/", 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