250 lines
7.3 KiB
Python
250 lines
7.3 KiB
Python
import os
|
|
import json
|
|
import requests
|
|
from functools import wraps
|
|
|
|
from flask import Flask, request, jsonify, redirect, session, url_for
|
|
from authlib.integrations.flask_client import OAuth
|
|
|
|
from config import Config
|
|
|
|
# All comments are in English.
|
|
|
|
app = Flask(__name__)
|
|
app.config.from_object(Config)
|
|
|
|
# Configure session secret key.
|
|
app.secret_key = app.config["SECRET_KEY"]
|
|
|
|
# Configure OAuth with Google.
|
|
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",
|
|
access_token_params=None,
|
|
authorize_url="https://accounts.google.com/o/oauth2/v2/auth",
|
|
authorize_params={
|
|
"access_type": "offline",
|
|
"prompt": "consent",
|
|
},
|
|
api_base_url="https://www.googleapis.com/oauth2/v2/",
|
|
client_kwargs={"scope": "openid email profile"},
|
|
)
|
|
|
|
|
|
def login_required(f):
|
|
"""Decorator to ensure the user is authenticated via Google OAuth."""
|
|
|
|
@wraps(f)
|
|
def wrapper(*args, **kwargs):
|
|
# Check if user info is stored in session.
|
|
if "user" not in session:
|
|
# Redirect to login if not authenticated.
|
|
return redirect(url_for("login"))
|
|
return f(*args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
|
|
@app.route("/login")
|
|
def login():
|
|
"""Start Google OAuth login flow."""
|
|
redirect_uri = app.config["GOOGLE_REDIRECT_URI"]
|
|
# Use the configured redirect URI for OAuth callback.
|
|
return google.authorize_redirect(redirect_uri)
|
|
|
|
|
|
@app.route("/oauth/callback")
|
|
def oauth_callback():
|
|
"""Handle Google OAuth callback and store user session."""
|
|
# Exchange authorization code for tokens.
|
|
token = google.authorize_access_token()
|
|
# Fetch user info from Google.
|
|
resp = google.get("userinfo")
|
|
user_info = resp.json()
|
|
|
|
# Store minimal user info in session.
|
|
session["user"] = {
|
|
"email": user_info.get("email"),
|
|
"name": user_info.get("name"),
|
|
}
|
|
|
|
# Redirect to a simple home or API docs page.
|
|
return redirect(url_for("home"))
|
|
|
|
|
|
@app.route("/")
|
|
def home():
|
|
"""Simple home endpoint explaining the service."""
|
|
# This endpoint is intentionally simple and easy to understand.
|
|
return jsonify(
|
|
{
|
|
"message": "Ollama Gateway is running.",
|
|
"auth": "Google OAuth required for /api/* endpoints.",
|
|
"endpoints": {
|
|
"health": "/health",
|
|
"generate": "/api/generate",
|
|
},
|
|
}
|
|
)
|
|
|
|
|
|
@app.route("/health")
|
|
def health():
|
|
"""Health check endpoint."""
|
|
return jsonify({"status": "ok"})
|
|
|
|
|
|
@app.route("/api/generate", methods=["POST"])
|
|
@login_required
|
|
def generate():
|
|
"""Generate text using a selected Ollama model.
|
|
|
|
Expected JSON body:
|
|
{
|
|
"model": "llama3.2",
|
|
"prompt": "Your prompt here"
|
|
}
|
|
"""
|
|
data = request.get_json(silent=True) or {}
|
|
|
|
# Extract model and prompt from request.
|
|
model_name = data.get("model", "llama3.2")
|
|
prompt = data.get("prompt")
|
|
|
|
if not prompt:
|
|
return jsonify({"error": "Missing 'prompt' field."}), 400
|
|
|
|
# Map user-facing model name to Ollama model identifier.
|
|
model_id = app.config["MODEL_MAP"].get(model_name)
|
|
if not model_id:
|
|
return jsonify({"error": f"Unsupported model '{model_name}'."}), 400
|
|
|
|
# Prepare request to Ollama API.
|
|
ollama_url = f"{app.config['OLLAMA_BASE_URL']}/api/generate"
|
|
payload = {
|
|
"model": model_id,
|
|
"prompt": prompt,
|
|
# You can add more Ollama-specific parameters here if needed.
|
|
}
|
|
|
|
try:
|
|
# Call Ollama HTTP API.
|
|
response = requests.post(ollama_url, json=payload, timeout=60)
|
|
response.raise_for_status()
|
|
except requests.RequestException as e:
|
|
# Handle network or API errors gracefully.
|
|
return jsonify({"error": "Failed to call Ollama API.", "details": str(e)}), 502
|
|
|
|
# Parse Ollama response.
|
|
try:
|
|
ollama_data = response.json()
|
|
except json.JSONDecodeError:
|
|
return jsonify({"error": "Invalid JSON response from Ollama."}), 502
|
|
|
|
# Extract output text (this depends on Ollama's response format).
|
|
output_text = ollama_data.get("response") or ollama_data.get("output") or ""
|
|
|
|
return jsonify(
|
|
{
|
|
"model": model_name,
|
|
"output": output_text,
|
|
}
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# This block is mainly for local debugging.
|
|
# In production, Gunicorn will run the app.
|
|
app.run(host="0.0.0.0", port=8000)
|
|
|
|
|
|
# ***
|
|
|
|
|
|
from flask import Flask, request, jsonify
|
|
from flask_cors import CORS
|
|
from db import Base, engine
|
|
from models import User
|
|
from services.intent_service import IntentService
|
|
from services.mode_service import ModeService
|
|
from services.preference_service import PreferenceService
|
|
from services.suggestion_service import SuggestionService
|
|
from services.document_service import DocumentService
|
|
from services.search_service import SearchService
|
|
from services.monitoring_service import MonitoringService
|
|
from services.export_service import ExportService
|
|
from services.audit_service import AuditService
|
|
from services.security_service import SecurityService
|
|
|
|
from services.ollama_service import OllamaService
|
|
ollama_service = OllamaService()
|
|
|
|
from services.model_router import ModelRouter
|
|
router = ModelRouter()
|
|
generated = router.generate(intent, mode, prompt)
|
|
|
|
|
|
app = Flask(__name__)
|
|
CORS(app)
|
|
|
|
Base.metadata.create_all(bind=engine)
|
|
|
|
intent_service = IntentService()
|
|
mode_service = ModeService()
|
|
preference_service = PreferenceService()
|
|
suggestion_service = SuggestionService()
|
|
document_service = DocumentService()
|
|
search_service = SearchService()
|
|
monitoring_service = MonitoringService()
|
|
export_service = ExportService()
|
|
audit_service = AuditService()
|
|
security_service = SecurityService()
|
|
|
|
USER_ID = 1
|
|
|
|
@app.route("/api/conversation", methods=["POST"])
|
|
def conversation():
|
|
data = request.json
|
|
message = data.get("message", "")
|
|
|
|
mode = mode_service.get_current_mode(USER_ID)
|
|
if not mode:
|
|
reply = (
|
|
"{user.name}, como você quer trabalhar hoje?\n\n"
|
|
"- Modo Advogado\n- Modo Corporativo\n- Apenas Conversar\n\n"
|
|
"Me diga qual modo você prefere agora."
|
|
)
|
|
return jsonify({"reply": reply, "suggestions": [], "quick_actions": []})
|
|
|
|
intent = intent_service.detect_intent(message)
|
|
prefs = preference_service.get_preferences(USER_ID)
|
|
|
|
if intent == "set_mode":
|
|
mode_service.set_mode(USER_ID, intent_service.extract_mode(message))
|
|
reply = f"Perfeito, vou atuar no modo {mode_service.get_current_mode(USER_ID)}. O que vamos fazer agora?"
|
|
return jsonify({"reply": reply, "suggestions": [], "quick_actions": []})
|
|
|
|
if intent == "document":
|
|
response = document_service.handle(message, mode, prefs)
|
|
elif intent == "search":
|
|
response = search_service.handle(message, mode, prefs)
|
|
elif intent == "monitoring":
|
|
response = monitoring_service.handle(message, mode, prefs)
|
|
else:
|
|
response = document_service.chat_like(message, mode, prefs)
|
|
|
|
suggestions = suggestion_service.generate_suggestions(response, intent, mode, prefs)
|
|
quick_actions = suggestion_service.generate_quick_actions(response, intent, mode, prefs)
|
|
|
|
return jsonify({
|
|
"reply": response["text"],
|
|
"suggestions": suggestions,
|
|
"quick_actions": quick_actions
|
|
})
|
|
|
|
if __name__ == "__main__":
|
|
app.run(debug=True)
|