chatbot.py

import os
import json
from datetime import datetime
import pytz
import traceback
from dateutil import parser as date_parser

# Local imports
from supabase_utils import get_supabase_client

# Attempt to import Gemini
try:
    import google.generativeai as genai
    GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
    if GEMINI_API_KEY:
        genai.configure(api_key=GEMINI_API_KEY)
        GEMINI_AVAILABLE = True
    else:
        GEMINI_AVAILABLE = False
except ImportError:
    GEMINI_AVAILABLE = False

# Station names (dependency for retrieve_context_from_db)
STATION_NAMES = {
    "DT_0001": "인천", "DT_0065": "평택", "DT_0008": "안산", "DT_0067": "대산",
    "DT_0043": "보령", "DT_0002": "군산", "DT_0050": "목포", "DT_0017": "제주",
    "DT_0052": "여수", "DT_0025": "마산", "DT_0051": "부산", "DT_0037": "포항",
    "DT_0068": "위도"
}

def parse_intent_with_llm(message: str) -> dict:
    """LLM을 사용해 사용자 질문에서 의도를 분석하고 JSON으로 반환"""
    if not GEMINI_AVAILABLE:
        return {"error": "Gemini API를 사용할 수 없습니다. API 키를 확인하세요."}
    
    prompt = f"""
    당신은 사용자의 자연어 질문을 분석하여 JSON 객체로 변환하는 전문가입니다.
    질문에서 '관측소 이름', '원하는 정보', '시작 시간', '종료 시간'을 추출해주세요.
    현재 시간은 {datetime.now(pytz.timezone('Asia/Seoul')).strftime('%Y-%m-%d %H:%M:%S')} KST 입니다.

    - '원하는 정보'는 '특정 시간 조위' 또는 '구간 조위' 중 하나여야 합니다.
    - '시작 시간'과 '종료 시간'은 'YYYY-MM-DD HH:MM:SS' 형식으로 변환해주세요.
    - 단일 시간이면 시작과 종료 시간을 동일하게 설정하고, 구간이면 그에 맞게 설정하세요.
    - 관측소 이름이 없으면 '인천'을 기본값으로 사용하세요.

    [사용자 질문]: {message}
    [JSON 출력]:
    """
    try:
        model = genai.GenerativeModel('gemini-1.5-flash', generation_config={"response_mime_type": "application/json"})
        response = model.generate_content(prompt)
        return json.loads(response.text)
    except Exception as e:
        return {"error": f"LLM 의도 분석 중 오류 발생: {e}"} 

def retrieve_context_from_db(intent: dict) -> str:
    """분석된 의도를 바탕으로 데이터베이스에서 정보 검색"""
    supabase = get_supabase_client()
    if not supabase: 
        return "데이터베이스에 연결할 수 없습니다."

    if "error" in intent:
        return f"의도 분석에 실패했습니다: {intent['error']}"
    
    station_name = intent.get("관측소 이름", "인천")
    start_time_str = intent.get("시작 시간")
    end_time_str = intent.get("종료 시간")
    
    station_id = next((sid for sid, name in STATION_NAMES.items() if name == station_name), "DT_0001")

    if not start_time_str or not end_time_str:
        return "질문에서 시간 정보를 찾을 수 없습니다."

    try:
        start_time = date_parser.parse(start_time_str)
        end_time = date_parser.parse(end_time_str)

        start_query_str = start_time.strftime('%Y-%m-%d %H:%M:%S')
        end_query_str = end_time.strftime('%Y-%m-%d %H:%M:%S')
        
        result = supabase.table('tide_predictions')\
            .select('*')\
            .eq('station_id', station_id)\
            .gte('predicted_at', start_query_str)\
            .lte('predicted_at', end_query_str)\
            .order('predicted_at')\
            .execute()

        if result.data:
            info_text = f"'{station_name}'의 '{start_time_str}'부터 '{end_time_str}'까지 조위 정보입니다.\n\n"
            
            if len(result.data) > 10:
                levels = [d['final_tide_level'] for d in result.data]
                max_level = max(levels)
                min_level = min(levels)
                info_text += f"- 최고 조위: {max_level:.1f}cm\n- 최저 조위: {min_level:.1f}cm"
            else:
                for d in result.data:
                    time_kst = date_parser.parse(d['predicted_at']).strftime('%H:%M')
                    info_text += f"- {time_kst}: 최종 조위 {d['final_tide_level']:.1f}cm (잔차 {d['predicted_residual']:.1f}cm)\n"
            return info_text
        else:
            return "해당 기간의 예측 데이터를 찾을 수 없습니다. '통합 조위 예측' 탭에서 먼저 예측을 실행해주세요."

    except Exception as e:
        return f"데이터 검색 중 오류 발생: {traceback.format_exc()}"

def process_chatbot_query_with_llm(message: str, history: list) -> str:
    """최종 RAG 파이프라인"""
    if not GEMINI_AVAILABLE:
        return "Gemini API를 사용할 수 없습니다. API 키를 확인하세요."
        
    intent = parse_intent_with_llm(message)
    retrieved_data = retrieve_context_from_db(intent)
    
    prompt = f"""당신은 친절한 해양 조위 정보 전문가입니다. 주어진 [검색된 데이터]를 바탕으로 사용자의 [질문]에 대해 자연스러운 문장으로 답변해주세요.
    [검색된 데이터]: {retrieved_data}
    [사용자 질문]: {message}
    [답변]:"""
    
    try:
        model = genai.GenerativeModel('gemini-1.5-flash')
        response = model.generate_content(prompt)
        return response.text
    except Exception as e:
        return f"Gemini 답변 생성 중 오류가 발생했습니다: {e}"