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

    def __init__(self):
        super().__init__()
        self.setWindowTitle('Ultimate Solana AI Assistant V2 - Competition 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  # Flag to check if model is loaded
        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)

    # ... (other methods)

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

    def start_model_loading(self):
        worker = AsyncWorker(self.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.")
        else:
            QMessageBox.critical(
                self,
                "Model Load Error",
                "Failed to load the AI model and tokenizer."
            ) INTO THIS SCRIPT FOR ME AND REWRITE IT WITH THIS EDIT PLEASE KEEP ALL THE CODE THE EXACT SAME JUST THESE 2 EDITS AND PROVIDE ME THE ENTIRETY OF THE CODE 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
from transformers import GPTNeoForCausalLM, GPT2Tokenizer, AutoTokenizer, AutoModel
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
from PyQt6.QtGui import QTextCursor, QColor, QSyntaxHighlighter, QTextCharFormat
from PyQt6.QtCore import QRegularExpression

import aiosqlite
from bs4 import BeautifulSoup
import psutil
from sentence_transformers import SentenceTransformer
from accelerate import Accelerator
import aiofiles
from aiolimiter import AsyncLimiter
import spacy
from rich.console import Console
from rich.logging import RichHandler
import bitsandbytes as bnb
from rake_nltk import Rake
import nltk
from groq import Groq
import ray
from pygments import highlight
from pygments.lexers import PythonLexer, RustLexer
from pygments.formatters import HtmlFormatter

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

from dotenv import load_dotenv
load_dotenv()

GPT_NEO_MODEL = r'C:\solana_rust_bot\local_gpt_neo_model'  # Direct local path
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')  # Initialize sentence_transformer

groq_client = Groq(api_key=GROQ_API_KEY)

ray.init(ignore_reinit_error=True)

async def load_model():
    global model, tokenizer
    try:
        print("Testing model loading...")
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        
        # Load tokenizer
        tokenizer = GPT2Tokenizer.from_pretrained(GPT_NEO_MODEL)
        print("Tokenizer loaded successfully.")

        # Load model
        model = GPTNeoForCausalLM.from_pretrained(GPT_NEO_MODEL)
        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),
                    torch.tensor(content_embedding).unsqueeze(0)
                ).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 V2 - Competition 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  # Flag to check if model is loaded
        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.")
        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)

            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"
            )

            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)

            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

            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)
        with torch.no_grad():
            outputs = model.generate(
                inputs,
                max_length=min(len(inputs[0]) + 500, self.MAX_TOKENS_LIMIT),
                num_return_sequences=1,
                no_repeat_ngram_size=2,
                temperature=0.7,
            )
        return tokenizer.decode(outputs[0], skip_special_tokens=True)

    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()

        # Use asyncio.create_task instead of asyncio.run
        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()