252 lines
9.3 KiB
Python
252 lines
9.3 KiB
Python
|
|
from __future__ import annotations
|
||
|
|
import json
|
||
|
|
import sys
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
sys.path.insert(0, str(Path(__file__).parent.parent))
|
||
|
|
|
||
|
|
import chromadb
|
||
|
|
from PySide6.QtCore import Qt, QEvent, QUrl
|
||
|
|
from PySide6.QtWidgets import (
|
||
|
|
QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
|
||
|
|
QLabel, QComboBox, QTextEdit, QPushButton,
|
||
|
|
)
|
||
|
|
from PySide6.QtWebEngineWidgets import QWebEngineView
|
||
|
|
|
||
|
|
from gui.worker import RagWorker
|
||
|
|
|
||
|
|
_HERE = Path(__file__).parent
|
||
|
|
_PROJECT = _HERE.parent
|
||
|
|
_CHROMA_DIR = _PROJECT / "chroma_db"
|
||
|
|
_CHAT_HTML = _HERE / "chat.html"
|
||
|
|
|
||
|
|
# ── Temi ──────────────────────────────────────────────────────────────────────
|
||
|
|
|
||
|
|
_THEMES = {
|
||
|
|
"dark": {
|
||
|
|
"topbar": "background:#171717; border-bottom:1px solid #2e2e2e;",
|
||
|
|
"lbl": "color:#aaa; font-size:13px;",
|
||
|
|
"combo": """
|
||
|
|
QComboBox {
|
||
|
|
background:#2a2a2a; color:#ececec;
|
||
|
|
border:1px solid #3a3a3a; border-radius:6px;
|
||
|
|
padding:4px 10px; font-size:13px;
|
||
|
|
}
|
||
|
|
QComboBox::drop-down { border:none; }
|
||
|
|
QComboBox QAbstractItemView {
|
||
|
|
background:#2a2a2a; color:#ececec;
|
||
|
|
selection-background-color:#3a3a3a;
|
||
|
|
}
|
||
|
|
""",
|
||
|
|
"inputbar": "background:#171717; border-top:1px solid #2e2e2e;",
|
||
|
|
"input": """
|
||
|
|
QTextEdit {
|
||
|
|
background:#2a2a2a; color:#ececec;
|
||
|
|
border:1px solid #3a3a3a; border-radius:10px;
|
||
|
|
padding:8px 12px; font-size:14px;
|
||
|
|
}
|
||
|
|
""",
|
||
|
|
"theme_btn": "background:#2a2a2a; color:#ececec; border:1px solid #3a3a3a; border-radius:6px; font-size:16px; padding:4px 8px;",
|
||
|
|
"icon": "☀️",
|
||
|
|
},
|
||
|
|
"light": {
|
||
|
|
"topbar": "background:#f5f5f5; border-bottom:1px solid #e0e0e0;",
|
||
|
|
"lbl": "color:#555; font-size:13px;",
|
||
|
|
"combo": """
|
||
|
|
QComboBox {
|
||
|
|
background:#ffffff; color:#1a1a1a;
|
||
|
|
border:1px solid #ddd; border-radius:6px;
|
||
|
|
padding:4px 10px; font-size:13px;
|
||
|
|
}
|
||
|
|
QComboBox::drop-down { border:none; }
|
||
|
|
QComboBox QAbstractItemView {
|
||
|
|
background:#ffffff; color:#1a1a1a;
|
||
|
|
selection-background-color:#e8e8e8;
|
||
|
|
}
|
||
|
|
""",
|
||
|
|
"inputbar": "background:#f5f5f5; border-top:1px solid #e0e0e0;",
|
||
|
|
"input": """
|
||
|
|
QTextEdit {
|
||
|
|
background:#ffffff; color:#1a1a1a;
|
||
|
|
border:1px solid #ddd; border-radius:10px;
|
||
|
|
padding:8px 12px; font-size:14px;
|
||
|
|
}
|
||
|
|
""",
|
||
|
|
"theme_btn": "background:#fff; color:#1a1a1a; border:1px solid #ddd; border-radius:6px; font-size:16px; padding:4px 8px;",
|
||
|
|
"icon": "🌙",
|
||
|
|
},
|
||
|
|
}
|
||
|
|
|
||
|
|
_BTN_SEND = """
|
||
|
|
QPushButton {
|
||
|
|
background:#10a37f; color:#fff;
|
||
|
|
border:none; border-radius:8px; font-size:14px; font-weight:600;
|
||
|
|
}
|
||
|
|
QPushButton:hover { background:#1ab58e; }
|
||
|
|
QPushButton:pressed { background:#0d8a6a; }
|
||
|
|
QPushButton:disabled { background:#2a2a2a; color:#555; }
|
||
|
|
"""
|
||
|
|
|
||
|
|
|
||
|
|
class ChatWindow(QMainWindow):
|
||
|
|
def __init__(self, default_collection: str = "") -> None:
|
||
|
|
super().__init__()
|
||
|
|
self.setWindowTitle("RAG Assistant")
|
||
|
|
self.resize(1000, 720)
|
||
|
|
|
||
|
|
self._dark = True
|
||
|
|
self._worker: RagWorker | None = None
|
||
|
|
self._collection: chromadb.Collection | None = None
|
||
|
|
self._chroma = chromadb.PersistentClient(path=str(_CHROMA_DIR))
|
||
|
|
|
||
|
|
self._build_ui()
|
||
|
|
self._load_collections(default_collection)
|
||
|
|
self._load_page()
|
||
|
|
|
||
|
|
# ── UI ────────────────────────────────────────────────────────────────────
|
||
|
|
|
||
|
|
def _build_ui(self) -> None:
|
||
|
|
root = QWidget()
|
||
|
|
self.setCentralWidget(root)
|
||
|
|
layout = QVBoxLayout(root)
|
||
|
|
layout.setSpacing(0)
|
||
|
|
layout.setContentsMargins(0, 0, 0, 0)
|
||
|
|
|
||
|
|
# Top bar
|
||
|
|
self._top = QWidget()
|
||
|
|
self._top.setFixedHeight(46)
|
||
|
|
tl = QHBoxLayout(self._top)
|
||
|
|
tl.setContentsMargins(14, 0, 14, 0)
|
||
|
|
tl.setSpacing(8)
|
||
|
|
|
||
|
|
self._lbl = QLabel("Collezione:")
|
||
|
|
tl.addWidget(self._lbl)
|
||
|
|
|
||
|
|
self._combo = QComboBox()
|
||
|
|
self._combo.setFixedWidth(260)
|
||
|
|
self._combo.currentTextChanged.connect(self._on_collection_changed)
|
||
|
|
tl.addWidget(self._combo)
|
||
|
|
|
||
|
|
tl.addStretch()
|
||
|
|
|
||
|
|
self._theme_btn = QPushButton()
|
||
|
|
self._theme_btn.setFixedSize(36, 30)
|
||
|
|
self._theme_btn.setToolTip("Cambia tema")
|
||
|
|
self._theme_btn.clicked.connect(self._toggle_theme)
|
||
|
|
tl.addWidget(self._theme_btn)
|
||
|
|
|
||
|
|
layout.addWidget(self._top)
|
||
|
|
|
||
|
|
# Web view
|
||
|
|
self._web = QWebEngineView()
|
||
|
|
s = self._web.settings()
|
||
|
|
s.setAttribute(s.WebAttribute.JavascriptEnabled, True)
|
||
|
|
s.setAttribute(s.WebAttribute.LocalContentCanAccessRemoteUrls, True)
|
||
|
|
layout.addWidget(self._web, stretch=1)
|
||
|
|
|
||
|
|
# Input bar
|
||
|
|
self._bar = QWidget()
|
||
|
|
bl = QHBoxLayout(self._bar)
|
||
|
|
bl.setContentsMargins(14, 10, 14, 12)
|
||
|
|
bl.setSpacing(10)
|
||
|
|
|
||
|
|
self._input = QTextEdit()
|
||
|
|
self._input.setPlaceholderText("Fai una domanda… (Invio per inviare, Shift+Invio per andare a capo)")
|
||
|
|
self._input.setFixedHeight(68)
|
||
|
|
self._input.installEventFilter(self)
|
||
|
|
bl.addWidget(self._input, stretch=1)
|
||
|
|
|
||
|
|
self._btn = QPushButton("Invia")
|
||
|
|
self._btn.setFixedSize(76, 42)
|
||
|
|
self._btn.setStyleSheet(_BTN_SEND)
|
||
|
|
self._btn.clicked.connect(self._on_send)
|
||
|
|
bl.addWidget(self._btn)
|
||
|
|
|
||
|
|
layout.addWidget(self._bar)
|
||
|
|
|
||
|
|
self._apply_theme()
|
||
|
|
|
||
|
|
def _apply_theme(self) -> None:
|
||
|
|
t = _THEMES["dark" if self._dark else "light"]
|
||
|
|
self._top.setStyleSheet(t["topbar"])
|
||
|
|
self._lbl.setStyleSheet(t["lbl"])
|
||
|
|
self._combo.setStyleSheet(t["combo"])
|
||
|
|
self._bar.setStyleSheet(t["inputbar"])
|
||
|
|
self._input.setStyleSheet(t["input"])
|
||
|
|
self._theme_btn.setStyleSheet(t["theme_btn"])
|
||
|
|
self._theme_btn.setText(t["icon"])
|
||
|
|
|
||
|
|
# ── Init ──────────────────────────────────────────────────────────────────
|
||
|
|
|
||
|
|
def _load_collections(self, default: str) -> None:
|
||
|
|
names = sorted(c.name for c in self._chroma.list_collections())
|
||
|
|
self._combo.addItems(names)
|
||
|
|
if default and default in names:
|
||
|
|
self._combo.setCurrentText(default)
|
||
|
|
|
||
|
|
def _load_page(self) -> None:
|
||
|
|
self._web.setUrl(QUrl.fromLocalFile(str(_CHAT_HTML)))
|
||
|
|
|
||
|
|
# ── Events ────────────────────────────────────────────────────────────────
|
||
|
|
|
||
|
|
def eventFilter(self, obj, event) -> bool:
|
||
|
|
if obj is self._input and event.type() == QEvent.Type.KeyPress:
|
||
|
|
if event.key() in (Qt.Key.Key_Return, Qt.Key.Key_Enter):
|
||
|
|
if not (event.modifiers() & Qt.KeyboardModifier.ShiftModifier):
|
||
|
|
self._on_send()
|
||
|
|
return True
|
||
|
|
return super().eventFilter(obj, event)
|
||
|
|
|
||
|
|
def _on_collection_changed(self, name: str) -> None:
|
||
|
|
if name:
|
||
|
|
self._collection = self._chroma.get_collection(name)
|
||
|
|
|
||
|
|
def _toggle_theme(self) -> None:
|
||
|
|
self._dark = not self._dark
|
||
|
|
self._apply_theme()
|
||
|
|
theme = "dark" if self._dark else "light"
|
||
|
|
self._js(f"setTheme({json.dumps(theme)})")
|
||
|
|
|
||
|
|
def _on_send(self) -> None:
|
||
|
|
question = self._input.toPlainText().strip()
|
||
|
|
if not question:
|
||
|
|
return
|
||
|
|
if not self._collection:
|
||
|
|
self._js("addMessage('assistant', '⚠️ Seleziona una collezione prima di inviare.', [])")
|
||
|
|
return
|
||
|
|
|
||
|
|
self._input.clear()
|
||
|
|
self._set_enabled(False)
|
||
|
|
self._js(f"addMessage('user', {json.dumps(question)}, [])")
|
||
|
|
self._js("addThinking()")
|
||
|
|
|
||
|
|
self._worker = RagWorker(question, self._collection, parent=self)
|
||
|
|
self._worker.finished.connect(self._on_response)
|
||
|
|
self._worker.error.connect(self._on_error)
|
||
|
|
self._worker.start()
|
||
|
|
|
||
|
|
def _on_response(self, response: str, chunks: list) -> None:
|
||
|
|
self._js("removeThinking()")
|
||
|
|
sources = [
|
||
|
|
{"source": c.get("source", ""), "header_path": c.get("header_path", "")}
|
||
|
|
for c in chunks
|
||
|
|
]
|
||
|
|
self._js(f"addMessage('assistant', {json.dumps(response)}, {json.dumps(sources)})")
|
||
|
|
self._set_enabled(True)
|
||
|
|
|
||
|
|
def _on_error(self, msg: str) -> None:
|
||
|
|
self._js("removeThinking()")
|
||
|
|
self._js(f"addMessage('assistant', {json.dumps('❌ Errore: ' + msg)}, [])")
|
||
|
|
self._set_enabled(True)
|
||
|
|
|
||
|
|
# ── Helpers ───────────────────────────────────────────────────────────────
|
||
|
|
|
||
|
|
def _js(self, script: str) -> None:
|
||
|
|
self._web.page().runJavaScript(script)
|
||
|
|
|
||
|
|
def _set_enabled(self, enabled: bool) -> None:
|
||
|
|
self._input.setEnabled(enabled)
|
||
|
|
self._btn.setEnabled(enabled)
|
||
|
|
if enabled:
|
||
|
|
self._input.setFocus()
|