File size: 7,224 Bytes
a3e0475
fc5f1c7
d6f5773
fc5f1c7
a3e0475
 
5c095c6
 
073538f
5c095c6
 
 
 
 
 
 
 
 
 
 
 
 
 
94ac9e7
a3e0475
5c095c6
 
 
 
 
 
fc5f1c7
5c095c6
 
 
a3e0475
 
 
5c095c6
 
a3e0475
 
94ac9e7
 
 
 
 
5c095c6
94ac9e7
5c095c6
073538f
fc5f1c7
 
5c095c6
fc5f1c7
 
 
 
 
 
 
 
5c095c6
fc5f1c7
 
5c095c6
 
 
fc5f1c7
5c095c6
 
fc5f1c7
5c095c6
 
 
 
fc5f1c7
 
 
 
5c095c6
 
2a239ae
5c095c6
fc5f1c7
 
 
5c095c6
2a239ae
5c095c6
2a239ae
 
ad0b8d6
5c095c6
 
 
 
ad0b8d6
2a239ae
 
 
5c095c6
2a239ae
 
 
fc5f1c7
 
5c095c6
fc5f1c7
 
 
 
5c095c6
2a239ae
5c095c6
 
 
2a239ae
5c095c6
2a239ae
 
5c095c6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f543f0b
ad0b8d6
5c095c6
ad0b8d6
 
5c095c6
ad0b8d6
b8a80ad
5c095c6
fc5f1c7
5c095c6
 
fc5f1c7
5c095c6
 
fc5f1c7
5c095c6
 
 
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
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
import os
import requests
import streamlit as st
from langchain_huggingface import HuggingFaceEndpoint
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from transformers import pipeline
from config import NASA_API_KEY  # Ensure this file exists with your NASA API Key

# Set up Streamlit UI
st.set_page_config(page_title="HAL - NASA ChatBot", page_icon="🚀")

# --- Ensure Session State Variables are Initialized ---
if "chat_history" not in st.session_state:
    st.session_state.chat_history = [{"role": "assistant", "content": "Hello! How can I assist you today?"}]

if "response_ready" not in st.session_state:
    st.session_state.response_ready = False  # Tracks whether HAL has responded

if "follow_up" not in st.session_state:
    st.session_state.follow_up = ""  # Stores follow-up question

# --- Set Up Model & API Functions ---
model_id = "mistralai/Mistral-7B-Instruct-v0.3"

# Initialize sentiment analysis pipeline with explicit model specification
sentiment_analyzer = pipeline(
    "sentiment-analysis",
    model="distilbert/distilbert-base-uncased-finetuned-sst-2-english",
    revision="714eb0f"
)

def get_llm_hf_inference(model_id=model_id, max_new_tokens=128, temperature=0.7):
    # Explicitly specify task="text-generation" so that the endpoint knows which task to run
    return HuggingFaceEndpoint(
        repo_id=model_id,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        token=os.getenv("HF_TOKEN"),
        task="text-generation"
    )

def get_nasa_apod():
    url = f"https://api.nasa.gov/planetary/apod?api_key={NASA_API_KEY}"
    response = requests.get(url)
    if response.status_code == 200:
        data = response.json()
        return data.get("url", ""), data.get("title", ""), data.get("explanation", "")
    else:
        return "", "NASA Data Unavailable", "I couldn't fetch data from NASA right now. Please try again later."

def analyze_sentiment(user_text):
    result = sentiment_analyzer(user_text)[0]
    return result['label']

def predict_action(user_text):
    if "NASA" in user_text or "space" in user_text:
        return "nasa_info"
    return "general_query"

def generate_follow_up(user_text):
    """
    Generates a concise and conversational follow-up question related to the user's input.
    """
    prompt_text = (
        f"Given the user's question: '{user_text}', generate a SHORT and SIMPLE follow-up question. "
        "Make it conversational and friendly. Example: 'Would you like to learn more about the six types of quarks?' "
        "Do NOT provide long explanations—just ask a friendly follow-up question."
    )
    hf = get_llm_hf_inference(max_new_tokens=32, temperature=0.7)
    return hf.invoke(input=prompt_text).strip()

def get_response(system_message, chat_history, user_text, max_new_tokens=256):
    """
    Generates HAL's response, making it more conversational and engaging.
    """
    sentiment = analyze_sentiment(user_text)
    action = predict_action(user_text)

    if action == "nasa_info":
        nasa_url, nasa_title, nasa_explanation = get_nasa_apod()
        response = f"**{nasa_title}**\n\n{nasa_explanation}"
        chat_history.append({'role': 'user', 'content': user_text})
        chat_history.append({'role': 'assistant', 'content': response})
        
        follow_up = generate_follow_up(user_text)
        chat_history.append({'role': 'assistant', 'content': follow_up})
        return response, follow_up, chat_history, nasa_url

    hf = get_llm_hf_inference(max_new_tokens=max_new_tokens, temperature=0.9)

    prompt = PromptTemplate.from_template(
        (
            "[INST] {system_message}\n\nCurrent Conversation:\n{chat_history}\n\nUser: {user_text}.\n [/INST]\n"
            "AI: Keep responses conversational and engaging. Start with a friendly phrase like "
            "'Certainly!', 'Of course!', or 'Great question!' before answering. "
            "Keep responses concise but engaging.\nHAL:"
        )
    )
    chat = prompt | hf.bind(skip_prompt=True) | StrOutputParser(output_key='content')
    response = chat.invoke(input=dict(system_message=system_message, user_text=user_text, chat_history=chat_history))
    response = response.split("HAL:")[-1].strip()

    chat_history.append({'role': 'user', 'content': user_text})
    chat_history.append({'role': 'assistant', 'content': response})

    if sentiment == "NEGATIVE":
        response = "I'm here to help. Let me know what I can do for you. 😊"

    follow_up = generate_follow_up(user_text)
    chat_history.append({'role': 'assistant', 'content': follow_up})

    return response, follow_up, chat_history, None

# --- Chat UI ---
st.title("🚀 HAL - Your NASA AI Assistant")
st.markdown("🌌 *Ask me about space, NASA, and beyond!*")

# Sidebar: Reset Chat
if st.sidebar.button("Reset Chat"):
    st.session_state.chat_history = [{"role": "assistant", "content": "Hello! How can I assist you today?"}]
    st.session_state.response_ready = False
    st.session_state.follow_up = ""
    st.experimental_rerun()

# Custom Chat Styling
st.markdown("""
    <style>
    .user-msg {
        background-color: #0078D7;
        color: white;
        padding: 10px;
        border-radius: 10px;
        margin-bottom: 5px;
        width: fit-content;
        max-width: 80%;
    }
    .assistant-msg {
        background-color: #333333;
        color: white;
        padding: 10px;
        border-radius: 10px;
        margin-bottom: 5px;
        width: fit-content;
        max-width: 80%;
    }
    .container {
        display: flex;
        flex-direction: column;
        align-items: flex-start;
    }
    @media (max-width: 600px) {
        .user-msg, .assistant-msg { font-size: 16px; max-width: 100%; }
    }
    </style>
""", unsafe_allow_html=True)

# Chat History Display
st.markdown("<div class='container'>", unsafe_allow_html=True)
for message in st.session_state.chat_history:
    if message["role"] == "user":
        st.markdown(f"<div class='user-msg'><strong>You:</strong> {message['content']}</div>", unsafe_allow_html=True)
    else:
        st.markdown(f"<div class='assistant-msg'><strong>HAL:</strong> {message['content']}</div>", unsafe_allow_html=True)
st.markdown("</div>", unsafe_allow_html=True)

# --- Single Input Box for Both Initial and Follow-Up Messages ---
user_input = st.chat_input("Type your message here...")  # Only ONE chat_input()

if user_input:
    response, follow_up, st.session_state.chat_history, image_url = get_response(
        system_message="You are a helpful AI assistant.",
        user_text=user_input,
        chat_history=st.session_state.chat_history
    )

    st.markdown(f"<div class='assistant-msg'><strong>HAL:</strong> {response}</div>", unsafe_allow_html=True)

    if image_url:
        st.image(image_url, caption="NASA Image of the Day")

    st.session_state.follow_up = follow_up
    st.session_state.response_ready = True

if st.session_state.response_ready and st.session_state.follow_up:
    st.markdown(f"<div class='assistant-msg'><strong>HAL:</strong> {st.session_state.follow_up}</div>", unsafe_allow_html=True)
    st.session_state.response_ready = False