Spaces:
Running
Running
| 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( | |
| """ | |
| <div class="stormy-header"> | |
| <h1>Stormy - AISdb Assistant</h1> | |
| <p>Your maritime data companion. Ask about AIS vessel tracking, data processing, machine learning, and more.</p> | |
| </div> | |
| """, | |
| elem_id="header", | |
| ) | |
| # Chat Interface | |
| chatbot = gr.Chatbot( | |
| placeholder=( | |
| "<strong>Welcome aboard!</strong><br>" | |
| "I'm Stormy, your AISdb documentation assistant.<br>" | |
| "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( | |
| '<div class="stormy-footer">Built with Gradio & LangChain | AISdb Documentation Assistant</div>' | |
| ) | |
| demo.launch() |