LLM Security Chatbot

app.py

import streamlit as st
from llama_cpp import Llama
import json

# Initialize the Llama model
llm = Llama(
    model_path="../llama.cpp/models/mistral-7B-v0.1/ggml-model-Q4_K_M.gguf",
    n_ctx=4096,
    n_gpu_layers=-1,
    chat_format="chatml"
)

st.title('Cybersecurity LLM ChatBot')

# Initialize or reset conversation history and debug info if not present
if 'conversation_history' not in st.session_state:
    st.session_state.conversation_history = []
if 'debug_info' not in st.session_state:
    st.session_state.debug_info = {}

user_query = st.text_area("Enter your query related to cybersecurity research", '')

prompt_template = """
Q: {user_query}
Context: The user is seeking in-depth analysis and actionable insights on cybersecurity vulnerabilities. 
Please provide detailed information, potential mitigation strategies, and reference relevant tools or resources.
A: """

def display_debug_info():
    if st.session_state.debug_info:
        st.json(st.session_state.debug_info)

if st.button('Submit'):
    if user_query:
        st.session_state.conversation_history.append(f"User: {user_query}")
        with st.spinner('Analyzing your query...'):
            prompt = prompt_template.format(user_query=user_query)
            output = llm(prompt, max_tokens=2048, stop=["Q:"], echo=True)
            
            if 'choices' in output and len(output['choices']) > 0:
                raw_response = output['choices'][0]['text']
                user_friendly_response = raw_response.split('A: ')[-1].strip()

                # Store formatted response to ensure it wraps properly
                formatted_response = user_friendly_response
                
                st.session_state.conversation_history.append(f"Agent: {formatted_response}")
                st.session_state.debug_info = output  # Store the most recent debug info
            else:
                st.error("The model did not return a valid response. Please try again.")
    else:
        st.warning("Please enter a query.")

st.write("### Conversation History:")
for exchange in st.session_state.conversation_history:
    # Use st.write or st.text for automatic text wrapping
    st.text(exchange)

if st.button('Reset Conversation'):
    st.session_state.conversation_history = []
    st.session_state.debug_info = {}

# Button to show debug information
if st.button("Show Debug Information"):
    display_debug_info()