Plank-AI/AIChat.py at main · garayot/Plank-AI · GitHub

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
from PySide6.QtWidgets import (
    QDialog, QVBoxLayout, QScrollArea, QLineEdit, QPushButton, QLabel,
    QSizePolicy, QHBoxLayout, QTextEdit, QWidget
)
from PySide6.QtCore import Qt, Signal, QTimer
from PySide6.QtGui import QFont
import google.generativeai as genai
from dotenv import load_dotenv
import os

class AIChatDialog(QDialog):
    response_received = Signal(str)

    def __init__(self, parent=None):
        super().__init__(parent)
        self.parent = parent
        self.setWindowTitle("Plank.AI Assistant")
        self.setFixedSize(350, 500)
        self.setStyleSheet("""
            QDialog {
                background-color: #1e1e1e;
                border-radius: 12px;
                padding: 10px;
            }
            QTextEdit, QLineEdit {
                background-color: #2e2e2e;
                color: white;
                border: 1px solid #444;
                border-radius: 6px;
                font-size: 13px;
                padding: 8px;
            }
            QPushButton {
                background-color: #3a82f7;
                color: white;
                border-radius: 6px;
                padding: 6px 12px;
            }
            QPushButton:hover {
                background-color: #5a9af9;
            }
        """)

        # AI Client setup
        load_dotenv()
        api_key = os.getenv("GEMINI_API_KEY")
        genai.configure(api_key=api_key)
        self.model = genai.GenerativeModel("gemini-2.0-flash")
        self.chat = self.model.start_chat(history=[])

        # Main layout
        self.layout = QVBoxLayout(self)
        self.layout.setContentsMargins(10, 10, 10, 10)

        # Close button
        close_btn = QPushButton("×")
        close_btn.setFixedSize(24, 24)
        close_btn.clicked.connect(self.hide)
        close_btn.setStyleSheet("background-color: transparent; font-size: 18px; color: white; border: none;")
        close_btn.setCursor(Qt.PointingHandCursor)

        top_bar = QHBoxLayout()
        top_bar.addStretch()
        top_bar.addWidget(close_btn)
        self.layout.addLayout(top_bar)

        # Chat history
        self.chat_history = QScrollArea()
        self.chat_history.setWidgetResizable(True)
        self.history_content = QWidget()
        self.history_layout = QVBoxLayout(self.history_content)
        self.chat_history.setWidget(self.history_content)
        self.layout.addWidget(self.chat_history)

        # Input area
        self.input_field = QLineEdit()
        self.input_field.setPlaceholderText("Ask about the data...")
        self.send_btn = QPushButton("Send")
        self.send_btn.clicked.connect(self.process_query)

        input_layout = QHBoxLayout()
        input_layout.addWidget(self.input_field)
        input_layout.addWidget(self.send_btn)
        self.layout.addLayout(input_layout)

        self.response_received.connect(self.display_response)

    def process_query(self):
        query = self.input_field.text()
        if not query:
            return

        self.add_message(query, is_user=True)
        self.input_field.clear()

        # Add context if available
        tally_data = self.parent.species_tally if self.parent else {}
        context = f"Current plankton tally: {tally_data}. "
        QTimer.singleShot(100, lambda: self.get_ai_response(context + query))

    def get_ai_response(self, query):
        try:
            response = self.chat.send_message(query)
            self.response_received.emit(response.text)
        except Exception as e:
            self.response_received.emit(f"Error: {str(e)}")

    def display_response(self, response):
        self.add_message(response, is_user=False)

    def add_message(self, text, is_user=True):
        message = QLabel(text)
        message.setWordWrap(True)
        message.setStyleSheet(f"""
            background-color: {'#3498db' if is_user else '#e0e0e0'};
            color: {'white' if is_user else 'black'};
            border-radius: 8px;
            padding: 8px;

        """)
        message.setSizePolicy(QSizePolicy.Maximum, QSizePolicy.Preferred)
        self.history_layout.addWidget(message, 0, Qt.AlignRight if is_user else Qt.AlignLeft)

        QTimer.singleShot(50, lambda: self.chat_history.verticalScrollBar().setValue(
            self.chat_history.verticalScrollBar().maximum()
        ))