File size: 6,676 Bytes
8a2dcce
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
"""
Gradio callback functions.

These call chat_service / auth / history directly (in-process) rather than
going over HTTP, since the Gradio Blocks app is mounted into the same Flask
process (see app.py). Because Gradio callbacks run outside of a normal
Flask request context, every function that touches the database wraps its
body in `with _flask_app.app_context():`.

Auth state for the UI is NOT the Flask session cookie (that's only for the
REST API in auth/routes.py) — it's tracked per-browser-tab in a gr.State
dict: {"user_id": int, "email": str} or None when logged out.
"""

import time

from auth.db import db
from auth.models import User
from chat_service import handle_chat_message
from history import service as history_service
from logs.logger import get_logger

logger = get_logger(__name__)

_flask_app = None


def init_app(flask_app):
    """Called once from app.py so callbacks can open an app context."""
    global _flask_app
    _flask_app = flask_app


# ---------------------------------------------------------------------------
# Auth
# ---------------------------------------------------------------------------

def do_signup(email: str, password: str):
    """Returns (auth_state, status_message, is_error)."""
    email = (email or "").strip().lower()
    password = password or ""

    if not email or "@" not in email:
        return None, "Please enter a valid email address.", True
    if len(password) < 8:
        return None, "Password must be at least 8 characters.", True

    with _flask_app.app_context():
        if User.query.filter_by(email=email).first() is not None:
            return None, "An account with this email already exists.", True

        user = User(email=email)
        user.set_password(password)
        db.session.add(user)
        db.session.commit()
        auth_state = {"user_id": user.id, "email": user.email}

    return auth_state, f"Account created. Welcome, {email}!", False


def do_login(email: str, password: str):
    """Returns (auth_state, status_message, is_error)."""
    email = (email or "").strip().lower()
    password = password or ""

    with _flask_app.app_context():
        user = User.query.filter_by(email=email).first()
        if user is None or not user.check_password(password):
            return None, "Invalid email or password.", True
        auth_state = {"user_id": user.id, "email": user.email}

    return auth_state, f"Logged in as {email}.", False


def do_logout():
    return None, "Logged out.", False


# ---------------------------------------------------------------------------
# Chat
# ---------------------------------------------------------------------------

def send_message(user_message: str, chat_display: list, auth_state, session_state, anon_history: list):
    """
    Args:
        user_message: raw text from the input box.
        chat_display: current [{"role": ..., "content": ...}, ...] shown in
                       gr.Chatbot(type="messages").
        auth_state: {"user_id", "email"} or None.
        session_state: current chat_session id (logged-in only) or None.
        anon_history: client-side rolling history used only when logged out.

    Returns:
        (chat_display, session_state, anon_history, cleared_textbox)
    """
    user_message = (user_message or "").strip()
    if not user_message:
        return chat_display, session_state, anon_history, ""

    chat_display = chat_display + [{"role": "user", "content": user_message}]

    is_logged_in = bool(auth_state)

    with _flask_app.app_context():
        if is_logged_in:
            user = User.query.get(auth_state["user_id"])
            result = handle_chat_message(
                query=user_message,
                user=user,
                session_id=session_state,
            )
            session_state = result["session_id"]
        else:
            result = handle_chat_message(
                query=user_message,
                user=None,
                anonymous_history=anon_history,
            )
            anon_history = (anon_history or []) + [
                {"role": "user", "content": user_message},
                {"role": "assistant", "content": result["response"]},
            ]

    chat_display = chat_display + [{"role": "assistant", "content": result["response"]}]
    return chat_display, session_state, anon_history, ""


def stream_last_response(chat_display: list):
    """
    Fakes token-by-token streaming for the assistant's last message so the
    UI feels responsive even though llm.generate() currently returns the
    full string at once. Yields progressively longer chat_display copies.
    """
    if not chat_display or chat_display[-1]["role"] != "assistant":
        yield chat_display
        return

    full_text = chat_display[-1]["content"]
    prefix_display = chat_display[:-1]

    step = max(1, len(full_text) // 60)
    for i in range(0, len(full_text) + step, step):
        partial = full_text[:i]
        yield prefix_display + [{"role": "assistant", "content": partial}]
        time.sleep(0.012)

    yield chat_display


def new_chat():
    """Resets the chat window + clears the active session id."""
    return [], None, []


# ---------------------------------------------------------------------------
# History panel
# ---------------------------------------------------------------------------

def refresh_sessions(auth_state):
    """Returns gr.update(...) choices for the session list, newest first."""
    import gradio as gr

    if not auth_state:
        return gr.update(choices=[], value=None)

    with _flask_app.app_context():
        sessions = history_service.list_sessions(auth_state["user_id"])

    choices = [(s["title"] or f"Session {s['id']}", s["id"]) for s in sessions]
    return gr.update(choices=choices, value=None)


def open_session(session_id, auth_state):
    """Loads a past session's full transcript into the chat window."""
    if not auth_state or session_id is None:
        return [], None

    with _flask_app.app_context():
        full_session = history_service.get_full_session(session_id, auth_state["user_id"])

    if full_session is None:
        return [], None

    chat_display = [
        {"role": m["role"], "content": m["content"]} for m in full_session["messages"]
    ]
    return chat_display, session_id


def delete_session(session_id, auth_state):
    """Deletes a session and returns a cleared chat window if it was active."""
    if not auth_state or session_id is None:
        return [], None

    with _flask_app.app_context():
        history_service.delete_session(session_id, auth_state["user_id"])

    return [], None