Certainly! I've reviewed your code and identified several issues that might prevent your functions from working correctly. Below are the problems along with suggestions and code corrections to fix them:
1. Incorrect Function Name:

    Issue: In your App class, the method analyze_sentiment_with_alberta should be analyze_sentiment_with_albert since you're using the AlbertModel.

    Fix: Rename the function to match the model used.

Corrected Function:

def analyze_sentiment_with_albert(self, text):
    inputs = albert_tokenizer(
        text,
        return_tensors="pt",
        max_length=512,
        truncation=True,
        padding=True
    ).to(device)
    outputs = albert_model(**inputs)
    logits = outputs.last_hidden_state.mean(dim=1)
    sentiment_score = torch.sigmoid(logits).item()
    if sentiment_score > 0.6:
        return "Positive"
    elif sentiment_score < 0.4:
        return "Negative"
    else:
        return "Neutral"

app.py
2. Model Loading Path Issue:

    Issue: In the load_model function, you're loading the model from a local path:

    GPT_J_MODEL = r'C:\solana_rust_bot\gpt-j-6B'

    This may cause issues if the model isn't located at that path.

    Fix: Load the model directly from Hugging Face's model hub. Also, consider using low_cpu_mem_usage=True to reduce memory usage.

Corrected load_model Function:

async def load_model():
    global model, tokenizer
    try:
        print("Testing model loading...")
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        tokenizer = AutoTokenizer.from_pretrained('EleutherAI/gpt-j-6B')
        print("Tokenizer loaded successfully.")

        model = GPTJForCausalLM.from_pretrained(
            'EleutherAI/gpt-j-6B',
            revision='float16',
            torch_dtype=torch.float16,
            low_cpu_mem_usage=True
        ).to(device)
        print("Model loaded successfully.")

        print("Model prepared successfully.")
        return True
    except Exception as e:
        print(f"Error loading model: {e}")
        return False

app.py
3. Groq Client Usage:

    Issue: The code imports from groq import Groq and uses groq_client.chat.completions.create(). However, there is no known groq client library that provides these methods.

    Fix: If you're intending to use OpenAI's GPT models, replace the groq client with the OpenAI client library.

Example Replacement Using OpenAI API (if applicable):

import openai

openai.api_key = os.getenv('OPENAI_API_KEY')

# Replace groq_client usage with openai.Completion or openai.ChatCompletion

app.py

    Note: Ensure you have the correct API keys and permissions to use the OpenAI API.

4. Asynchronous Function Execution in PyQt6:

    Issue: You are creating new event loops within the AsyncWorker, which can conflict with the main event loop used by PyQt6.

    Fix: Use asyncio.run or restructure your code to avoid creating new event loops.

Modified AsyncWorker Class:

class AsyncWorker(QRunnable):
    def __init__(self, coro):
        super().__init__()
        self.coro = coro
        self.signals = WorkerSignals()

    @pyqtSlot()
    def run(self):
        try:
            loop = asyncio.get_event_loop()
            if loop.is_running():
                result = asyncio.ensure_future(self.coro)
                result.add_done_callback(self.handle_result)
            else:
                result = loop.run_until_complete(self.coro)
                self.signals.result.emit(result)
        except Exception as e:
            self.signals.error.emit(str(e))
        finally:
            self.signals.finished.emit()

    def handle_result(self, future):
        try:
            result = future.result()
            self.signals.result.emit(result)
        except Exception as e:
            self.signals.error.emit(str(e))

app.py
5. Event Handling in eventFilter:

    Issue: Incorrect checking of modifiers and key events in eventFilter.

    Fix: Ensure that you're correctly checking for the Enter key without modifiers.

Corrected eventFilter Method:

def eventFilter(self, source, event):
    if (source is self.chat_input and
        event.type() == QEvent.Type.KeyPress and
        event.key() == Qt.Key_Return and
        not (event.modifiers() & Qt.KeyboardModifier.ShiftModifier)):
        self.send_chat()
        return True
    return super().eventFilter(source, event)

app.py
6. Async Function Calls During Shutdown:

    Issue: Asynchronous functions in graceful_shutdown may not execute properly if the event loop isn't running.

    Fix: Modify graceful_shutdown to await the coroutines properly.

Modified graceful_shutdown Method:

def graceful_shutdown(self):
    logger.info("Initiating graceful shutdown...")
    self.knowledge_update_timer.stop()
    self.threadpool.waitForDone()

    loop = asyncio.get_event_loop()
    if loop.is_running():
        asyncio.ensure_future(self.async_shutdown())
    else:
        loop.run_until_complete(self.async_shutdown())

def async_shutdown(self):
    await self.web_crawler.close_session()
    await self.final_database_operations()
    self.save_chat_memory()
    logger.info("Graceful shutdown completed.")

app.py
7. Package Installation:

    Issue: Missing packages or incorrect versions can cause import errors.

    Fix: Ensure all required packages are installed. Install them using the following command:

pip install bitsandbytes accelerate sentence-transformers transformers aiosqlite aiofiles aiolimiter beautifulsoup4 spacy rake-nltk nltk keybert ray psutil rich pygments python-dotenv PyQt6

    Additional Steps: For spacy and nltk, download the necessary models:

python -m spacy download en_core_web_sm
python -m nltk.downloader punkt stopwords

8. Potential Memory Issues with Large Models:

    Issue: Loading models like GPT-J-6B requires a significant amount of GPU memory.

    Fix: If you encounter memory errors, consider using model quantization techniques or offloading to CPU with slower performance.

9. Tensor Device and Shape Consistency:

    Issue: In search_knowledge, tensors might not have the correct shapes or devices.

    Fix: Ensure that tensors are correctly shaped and moved to the same device.

Corrected Code in search_knowledge:

similarity = F.cosine_similarity(
    torch.tensor(query_embedding).to(device),
    torch.tensor(content_embedding).to(device),
    dim=0
).item()

10. Regular Expression Patterns:

    Issue: Some regular expressions might not be correctly defined, potentially causing unexpected behavior.

    Fix: Verify and adjust your regular expressions as needed.

11. Ensure All Functions are Properly Defined and Called:

    Issue: Some asynchronous functions may not be awaited or used correctly.

    Fix: Go through your code to ensure that all async functions are awaited properly.

12. Handling NLTK Downloads:

    Issue: Downloading NLTK data at runtime can cause delays or failures if there's no internet connection.

    Fix: Consider downloading NLTK data separately or handle exceptions accordingly.

13. Logging Configuration:

    Issue: Logging might not capture the desired output or levels.

    Fix: Adjust your logging configuration to suit your needs.

14. Use of Undefined Variables:

    Issue: Variables like model and tokenizer are declared as global but might not be initialized before use.

    Fix: Ensure that these variables are properly initialized before they're accessed.

15. Exception Handling in Async Functions:

    Issue: Exceptions in asynchronous functions might not be caught properly.

    Fix: Use try-except blocks within your async functions to handle exceptions.

16. GUI Responsiveness:

    Issue: Long-running tasks might freeze the GUI.

    Fix: Ensure that all blocking operations are run asynchronously or in separate threads.

Final Recommendations:

    Testing: Write unit tests for individual functions to ensure they work as expected.

    Code Formatting: Use a code formatter like black or flake8 to maintain consistent code style.

    Code Comments: Add comments to explain complex logic, which aids in future maintenance.

    Performance Considerations: Be mindful of the performance implications of loading large models and performing heavy computations.

Let me know if you have any questions or need further assistance in implementing these changes!

import os
import sys
import asyncio
import aiohttp
import json
import logging
import hashlib
import re
import gc
from datetime import datetime
from typing import List, Dict, Any
from collections import deque

import torch
import torch.nn.functional as F
import bitsandbytes as bnb
from accelerate import Accelerator
from sentence_transformers import SentenceTransformer
from transformers import (
    GPTJForCausalLM,
    AutoTokenizer,
    AutoModel,
    BartForConditionalGeneration,
    BartTokenizer,
    AlbertModel,
    AlbertTokenizer,
)

from PyQt6.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, QPushButton,
    QTextEdit, QLabel, QProgressBar, QFileDialog, QMessageBox, QTabWidget, QSplitter
)
from PyQt6.QtCore import (
    Qt, QRunnable, QThreadPool, QObject, pyqtSignal, pyqtSlot, QTimer, QEvent, QRegularExpression
)
from PyQt6.QtGui import QTextCursor, QColor, QSyntaxHighlighter, QTextCharFormat

import aiosqlite
import aiofiles
from aiolimiter import AsyncLimiter
from bs4 import BeautifulSoup
import spacy
from rake_nltk import Rake
import nltk
from keybert import KeyBERT
from groq import Groq
import ray
import psutil
from rich.console import Console
from rich.logging import RichHandler
from pygments import highlight
from pygments.lexers import PythonLexer, RustLexer
from pygments.formatters import HtmlFormatter
from dotenv import load_dotenv

console = Console()
logging.basicConfig(
    level="INFO",
    format="%(message)s",
    datefmt="[%X]",
    handlers=[RichHandler(console=console)]
)
logger = logging.getLogger("rich")

load_dotenv()

GPT_J_MODEL = r'C:\solana_rust_bot\gpt-j-6B'
KNOWLEDGE_DB_PATH = 'solana_knowledge_base.db'
MEMORY_PATH = 'memory.json'
GROQ_API_KEY = os.getenv('GROQ_API_KEY')

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logger.info(f"Using device: {device}")

accelerator = Accelerator()

MAX_GROQ_CALLS_PER_MINUTE = 30
groq_limiter = AsyncLimiter(MAX_GROQ_CALLS_PER_MINUTE, 60)

nlp = spacy.load("en_core_web_sm")

rake = Rake()

nltk.download('punkt', quiet=True)
nltk.download('stopwords', quiet=True)

sentence_transformer = SentenceTransformer('all-MiniLM-L6-v2')

groq_client = Groq(api_key=GROQ_API_KEY)

ray.init(ignore_reinit_error=True)

key_bert = KeyBERT()

albert_tokenizer = AlbertTokenizer.from_pretrained('albert-base-v2')
albert_model = AlbertModel.from_pretrained('albert-base-v2').to(device)

bart_tokenizer = BartTokenizer.from_pretrained('facebook/bart-large')
bart_model = BartForConditionalGeneration.from_pretrained('facebook/bart-large').to(device)

async def load_model():
    global model, tokenizer
    try:
        print("Testing model loading...")
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        tokenizer = AutoTokenizer.from_pretrained(GPT_J_MODEL)
        print("Tokenizer loaded successfully.")

        model = GPTJForCausalLM.from_pretrained(GPT_J_MODEL).to(device)
        print("Model loaded successfully.")

        print("Model prepared successfully.")
        return True
    except Exception as e:
        print(f"Error loading model: {e}")
        return False

class SolanaKnowledgeDB:
    def __init__(self, db_path: str):
        self.db_path = db_path

    async def init_db(self):
        try:
            async with aiosqlite.connect(self.db_path) as db:
                await db.execute('''
                    CREATE TABLE IF NOT EXISTS knowledge_base
                        (id TEXT PRIMARY KEY, content TEXT, category TEXT, timestamp DATETIME DEFAULT CURRENT_TIMESTAMP)
                ''')
                await db.commit()
        except Exception as e:
            logger.error(f"Failed to initialize database: {e}")

    async def add_knowledge(self, content: str, category: str):
        content_hash = hashlib.sha256(content.encode()).hexdigest()
        try:
            async with aiosqlite.connect(self.db_path) as db:
                await db.execute(
                    'INSERT OR REPLACE INTO knowledge_base (id, content, category) VALUES (?, ?, ?)',
                    (content_hash, content, category)
                )
                await db.commit()
        except Exception as e:
            logger.error(f"Failed to add knowledge: {e}")

    async def search_knowledge(self, query: str) -> List[Dict[str, Any]]:
        query_embedding = sentence_transformer.encode(query)
        async with aiosqlite.connect(self.db_path) as db:
            cursor = await db.execute('SELECT * FROM knowledge_base')
            rows = await cursor.fetchall()
            results = []
            for row in rows:
                content_embedding = sentence_transformer.encode(row[1])
                similarity = F.cosine_similarity(
                    torch.tensor(query_embedding).unsqueeze(0).to(device),
                    torch.tensor(content_embedding).unsqueeze(0).to(device)
                ).item()
                results.append((similarity, {
                    'id': row[0],
                    'content': row[1],
                    'category': row[2],
                    'timestamp': row[3]
                }))
            results.sort(key=lambda x: x[0], reverse=True)
            return [item[1] for item in results[:5]]

    async def close_connection(self):
        try:
            async with aiosqlite.connect(self.db_path) as db:
                await db.execute('VACUUM')
                await db.commit()
            logger.info("Database connection closed and optimized.")
        except Exception as e:
            logger.error(f"Error closing database connection: {e}")

class WebCrawler:
    def __init__(self):
        self.session = None

    async def create_session(self):
        if self.session is None:
            self.session = aiohttp.ClientSession()

    async def close_session(self):
        if self.session:
            await self.session.close()
            self.session = None

    async def crawl_sourcegraph(self, query):
        await self.create_session()
        url = f"https://sourcegraph.com/search?q={query}&patternType=literal"
        try:
            async with self.session.get(url, timeout=30) as response:
                if response.status == 200:
                    html = await response.text()
                    soup = BeautifulSoup(html, 'html.parser')
                    results = soup.find_all('div', class_='result-container')
                    return [result.get_text() for result in results]
                else:
                    return []
        except asyncio.TimeoutError:
            return ["Error: Request timed out"]
        except Exception as e:
            return [f"Error: {str(e)}"]

    async def crawl_solana_docs(self, topic):
        await self.create_session()
        url = f"https://docs.solana.com/search?q={topic}"
        try:
            async with self.session.get(url, timeout=30) as response:
                if response.status == 200:
                    html = await response.text()
                    soup = BeautifulSoup(html, 'html.parser')
                    results = soup.find_all('div', class_='search-result-item')
                    return [result.get_text() for result in results]
                else:
                    return []
        except asyncio.TimeoutError:
            return ["Error: Request timed out"]
        except Exception as e:
            return [f"Error: {str(e)}"]

class KnowledgeUpdater:
    def __init__(self, knowledge_db: SolanaKnowledgeDB, web_crawler: WebCrawler):
        self.knowledge_db = knowledge_db
        self.web_crawler = web_crawler
        self.topics = [
            "Solana API", "Solana trading", "Solana RPC", "Solana websockets",
            "Helius API", "Jupiter V6", "Raydium DEX", "SPL tokens",
            "Solana Rust programming", "Solana transaction structure",
            "Solana program development", "Solana account model",
            "Solana staking", "Solana wallet integration", "Solana NFTs",
            "Solana DeFi protocols", "Solana security best practices",
            "Solana performance optimization", "Solana cluster management",
            "Solana ecosystem tools and frameworks"
        ]

    async def update_knowledge_base(self):
        for topic in self.topics:
            logger.info(f"Updating knowledge base for topic: {topic}")
            groq_info = await self.get_groq_information(topic)
            sourcegraph_info = await self.web_crawler.crawl_sourcegraph(topic)
            solana_docs_info = await self.web_crawler.crawl_solana_docs(topic)

            combined_info = (
                f"GROQ Info: {groq_info}\n\n"
                f"Sourcegraph Info: {' '.join(sourcegraph_info)}\n\n"
                f"Solana Docs Info: {' '.join(solana_docs_info)}"
            )
            await self.knowledge_db.add_knowledge(combined_info, topic)
            logger.info(f"Knowledge base updated for topic: {topic}")

    async def get_groq_information(self, topic):
        prompt = (
            f"Provide detailed, technical information about {topic} in the context of Solana blockchain development and trading. "
            f"Include code examples and best practices where applicable."
        )
        try:
            async with groq_limiter:
                response = await groq_client.chat.completions.create(
                    model="mixtral-8x7b-32768",
                    messages=[{"role": "user", "content": prompt}],
                    max_tokens=2000
                )
            return response.choices[0].message.content
        except Exception as e:
            logger.error(f"Error getting information from GROQ: {e}")
            return ""

class WorkerSignals(QObject):
    finished = pyqtSignal()
    result = pyqtSignal(object)
    error = pyqtSignal(str)

class AsyncWorker(QRunnable):
    def __init__(self, fn, *args, **kwargs):
        super().__init__()
        self.fn = fn
        self.args = args
        self.kwargs = kwargs
        self.signals = WorkerSignals()

    @pyqtSlot()
    def run(self):
        try:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            result = loop.run_until_complete(self.fn(*self.args, **self.kwargs))
            self.signals.result.emit(result)
        except Exception as e:
            self.signals.error.emit(str(e))
        finally:
            self.signals.finished.emit()

class CodeHighlighter(QSyntaxHighlighter):
    def __init__(self, parent, language):
        super().__init__(parent)
        self.language = language
        self.highlighting_rules = []

        keyword_format = QTextCharFormat()
        keyword_format.setForeground(Qt.GlobalColor.darkBlue)
        keyword_format.setFontWeight(700)

        keyword_patterns = [
            "\\bconst\\b", "\\blet\\b", "\\bfn\\b", "\\bpub\\b",
            "\\buse\\b", "\\bmod\\b", "\\bstruct\\b", "\\benum\\b",
            "\\bself\\b", "\\bmut\\b", "\\bimpl\\b", "\\btrait\\b"
        ]

        for pattern in keyword_patterns:
            self.highlighting_rules.append((QRegularExpression(pattern), keyword_format))

        class_format = QTextCharFormat()
        class_format.setFontWeight(700)
        class_format.setForeground(Qt.GlobalColor.darkMagenta)
        self.highlighting_rules.append((QRegularExpression("\\bQ[A-Za-z]+\\b"), class_format))
        self.highlighting_rules.append((QRegularExpression("\\b[A-Z][A-Za-z]*\\b"), class_format))

        quotation_format = QTextCharFormat()
        quotation_format.setForeground(Qt.GlobalColor.darkGreen)
        self.highlighting_rules.append((QRegularExpression("\".*\""), quotation_format))

        function_format = QTextCharFormat()
        function_format.setFontItalic(True)
        function_format.setForeground(Qt.GlobalColor.blue)
        self.highlighting_rules.append((QRegularExpression("\\b[A-Za-z0-9_]+(?=\\()"), function_format))

        comment_format = QTextCharFormat()
        comment_format.setForeground(Qt.GlobalColor.darkGreen)
        self.highlighting_rules.append((QRegularExpression("//[^\n]*"), comment_format))

    def highlightBlock(self, text):
        for pattern, format in self.highlighting_rules:
            expression = QRegularExpression(pattern)
            it = expression.globalMatch(text)
            while it.hasNext():
                match = it.next()
                self.setFormat(match.capturedStart(), match.capturedLength(), format)

class App(QMainWindow):
    MAX_CHAT_MEMORY = 100
    MAX_TOKENS_LIMIT = 4000

    def __init__(self):
        super().__init__()
        self.setWindowTitle('Ultimate Solana AI Assistant V3 - Enhanced Edition')
        self.setGeometry(100, 100, 1200, 800)

        self.threadpool = QThreadPool()
        self.knowledge_db = SolanaKnowledgeDB(KNOWLEDGE_DB_PATH)
        self.web_crawler = WebCrawler()
        self.knowledge_updater = KnowledgeUpdater(self.knowledge_db, self.web_crawler)
        self.chat_memory = self.load_chat_memory()
        self.current_token_count = 0

        self.model_loaded = False
        self.init_ui()
        self.init_knowledge_db()
        self.start_model_loading()
        self.start_knowledge_update_timer()

        app = QApplication.instance()
        app.aboutToQuit.connect(self.graceful_shutdown)

    def load_chat_memory(self):
        try:
            with open(MEMORY_PATH, 'r') as f:
                return deque(json.load(f), maxlen=self.MAX_CHAT_MEMORY)
        except FileNotFoundError:
            return deque(maxlen=self.MAX_CHAT_MEMORY)

    def save_chat_memory(self):
        with open(MEMORY_PATH, 'w') as f:
            json.dump(list(self.chat_memory), f)

    def init_knowledge_db(self):
        worker = AsyncWorker(self.knowledge_db.init_db)
        worker.signals.error.connect(self.handle_db_init_error)
        self.threadpool.start(worker)

    def handle_db_init_error(self, error_message):
        QMessageBox.critical(
            self,
            "Database Error",
            f"Failed to initialize the knowledge database: {error_message}"
        )
        logger.error(f"Database initialization error: {error_message}")

    def start_model_loading(self):
        async def load_model_wrapper():
            result = await load_model()
            if result:
                self.on_model_loaded()
            else:
                self.handle_model_load_error("Failed to load model")

        worker = AsyncWorker(load_model_wrapper)
        self.threadpool.start(worker)

    def on_model_loaded(self):
        global model, tokenizer
        if model is not None and tokenizer is not None:
            self.model_loaded = True
            logger.info("Model and tokenizer loaded successfully.")
            print(f"Model device: {next(model.parameters()).device}")
            print(f"Albert model device: {next(albert_model.parameters()).device}")
            print(f"Bart model device: {next(bart_model.parameters()).device}")
        else:
            QMessageBox.critical(
                self,
                "Model Load Error",
                "Failed to load the AI model and tokenizer."
            )

    def handle_model_load_error(self, error_message):
        QMessageBox.critical(
            self,
            "Model Load Error",
            f"An error occurred while loading the model: {error_message}"
        )
        logger.error(f"Model loading error: {error_message}")

    def start_knowledge_update_timer(self):
        self.knowledge_update_timer = QTimer(self)
        self.knowledge_update_timer.timeout.connect(self.update_knowledge_base)
        self.knowledge_update_timer.start(3600000)  # Run every hour

    def update_knowledge_base(self):
        worker = AsyncWorker(self.knowledge_updater.update_knowledge_base)
        self.threadpool.start(worker)

    def init_ui(self):
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QVBoxLayout(central_widget)

        self.tabs = QTabWidget()
        main_layout.addWidget(self.tabs)

        # AI Chat Tab
        chat_tab = QWidget()
        chat_layout = QVBoxLayout(chat_tab)
        self.tabs.addTab(chat_tab, "AI Chat")

        self.chat_display = QTextEdit()
        self.chat_display.setReadOnly(True)
        chat_layout.addWidget(self.chat_display)

        input_layout = QHBoxLayout()
        self.chat_input = QTextEdit()
        self.chat_input.setFixedHeight(50)
        self.chat_input.setPlaceholderText("Type your message here...")
        self.chat_input.installEventFilter(self)
        input_layout.addWidget(self.chat_input)

        send_btn = QPushButton('Send')
        send_btn.clicked.connect(self.send_chat)
        input_layout.addWidget(send_btn)

        chat_layout.addLayout(input_layout)

        self.clear_chat_btn = QPushButton('Clear Chat')
        self.clear_chat_btn.clicked.connect(self.clear_chat)
        chat_layout.addWidget(self.clear_chat_btn)

        # Code Editor Tab
        code_tab = QWidget()
        code_layout = QVBoxLayout(code_tab)
        self.tabs.addTab(code_tab, "Code Editor")

        self.code_editor = QTextEdit()
        self.code_editor.setPlaceholderText("Enter your Solana code here...")
        self.code_highlighter = CodeHighlighter(self.code_editor.document(), "rust")
        code_layout.addWidget(self.code_editor)

        analyze_btn = QPushButton('Analyze Code')
        analyze_btn.clicked.connect(self.analyze_code)
        code_layout.addWidget(analyze_btn)

        self.code_analysis_display = QTextEdit()
        self.code_analysis_display.setReadOnly(True)
        code_layout.addWidget(self.code_analysis_display)

        # Knowledge Base Tab
        kb_tab = QWidget()
        kb_layout = QVBoxLayout(kb_tab)
        self.tabs.addTab(kb_tab, "Knowledge Base")

        self.kb_search_input = QTextEdit()
        self.kb_search_input.setFixedHeight(50)
        self.kb_search_input.setPlaceholderText("Search the knowledge base...")
        kb_layout.addWidget(self.kb_search_input)

        kb_search_btn = QPushButton('Search')
        kb_search_btn.clicked.connect(self.search_knowledge_base)
        kb_layout.addWidget(kb_search_btn)

        self.kb_display = QTextEdit()
        self.kb_display.setReadOnly(True)
        kb_layout.addWidget(self.kb_display)

    def eventFilter(self, source, event):
        if (source is self.chat_input and
            event.type() == QEvent.Type.KeyPress and
            event.key() == Qt.Key.Key_Return and
            event.modifiers() == Qt.KeyboardModifier.NoModifier):
            self.send_chat()
            return True
        return super().eventFilter(source, event)

    def send_chat(self):
        if not self.model_loaded:
            QMessageBox.warning(self, "Warning", "AI model is still loading. Please wait.")
            return

        message = self.chat_input.toPlainText().strip()
        if message:
            if len(tokenizer.encode(message)) > 150:
                message = ' '.join(message.split()[:150])

            self.chat_display.append(f"User: {message}")
            self.chat_input.clear()
            self.chat_display.append("AI: Processing...")
            QApplication.processEvents()

            worker = AsyncWorker(self.handle_ai_chat, message)
            worker.signals.result.connect(self.update_chat_display)
            worker.signals.error.connect(self.handle_chat_error)

            self.threadpool.start(worker)

    async def handle_ai_chat(self, message):
        try:
            self.chat_memory.append({"role": "user", "content": message})
            self.save_chat_memory()

            knowledge_context = await self.get_relevant_knowledge(message)

            # Use KeyBERT to extract keywords
            keywords = key_bert.extract_keywords(message, keyphrase_ngram_range=(1, 2), stop_words='english', top_n=5)
            keyword_str = ", ".join([kw for kw, _ in keywords])

            context = (
                "You are an expert AI assistant specializing in Solana development and trading bots. "
                "Provide detailed, comprehensive responses with code examples where applicable. "
                f"Knowledge context: {knowledge_context}\n\n"
                f"Extracted keywords: {keyword_str}\n\n"
            )

            context += "\n".join([
                f"{msg['role'].capitalize()}: {msg['content']}" for msg in list(self.chat_memory)[-5:]
            ])

            context += f"\nUser: {message}\nAI:"

            if any(keyword in message.lower() for keyword in ["code", "script", "function", "class"]):
                sourcegraph_results = await self.web_crawler.crawl_sourcegraph(message)
                if sourcegraph_results:
                    context += f"\n\nRelevant code from Sourcegraph:\n" + "\n".join(sourcegraph_results[:3])

            response = await self.generate_response(context)

            # Use BART for summarization if the response is too long
            if len(response.split()) > 100:
                summarized_response = self.summarize_with_bart(response)
                response += f"\n\nSummary: {summarized_response}"

            self.chat_memory.append({"role": "assistant", "content": response})
            self.save_chat_memory()

            await self.knowledge_db.add_knowledge(response, "ai_generated")

            return response
        except Exception as e:
            logger.error(f"Error in handle_ai_chat: {str(e)}")
            return f"An error occurred: {str(e)}"

    async def generate_response(self, prompt):
        try:
            local_response = await self.generate_with_local_model(prompt)

            if len(local_response) < 100 or "I don't have enough information" in local_response:
                logger.info("Local model response insufficient. Using GROQ.")
                async with groq_limiter:
                    groq_response = await groq_client.chat.completions.create(
                        model="mixtral-8x7b-32768",
                        messages=[{"role": "user", "content": prompt}],
                        max_tokens=2000
                    )
                response = groq_response.choices[0].message.content
            else:
                response = local_response

            # Use AlBERTA for sentiment analysis
            sentiment = self.analyze_sentiment_with_alberta(response)
            response += f"\n\nSentiment: {sentiment}"

            return response
        except Exception as e:
            logger.error(f"Error generating response: {e}")
            return "I'm sorry, I encountered an error while processing your request."

    async def generate_with_local_model(self, prompt):
        inputs = tokenizer.encode(prompt, return_tensors="pt").to(device)
        attention_mask = (inputs != tokenizer.pad_token_id).long().to(device)

        with torch.no_grad():
            outputs = model.generate(
                inputs,
                attention_mask=attention_mask,
                max_length=min(len(inputs[0]) + 500, self.MAX_TOKENS_LIMIT),
                num_return_sequences=1,
                no_repeat_ngram_size=2,
                do_sample=True,  # Enabled sampling as per warning
                temperature=0.7,
                pad_token_id=tokenizer.eos_token_id  # Set pad_token_id
            )
        return tokenizer.decode(outputs[0], skip_special_tokens=True)

    def summarize_with_bart(self, text):
        inputs = bart_tokenizer(text, return_tensors="pt", max_length=1024, truncation=True).to(device)
        summary_ids = bart_model.generate(inputs["input_ids"], num_beams=4, min_length=30, max_length=100, early_stopping=True)
        return bart_tokenizer.decode(summary_ids[0], skip_special_tokens=True)

    def analyze_sentiment_with_alberta(self, text):
        inputs = albert_tokenizer(text, return_tensors="pt", max_length=512, truncation=True, padding=True).to(device)
        outputs = albert_model(**inputs)
        logits = outputs.last_hidden_state.mean(dim=1)
        sentiment_score = torch.sigmoid(logits).item()
        if sentiment_score > 0.6:
            return "Positive"
        elif sentiment_score < 0.4:
            return "Negative"
        else:
            return "Neutral"

    async def get_relevant_knowledge(self, query):
        results = await self.knowledge_db.search_knowledge(query)
        return "\n".join([result['content'] for result in results])

    def update_chat_display(self, message):
        self.chat_display.append(f"AI: {message}")
        self.chat_display.verticalScrollBar().setValue(
            self.chat_display.verticalScrollBar().maximum()
        )
        QApplication.processEvents()

    def handle_chat_error(self, error_message):
        self.chat_display.append(f"Error: {error_message}")
        QApplication.processEvents()

    def clear_chat(self):
        self.chat_display.clear()
        self.chat_memory.clear()
        self.current_token_count = 0
        self.chat_display.append("Chat cleared. Ready for new conversation.")
        self.save_chat_memory()
        QApplication.processEvents()

    def analyze_code(self):
        code = self.code_editor.toPlainText()
        if not code:
            QMessageBox.warning(self, "Warning", "Please enter some code to analyze.")
            return

        worker = AsyncWorker(self.perform_code_analysis, code)
        worker.signals.result.connect(self.update_code_analysis)
        worker.signals.error.connect(self.handle_analysis_error)

        self.threadpool.start(worker)

    async def perform_code_analysis(self, code):
        analysis = {
            "lines_of_code": len(code.split("\n")),
            "function_count": len(re.findall(r'\bfn\b', code)),
            "struct_count": len(re.findall(r'\bstruct\b', code)),
            "potential_issues": []
        }

        if "unwrap()" in code:
            analysis["potential_issues"].append("Use of unwrap() detected. Consider proper error handling.")

        if not re.search(r'#\[derive\(Debug\)\]', code):
            analysis["potential_issues"].append("Consider adding #[derive(Debug)] to structs for better debugging.")

        groq_analysis = await self.get_groq_code_analysis(code)
        analysis["groq_analysis"] = groq_analysis

        return analysis

    async def get_groq_code_analysis(self, code):
        prompt = (
            f"Analyze the following Solana Rust code and provide insights on best practices, potential improvements, "
            f"and any security considerations:\n\n{code}"
        )
        try:
            async with groq_limiter:
                response = await groq_client.chat.completions.create(
                    model="mixtral-8x7b-32768",
                    messages=[{"role": "user", "content": prompt}],
                    max_tokens=1000
                )
            return response.choices[0].message.content
        except Exception as e:
            logger.error(f"Error getting code analysis from GROQ: {e}")
            return "Unable to perform advanced analysis at this time."

    def update_code_analysis(self, analysis):
        self.code_analysis_display.clear()
        self.code_analysis_display.append(f"Lines of code: {analysis['lines_of_code']}")
        self.code_analysis_display.append(f"Function count: {analysis['function_count']}")
        self.code_analysis_display.append(f"Struct count: {analysis['struct_count']}")

        if analysis['potential_issues']:
            self.code_analysis_display.append("\nPotential issues:")
            for issue in analysis['potential_issues']:
                self.code_analysis_display.append(f"- {issue}")

        self.code_analysis_display.append(f"\nAdvanced Analysis:\n{analysis['groq_analysis']}")

    def handle_analysis_error(self, error_message):
        self.code_analysis_display.append(f"Error during analysis: {error_message}")

    def search_knowledge_base(self):
        query = self.kb_search_input.toPlainText().strip()
        if not query:
            QMessageBox.warning(self, "Warning", "Please enter a search query.")
            return

        worker = AsyncWorker(self.perform_kb_search, query)
        worker.signals.result.connect(self.update_kb_display)
        worker.signals.error.connect(self.handle_kb_search_error)

        self.threadpool.start(worker)

    async def perform_kb_search(self, query):
        results = await self.knowledge_db.search_knowledge(query)
        return results

    def update_kb_display(self, results):
        self.kb_display.clear()
        for result in results:
            self.kb_display.append(f"Category: {result['category']}")
            self.kb_display.append(f"Content: {result['content']}")
            self.kb_display.append(f"Timestamp: {result['timestamp']}")
            self.kb_display.append("\n" + "-"*50 + "\n")

    def handle_kb_search_error(self, error_message):
        self.kb_display.append(f"Error during search: {error_message}")

    def graceful_shutdown(self):
        logger.info("Initiating graceful shutdown...")
        self.knowledge_update_timer.stop()
        self.threadpool.waitForDone()

        asyncio.create_task(self.web_crawler.close_session())
        asyncio.create_task(self.final_database_operations())

        self.save_chat_memory()
        logger.info("Graceful shutdown completed.")

    async def final_database_operations(self):
        try:
            await self.knowledge_db.add_knowledge("Shutdown occurred", "system_log")
            await self.knowledge_db.close_connection()
        except Exception as e:
            logger.error(f"Error during final database operations: {e}")

def main():
    app = QApplication(sys.argv)
    ex = App()
    ex.show()
    sys.exit(app.exec())

if __name__ == '__main__':
    main()