273 lines
11 KiB
Python
273 lines
11 KiB
Python
import sqlite3
|
|
import os
|
|
import json
|
|
from typing import Dict, List, Any, Optional
|
|
from datetime import datetime
|
|
|
|
class DatabaseManager:
|
|
def __init__(self, db_path: str = "railway_data.db"):
|
|
self.db_path = db_path
|
|
self.initialize_db()
|
|
|
|
# 임시 데이터 생성 (실제 구현에서는 DB에서 로드)
|
|
self.sample_data = {
|
|
"train_schedules": {
|
|
"2212": {"formation": "24", "type": "일반", "line": "1호선"},
|
|
"2242": {"formation": "24", "type": "일반", "line": "1호선"},
|
|
"2312": {"formation": "31", "type": "급행", "line": "1호선"},
|
|
},
|
|
"train_formations": {
|
|
"24": {
|
|
"manufacturer": "현대로템",
|
|
"introduction_date": "2019-05",
|
|
"recent_faults": [
|
|
{"date": "2023-11-15", "car": "3", "component": "HVAC", "detail": "배기팬 고장", "status": "완료"},
|
|
{"date": "2023-12-03", "car": "5", "component": "출입문", "detail": "개방 지연", "status": "완료"},
|
|
],
|
|
"maintenance": [
|
|
{"type": "일상", "date": "2023-12-20", "details": "각종 센서 점검 및 배터리 교체"},
|
|
{"type": "월상", "date": "2023-11-25", "details": "공조장치 종합 점검"},
|
|
{"type": "중정비", "date": "2023-10-05", "details": "차량 종합 정밀검사"}
|
|
],
|
|
"reports": [
|
|
{"type": "동향", "date": "2023-12-01", "file": "reports/24_trend_202312.pdf"},
|
|
{"type": "조치결과", "date": "2023-11-20", "file": "reports/24_fix_20231120.pdf"}
|
|
]
|
|
},
|
|
"31": {
|
|
"manufacturer": "대우중공업",
|
|
"introduction_date": "2012-08",
|
|
"recent_faults": [
|
|
{"date": "2023-10-25", "car": "2", "component": "제동장치", "detail": "공기압 저하", "status": "완료"},
|
|
],
|
|
"maintenance": [
|
|
{"type": "일상", "date": "2023-12-15", "details": "전기장치 점검"},
|
|
{"type": "월상", "date": "2023-11-10", "details": "차체 및 대차 검사"},
|
|
{"type": "중정비", "date": "2023-06-20", "details": "차량 전체 검사 및 부품 교체"}
|
|
],
|
|
"reports": [
|
|
{"type": "동향", "date": "2023-12-01", "file": "reports/31_trend_202312.pdf"}
|
|
]
|
|
}
|
|
},
|
|
"fault_history": [
|
|
{"date": "2023-12-05", "formation": "15", "car": "3", "component": "HVAC", "detail": "냉방 불량", "status": "완료"},
|
|
{"date": "2023-12-01", "formation": "22", "car": "1", "component": "HVAC", "detail": "과열", "status": "완료"},
|
|
{"date": "2023-11-28", "formation": "17", "car": "4", "component": "HVAC", "detail": "배기팬 고장", "status": "완료"},
|
|
{"date": "2023-11-25", "formation": "24", "car": "3", "component": "HVAC", "detail": "배기팬 고장", "status": "완료"},
|
|
{"date": "2023-11-20", "formation": "08", "car": "2", "component": "HVAC", "detail": "제어기 오류", "status": "완료"},
|
|
{"date": "2023-11-15", "formation": "19", "car": "5", "component": "HVAC", "detail": "필터 교체 필요", "status": "완료"},
|
|
{"date": "2023-11-10", "formation": "28", "car": "6", "component": "HVAC", "detail": "모터 소음", "status": "완료"},
|
|
]
|
|
}
|
|
|
|
def initialize_db(self):
|
|
"""데이터베이스 초기화 및 필요한 테이블 생성"""
|
|
if not os.path.exists(self.db_path):
|
|
conn = sqlite3.connect(self.db_path)
|
|
cursor = conn.cursor()
|
|
|
|
# 열차다이아표 테이블
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS train_schedules (
|
|
train_number TEXT PRIMARY KEY,
|
|
formation_number TEXT,
|
|
train_type TEXT,
|
|
line TEXT
|
|
)
|
|
''')
|
|
|
|
# 편성데이터 테이블
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS train_formations (
|
|
formation_number TEXT PRIMARY KEY,
|
|
manufacturer TEXT,
|
|
introduction_date TEXT,
|
|
details TEXT
|
|
)
|
|
''')
|
|
|
|
# 고장이력 테이블
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS fault_history (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
date TEXT,
|
|
formation_number TEXT,
|
|
car_number TEXT,
|
|
component TEXT,
|
|
detail TEXT,
|
|
status TEXT
|
|
)
|
|
''')
|
|
|
|
# 대화 기록 테이블
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS conversation_history (
|
|
id TEXT PRIMARY KEY,
|
|
timestamp TEXT,
|
|
train_number TEXT,
|
|
formation_number TEXT,
|
|
content TEXT
|
|
)
|
|
''')
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
# 샘플 데이터 추가
|
|
self._add_sample_data()
|
|
|
|
def _add_sample_data(self):
|
|
"""샘플 데이터를 DB에 추가합니다."""
|
|
conn = sqlite3.connect(self.db_path)
|
|
cursor = conn.cursor()
|
|
|
|
# 열차다이아표 데이터
|
|
for train_number, data in self.sample_data["train_schedules"].items():
|
|
cursor.execute(
|
|
"INSERT OR REPLACE INTO train_schedules VALUES (?, ?, ?, ?)",
|
|
(train_number, data["formation"], data["type"], data["line"])
|
|
)
|
|
|
|
# 편성데이터
|
|
for formation_number, data in self.sample_data["train_formations"].items():
|
|
cursor.execute(
|
|
"INSERT OR REPLACE INTO train_formations VALUES (?, ?, ?, ?)",
|
|
(formation_number, data["manufacturer"], data["introduction_date"], json.dumps(data))
|
|
)
|
|
|
|
# 고장이력 데이터
|
|
for fault in self.sample_data["fault_history"]:
|
|
cursor.execute(
|
|
"INSERT OR IGNORE INTO fault_history (date, formation_number, car_number, component, detail, status) VALUES (?, ?, ?, ?, ?, ?)",
|
|
(fault["date"], fault["formation"], fault["car"], fault["component"], fault["detail"], fault["status"])
|
|
)
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def get_train_formation(self, train_number: str) -> Optional[str]:
|
|
"""열차번호로 편성번호를 조회합니다."""
|
|
conn = sqlite3.connect(self.db_path)
|
|
cursor = conn.cursor()
|
|
|
|
cursor.execute("SELECT formation_number FROM train_schedules WHERE train_number = ?", (train_number,))
|
|
result = cursor.fetchone()
|
|
|
|
conn.close()
|
|
|
|
if result:
|
|
return result[0]
|
|
return None
|
|
|
|
def get_formation_details(self, formation_number: str) -> Optional[Dict[str, Any]]:
|
|
"""편성번호로 편성 상세정보를 조회합니다."""
|
|
conn = sqlite3.connect(self.db_path)
|
|
cursor = conn.cursor()
|
|
|
|
cursor.execute("SELECT details FROM train_formations WHERE formation_number = ?", (formation_number,))
|
|
result = cursor.fetchone()
|
|
|
|
conn.close()
|
|
|
|
if result:
|
|
return json.loads(result[0])
|
|
return None
|
|
|
|
def get_fault_history(self, component: str = None, limit: int = 50) -> List[Dict[str, Any]]:
|
|
"""특정 컴포넌트의 고장이력을 조회합니다."""
|
|
conn = sqlite3.connect(self.db_path)
|
|
cursor = conn.cursor()
|
|
|
|
query = "SELECT * FROM fault_history"
|
|
params = []
|
|
|
|
if component:
|
|
query += " WHERE component = ?"
|
|
params.append(component)
|
|
|
|
query += " ORDER BY date DESC LIMIT ?"
|
|
params.append(limit)
|
|
|
|
cursor.execute(query, params)
|
|
results = cursor.fetchall()
|
|
|
|
conn.close()
|
|
|
|
fault_history = []
|
|
for row in results:
|
|
fault_history.append({
|
|
"id": row[0],
|
|
"date": row[1],
|
|
"formation": row[2],
|
|
"car": row[3],
|
|
"component": row[4],
|
|
"detail": row[5],
|
|
"status": row[6]
|
|
})
|
|
|
|
return fault_history
|
|
|
|
def save_conversation(self, conversation: Dict[str, Any]):
|
|
"""대화 기록을 저장합니다."""
|
|
conn = sqlite3.connect(self.db_path)
|
|
cursor = conn.cursor()
|
|
|
|
cursor.execute(
|
|
"INSERT OR REPLACE INTO conversation_history VALUES (?, ?, ?, ?, ?)",
|
|
(
|
|
conversation["id"],
|
|
datetime.fromtimestamp(conversation["start_time"]).isoformat() if conversation["start_time"] else None,
|
|
conversation["train_number"],
|
|
conversation["train_formation"],
|
|
json.dumps(conversation)
|
|
)
|
|
)
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def get_train_info(self, conversation: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""대화 컨텍스트를 기반으로 열차 정보를 검색합니다."""
|
|
train_info = {"found": False}
|
|
|
|
train_number = conversation.get("train_number")
|
|
if not train_number:
|
|
return train_info
|
|
|
|
# 편성번호 조회
|
|
formation_number = conversation.get("train_formation")
|
|
if not formation_number:
|
|
formation_number = self.get_train_formation(train_number)
|
|
|
|
if formation_number:
|
|
# 편성 상세정보 조회
|
|
formation_details = self.get_formation_details(formation_number)
|
|
if formation_details:
|
|
train_info.update({
|
|
"found": True,
|
|
"train_number": train_number,
|
|
"formation_number": formation_number,
|
|
"details": formation_details
|
|
})
|
|
|
|
return train_info
|
|
|
|
def get_fault_info(self, conversation: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""대화 컨텍스트를 기반으로 고장 정보를 검색합니다."""
|
|
fault_info = {"found": False, "faults": []}
|
|
|
|
if not conversation.get("faults"):
|
|
return fault_info
|
|
|
|
for fault_type in conversation["faults"]:
|
|
component = "HVAC" if fault_type.upper() == "HVAC" else fault_type
|
|
fault_history = self.get_fault_history(component)
|
|
|
|
if fault_history:
|
|
fault_info["found"] = True
|
|
fault_info["component"] = component
|
|
fault_info["faults"] = fault_history
|
|
break
|
|
|
|
return fault_info |