Spaces:
Sleeping
Sleeping
File size: 8,445 Bytes
09aa2b8 |
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 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 |
import os
import sys
import logging
import uuid
from typing import List, Dict, Any, Tuple
from logging.handlers import RotatingFileHandler
import gradio as gr
from tenacity import retry, stop_after_attempt, wait_exponential
from core.rag_engine import RAGEngine
from core.user_profile import UserProfile
from config.config import settings
# ======================================================================================
# Logging Setup
# ======================================================================================
os.makedirs("logs", exist_ok=True)
logging.basicConfig(
level=getattr(logging, settings.LOG_LEVEL.upper(), logging.INFO),
format=settings.LOG_FORMAT,
handlers=[
logging.StreamHandler(sys.stdout),
RotatingFileHandler(
settings.LOG_FILE_PATH,
maxBytes=settings.LOG_FILE_MAX_BYTES,
backupCount=settings.LOG_FILE_BACKUP_COUNT
),
],
)
logger = logging.getLogger(__name__)
# ======================================================================================
# Core Module Initialization
# ======================================================================================
@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
def initialize_with_retry(func):
"""Initializes a component with retry logic."""
try:
return func()
except Exception as e:
logger.error(f"Initialization failed: {e}", exc_info=True)
raise
try:
user_profile = initialize_with_retry(UserProfile)
rag_engine = initialize_with_retry(lambda: RAGEngine(user_profile=user_profile))
logger.info("Core modules initialized successfully.")
except Exception as e:
logger.critical(f"Fatal: Could not initialize core modules: {e}. Exiting.", exc_info=True)
sys.exit(1)
# ======================================================================================
# Business Logic
# ======================================================================================
async def handle_chat_interaction(
message: str, chat_history: List[List[str]], user_id: str, categories: List[str]
) -> List[List[str]]:
"""Handles the user's chat message, processes it, and updates the history."""
if not message.strip():
gr.Warning("Message cannot be empty. Please type a question.")
return chat_history
try:
profile = user_profile.get_profile(user_id)
profile["preferences"]["favorite_categories"] = categories
user_profile.update_profile(user_id, profile)
logger.info(f"Updated preferences for user {user_id}: {categories}")
result = await rag_engine.process_query(query=message, user_id=user_id)
response = result.get("answer", "Sorry, I could not find an answer.")
sources = result.get("sources")
if sources:
response += "\n\n**Sources:**\n" + format_sources(sources)
chat_history.append((message, response))
logger.info(f"User {user_id} received response.")
return chat_history
except Exception as e:
error_message = f"An unexpected error occurred: {str(e)}"
logger.error(f"Error for user {user_id}: {error_message}", exc_info=True)
gr.Warning("Sorry, I encountered a problem. Please try again or rephrase your question.")
return chat_history
def format_sources(sources: List[Dict[str, Any]]) -> str:
"""Formats the source documents into a readable string."""
if not sources:
return ""
formatted_list = [f"- **{source.get('title', 'Unknown Source')}** (Category: {source.get('category', 'N/A')})" for source in sources]
return "\n".join(formatted_list)
# ======================================================================================
# Gradio UI Definition
# ======================================================================================
def handle_slider_change(value: int) -> None:
"""
Handles the change event for the document loader slider.
Note: This currently only shows a notification. A restart is required.
"""
gr.Info(f"Document limit set to {int(value)}. Please restart the app for changes to take effect.")
def create_interface() -> gr.Blocks:
"""Creates and configures the Gradio web interface."""
with gr.Blocks(
title="TravelMate - Your AI Travel Assistant",
theme=gr.themes.Base(),
) as demo:
user_id = gr.State(lambda: str(uuid.uuid4()))
gr.Markdown("""
<div style="text-align: center;">
<h1 style="font-size: 2.5em;">✈️ TravelMate</h1>
<p style="font-size: 1.1em; color: #333;">Your AI-powered travel assistant. Ask me anything to plan your next trip!</p>
</div>
""")
with gr.Accordion("Advanced Settings", open=False):
doc_load_slider = gr.Slider(
minimum=100,
maximum=5000,
value=settings.MAX_DOCUMENTS_TO_LOAD,
step=100,
label="Documents to Load",
info="Controls how many documents are loaded for the RAG engine. Higher values may increase startup time.",
)
doc_load_slider.change(
fn=handle_slider_change, inputs=[doc_load_slider], outputs=None
)
with gr.Row():
with gr.Column(scale=2):
chatbot = gr.Chatbot(
elem_id="chatbot",
label="TravelMate Chat",
height=600,
show_label=False,
show_copy_button=True,
bubble_full_width=False,
avatar_images=("assets/user_avatar.png", "assets/bot_avatar.png"),
)
with gr.Row():
msg = gr.Textbox(
placeholder="Ask me about destinations, flights, hotels...",
show_label=False,
container=False,
scale=8,
)
submit_btn = gr.Button("Send", variant="primary", scale=1)
with gr.Column(scale=1):
gr.Markdown("### Select Your Interests")
categories = gr.CheckboxGroup(
choices=[
"Flights", "Hotels", "Destinations", "Activities",
"Transportation", "Food & Dining", "Shopping",
"Health & Safety", "Budget Planning"
],
value=["Flights", "Hotels"],
label="Travel Categories",
)
gr.Markdown("### Example Questions")
gr.Examples(
examples=[
"What are the best places to visit in Japan?",
"How do I find cheap flights to Europe?",
"What should I pack for a beach vacation?",
"Tell me about local customs in Thailand",
"What's the best time to visit Paris?",
],
inputs=msg,
)
async def on_submit(message: str, history: List[List[str]], uid: str, cats: List[str]) -> Tuple[str, List[List[str]]]:
"""Handles submission and returns updated values for the message box and chatbot."""
updated_history = await handle_chat_interaction(message, history, uid, cats)
return "", updated_history
submit_btn.click(on_submit, [msg, chatbot, user_id, categories], [msg, chatbot])
msg.submit(on_submit, [msg, chatbot, user_id, categories], [msg, chatbot])
return demo
# ======================================================================================
# Application Launch
# ======================================================================================
if __name__ == "__main__":
try:
app = create_interface()
app.queue(default_concurrency_limit=settings.GRADIO_CONCURRENCY_COUNT)
app.launch(
server_name=settings.GRADIO_SERVER_NAME,
server_port=settings.GRADIO_SERVER_PORT,
share=settings.GRADIO_SHARE,
show_error=True,
show_api=False,
)
except Exception as e:
logger.critical(f"Failed to launch Gradio app: {e}", exc_info=True)
raise |