import logging
import gradio as gr
import configs.config as config
import services.scraper
import stores.chroma
from llm_setup.llm_setup import LLMService
from caching.lfu import LFUCache
from configs.config import MODEL_REGISTRY, DEFAULT_PROVIDER
import time
logger = logging.getLogger() # Create a logger object
logger.setLevel(logging.INFO) # Set the logging level to INFO
config.set_envs() # Set environment variables using the config module
store = stores.chroma.ChromaDB(config.EMBEDDINGS)
service = services.scraper.Service(store)
# Scrape data and get the store vector retriever
service.scrape_and_get_store_vector_retriever(config.URLS)
# Initialize the LLMService with logger, prompt, and store vector retriever
llm_svc = LLMService(
logger=logger,
system_prompt=config.SYSTEM_PROMPT,
web_retriever=store.get_chroma_instance().as_retriever(),
provider=config.DEFAULT_PROVIDER,
llm_model_name=config.LLM_MODEL_NAME,
)
def respond(user_input,session_hash):
if user_input == "clear_chat_history_aisdb_override":
llm_svc.store={}
return "Memory Cache cleared"
response = llm_svc.conversational_rag_chain().invoke(
{"input": user_input},
config={"configurable": {"session_id": session_hash}},
)["answer"]
return response
def echo(text, chat_history, request: gr.Request):
if request:
session_hash = request.session_hash
resp = respond(text, session_hash)
for i in range(len(resp)):
time.sleep(0.01)
yield resp[: i + 1]
else:
return "No request object received."
def on_reset_button_click():
llm_svc.store=LFUCache(capacity=50)
def on_apply_model(provider, model_name, api_key):
key = api_key.strip() if api_key and api_key.strip() else None
try:
llm_svc.update_llm(provider, model_name, key)
return f"Switched to {provider} / {model_name}"
except Exception as e:
return f"Error: {str(e)}"
def on_provider_change(provider):
models = MODEL_REGISTRY.get(provider, [])
return gr.update(choices=models, value=models[0] if models else None)
# --- Maritime Theme ---
maritime_blue = gr.themes.Color(
c50="#f0f9ff", c100="#e0f2fe", c200="#b9e6fe", c300="#7dd4fc",
c400="#38bdf8", c500="#0ea5e9", c600="#0284c7", c700="#0369a1",
c800="#075985", c900="#0c4a6e", c950="#082f49",
name="maritime-blue",
)
teal_accent = gr.themes.Color(
c50="#f0fdfa", c100="#ccfbf1", c200="#99f6e4", c300="#5eead4",
c400="#2dd4bf", c500="#14b8a6", c600="#0d9488", c700="#0f766e",
c800="#115e59", c900="#134e4a", c950="#042f2e",
name="teal-accent",
)
try:
stormy_theme = gr.themes.Ocean(
primary_hue=maritime_blue,
secondary_hue=teal_accent,
neutral_hue="slate",
spacing_size=gr.themes.sizes.spacing_md,
radius_size=gr.themes.sizes.radius_lg,
text_size=gr.themes.sizes.text_md,
font=(gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"),
font_mono=(gr.themes.GoogleFont("JetBrains Mono"), "ui-monospace", "Consolas", "monospace"),
)
except AttributeError:
stormy_theme = gr.themes.Soft(
primary_hue=maritime_blue,
secondary_hue=teal_accent,
neutral_hue="slate",
spacing_size=gr.themes.sizes.spacing_md,
radius_size=gr.themes.sizes.radius_lg,
text_size=gr.themes.sizes.text_md,
font=(gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"),
font_mono=(gr.themes.GoogleFont("JetBrains Mono"), "ui-monospace", "Consolas", "monospace"),
)
stormy_theme = stormy_theme.set(
body_background_fill="#f0f9ff",
body_background_fill_dark="#0c1929",
body_text_color="#0c4a6e",
body_text_color_dark="#e0f2fe",
block_background_fill="#ffffff",
block_background_fill_dark="#0f2942",
block_border_color="#b9e6fe",
block_border_color_dark="#0369a1",
button_primary_background_fill="linear-gradient(135deg, #0ea5e9, #0d9488)",
button_primary_background_fill_hover="linear-gradient(135deg, #38bdf8, #14b8a6)",
button_primary_background_fill_dark="linear-gradient(135deg, #0369a1, #0f766e)",
button_primary_text_color="#ffffff",
button_secondary_background_fill="#e0f2fe",
button_secondary_background_fill_hover="#b9e6fe",
button_secondary_background_fill_dark="#0f2942",
button_secondary_text_color="#0c4a6e",
button_secondary_text_color_dark="#7dd4fc",
input_background_fill="#f8fafc",
input_background_fill_dark="#0f2942",
input_border_color="#b9e6fe",
input_border_color_focus="#0ea5e9",
input_border_color_dark="#0369a1",
shadow_drop="0 2px 8px rgba(14, 165, 233, 0.08)",
shadow_drop_lg="0 4px 16px rgba(14, 165, 233, 0.12)",
)
custom_css = """
.stormy-header {
text-align: center;
padding: 1.5rem 1rem 1rem 1rem;
background: linear-gradient(135deg, #0c4a6e 0%, #0ea5e9 50%, #0d9488 100%);
border-radius: 12px;
margin-bottom: 0.5rem;
color: white;
}
.stormy-header h1 {
font-size: 1.8rem;
margin: 0 0 0.25rem 0;
font-weight: 700;
color: #ffffff !important;
}
.stormy-header p {
font-size: 0.95rem;
margin: 0;
color: #e0f2fe !important;
opacity: 0.9;
}
.reset-btn {
max-width: 200px !important;
}
.stormy-footer {
text-align: center;
font-size: 0.8rem;
color: #64748b;
padding-top: 0.5rem;
}
"""
if __name__ == '__main__':
logging.info("Starting AIVIz Bot")
with gr.Blocks(theme=stormy_theme, css=custom_css, title="Stormy - AISdb Assistant") as demo:
# Branding Header
gr.Markdown(
"""
""",
elem_id="header",
)
# Chat Interface
chatbot = gr.Chatbot(
placeholder=(
"Welcome aboard!
"
"I'm Stormy, your AISdb documentation assistant.
"
"Ask me about vessel tracking, data queries, or machine learning with AIS data."
),
height=500,
type="messages",
show_copy_button=True,
)
gr.ChatInterface(
fn=echo,
type="messages",
chatbot=chatbot,
textbox=gr.Textbox(
placeholder="Ask Stormy about AISdb...",
container=False,
scale=7,
),
examples=[
"How do I get started with AISdb?",
"How can I query vessel tracks by MMSI?",
"What machine learning models work with AIS data?",
"How do I visualize ship trajectories on a map?",
],
)
# Action Bar
with gr.Row():
with gr.Column(scale=3):
with gr.Accordion("About Stormy & AISdb", open=False):
gr.Markdown(
"""
**Stormy** is an AI assistant built on the AISdb (Automatic Identification System Database)
documentation. It can help you with:
- **Data Access**: Loading AIS data, creating databases, CSV export
- **Querying**: SQL queries, filtering by MMSI, time ranges, geographic areas
- **Processing**: Data cleaning, track interpolation, decimation
- **Visualization**: Plotting vessel trajectories, hexagon discretization
- **Machine Learning**: Seq2Seq models, autoencoders for AIS data
- **Geospatial**: Haversine distance, shore distance, bathymetric data
Powered by AISdb documentation from [aisviz.gitbook.io](https://aisviz.gitbook.io/documentation)
and [MAPS Lab](https://mapslab.tech/).
"""
)
with gr.Column(scale=1, min_width=200):
reset_button = gr.Button(
"Reset Chat Memory",
variant="secondary",
size="sm",
elem_classes=["reset-btn"],
)
reset_button.click(on_reset_button_click)
# Model Settings Panel
with gr.Accordion("Model Settings", open=False):
with gr.Row():
provider_dropdown = gr.Dropdown(
choices=list(MODEL_REGISTRY.keys()),
value=DEFAULT_PROVIDER,
label="Provider",
interactive=True,
scale=1,
)
model_dropdown = gr.Dropdown(
choices=MODEL_REGISTRY[DEFAULT_PROVIDER],
value=config.LLM_MODEL_NAME,
label="Model",
interactive=True,
scale=1,
)
with gr.Row():
api_key_input = gr.Textbox(
label="API Key (optional override)",
placeholder="Leave blank to use environment variable",
type="password",
scale=3,
)
apply_button = gr.Button(
"Apply",
variant="primary",
size="sm",
scale=1,
)
status_text = gr.Textbox(
label="Status",
interactive=False,
value=f"Active: {DEFAULT_PROVIDER} / {config.LLM_MODEL_NAME}",
max_lines=1,
)
provider_dropdown.change(
fn=on_provider_change,
inputs=[provider_dropdown],
outputs=[model_dropdown],
)
apply_button.click(
fn=on_apply_model,
inputs=[provider_dropdown, model_dropdown, api_key_input],
outputs=[status_text],
)
# Footer
gr.Markdown(
''
)
demo.launch()