353 lines
13 KiB
Python
353 lines
13 KiB
Python
import json
|
|
import sqlite3
|
|
import logging
|
|
from typing import List, Dict, Any, Optional
|
|
from datetime import datetime, timedelta
|
|
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
|
|
logger = logging.getLogger("DB_TOOLS")
|
|
|
|
DB_PATH = "/home/puh/scud_orion_ai_v2/data/scud_orion_ai.db"
|
|
|
|
def db_get_current_server_time() -> Dict[str, Any]:
|
|
now = datetime.now()
|
|
days_ru = ["понедельник", "вторник", "среда", "четверг", "пятница", "суббота", "воскресенье"]
|
|
return {
|
|
"current_date": now.strftime("%d.%m.%Y"),
|
|
"current_time": now.strftime("%H:%M:%S"),
|
|
"day_of_week": days_ru[now.weekday()],
|
|
"iso_date": now.strftime("%Y-%m-%d")
|
|
}
|
|
|
|
def smart_parse_date(date_str: Optional[str], original_user_message: str = "") -> Optional[str]:
|
|
"""
|
|
Дата уже точно подготовлена моделью на основе системного календаря.
|
|
Возвращаем date_str без повторной тяжелой фильтрации.
|
|
"""
|
|
return date_str
|
|
|
|
def db_save_chat_message(session_id: str, role: str, content: str):
|
|
if not content:
|
|
return
|
|
conn = get_db_connection()
|
|
cursor = conn.cursor()
|
|
cursor.execute("""
|
|
INSERT INTO chat_messages (session_id, role, content, created_at)
|
|
VALUES (?, ?, ?, CURRENT_TIMESTAMP)
|
|
""", (session_id, role, content))
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def db_get_chat_history(session_id: str, limit: int = 20) -> List[Dict[str, Any]]:
|
|
conn = get_db_connection()
|
|
cursor = conn.cursor()
|
|
cursor.execute("""
|
|
SELECT role, content FROM chat_messages
|
|
WHERE session_id = ?
|
|
ORDER BY id DESC LIMIT ?
|
|
""", (session_id, limit))
|
|
rows = cursor.fetchall()
|
|
conn.close()
|
|
return [{"role": r["role"], "content": r["content"]} for r in reversed(rows)]
|
|
|
|
def db_get_snapshots(session_id: str = "web_session_main", date_str: Optional[str] = None, original_user_message: str = "") -> Dict[str, Any]:
|
|
date_str = smart_parse_date(date_str, original_user_message)
|
|
|
|
conn = get_db_connection()
|
|
cursor = conn.cursor()
|
|
|
|
query = """
|
|
SELECT snapshot_id, log_date, snapshot_time, COUNT(*) as record_count
|
|
FROM scud_logs
|
|
"""
|
|
params = []
|
|
|
|
if date_str:
|
|
# Приводим дату ДД.ММ.ГГГГ к ISO YYYY-MM-DD
|
|
iso_date = date_str
|
|
if "." in date_str:
|
|
parts = date_str.split(".")
|
|
if len(parts) == 3:
|
|
iso_date = f"{parts[2]}-{parts[1]}-{parts[0]}"
|
|
|
|
# Строгий поиск: ищем совпадение строго по log_date или началу snapshot_time/created_at
|
|
query += " WHERE log_date = ? OR log_date = ? OR snapshot_time LIKE ? OR created_at LIKE ? "
|
|
params.extend([date_str, iso_date, f"{iso_date}%", f"{iso_date}%"])
|
|
|
|
query += " GROUP BY snapshot_id ORDER BY id DESC LIMIT 20"
|
|
|
|
cursor.execute(query, params)
|
|
rows = cursor.fetchall()
|
|
snapshots = [dict(r) for r in rows]
|
|
|
|
result_data = {
|
|
"query_date": date_str or "все",
|
|
"snapshots_count": len(snapshots),
|
|
"snapshots": snapshots
|
|
}
|
|
|
|
db_set_session_state(
|
|
session_id=session_id,
|
|
state_type="SNAPSHOTS_VIEW",
|
|
data=json.dumps(result_data, ensure_ascii=False)
|
|
)
|
|
|
|
conn.close()
|
|
return result_data
|
|
|
|
def get_db_connection():
|
|
conn = sqlite3.connect(DB_PATH, timeout=30.0)
|
|
conn.row_factory = sqlite3.Row
|
|
conn.execute("PRAGMA journal_mode = WAL;")
|
|
conn.execute("PRAGMA synchronous = NORMAL;")
|
|
return conn
|
|
|
|
def normalize_task_id(task_id_input: str) -> str:
|
|
if not task_id_input:
|
|
return ""
|
|
clean_id = str(task_id_input).strip().upper().replace("TASK-", "").replace("TASK", "")
|
|
if clean_id.isdigit():
|
|
num = int(clean_id)
|
|
return f"TASK-{(num):02d}" if num < 100 else f"TASK-{(num):03d}"
|
|
return f"TASK-{clean_id}"
|
|
|
|
def db_get_tasks(user_id: int) -> List[Dict[str, Any]]:
|
|
conn = get_db_connection()
|
|
cursor = conn.cursor()
|
|
cursor.execute("""
|
|
SELECT id, task_id, module, title, priority, status, due_date, created_at
|
|
FROM tasks
|
|
WHERE user_id = ?
|
|
ORDER BY id DESC
|
|
""", (user_id,))
|
|
rows = cursor.fetchall()
|
|
conn.close()
|
|
return [dict(r) for r in rows]
|
|
|
|
def db_add_task(user_id: int, module: str, title: str, priority: str = "MEDIUM", due_date: Optional[str] = None) -> Dict[str, Any]:
|
|
conn = get_db_connection()
|
|
cursor = conn.cursor()
|
|
|
|
cursor.execute("SELECT MAX(id) FROM tasks")
|
|
max_id = cursor.fetchone()[0] or 0
|
|
new_task_id = f"TASK-{(max_id + 1):02d}"
|
|
|
|
cursor.execute("""
|
|
INSERT INTO tasks (task_id, module, title, priority, status, due_date, user_id)
|
|
VALUES (?, ?, ?, ?, 'BACKLOG', ?, ?)
|
|
""", (new_task_id, module, title, priority.upper(), due_date, user_id))
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
return {"status": "success", "task_id": new_task_id, "message": f"Задача {new_task_id} создана"}
|
|
|
|
def db_update_task_status(user_id: int, task_id: str, status: str = "COMPLETED", due_date: Optional[str] = None) -> Dict[str, Any]:
|
|
conn = get_db_connection()
|
|
cursor = conn.cursor()
|
|
|
|
formatted_id = normalize_task_id(task_id)
|
|
|
|
if due_date:
|
|
cursor.execute("""
|
|
UPDATE tasks
|
|
SET status = ?, due_date = ?
|
|
WHERE (UPPER(task_id) = ? OR task_id LIKE ?) AND user_id = ?
|
|
""", (status.upper(), due_date, formatted_id, f"%{task_id.strip()}", user_id))
|
|
else:
|
|
cursor.execute("""
|
|
UPDATE tasks
|
|
SET status = ?
|
|
WHERE (UPPER(task_id) = ? OR task_id LIKE ?) AND user_id = ?
|
|
""", (status.upper(), formatted_id, f"%{task_id.strip()}", user_id))
|
|
|
|
if cursor.rowcount == 0:
|
|
conn.close()
|
|
return {"error": f"Задача {task_id} не найдена или принадлежит другому пользователю"}
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
return {"status": "success", "message": f"Статус задачи {formatted_id} обновлен на {status.upper()}"}
|
|
|
|
def db_delete_task(user_id: int, task_id: str) -> Dict[str, Any]:
|
|
conn = get_db_connection()
|
|
cursor = conn.cursor()
|
|
|
|
formatted_id = normalize_task_id(task_id)
|
|
|
|
cursor.execute("""
|
|
DELETE FROM tasks
|
|
WHERE (UPPER(task_id) = ? OR task_id LIKE ?) AND user_id = ?
|
|
""", (formatted_id, f"%{task_id.strip()}", user_id))
|
|
|
|
if cursor.rowcount == 0:
|
|
conn.close()
|
|
return {"error": f"Задача {task_id} не найдена"}
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
return {"status": "success", "message": f"Задача {formatted_id} удалена"}
|
|
|
|
def db_get_active_system_prompt() -> str:
|
|
conn = get_db_connection()
|
|
cursor = conn.cursor()
|
|
cursor.execute("SELECT prompt_text FROM system_prompts WHERE is_active = 1 ORDER BY id DESC LIMIT 1")
|
|
row = cursor.fetchone()
|
|
conn.close()
|
|
return row["prompt_text"] if row else "Ты — ИИ-ассистент SCUD Orion AI."
|
|
|
|
def db_add_system_prompt(name: str, prompt_text: str) -> Dict[str, Any]:
|
|
try:
|
|
with get_db_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute("BEGIN IMMEDIATE;")
|
|
|
|
cursor.execute("SELECT id FROM system_prompts WHERE name = ?", (name,))
|
|
existing = cursor.fetchone()
|
|
|
|
if existing:
|
|
cursor.execute(
|
|
"UPDATE system_prompts SET prompt_text = ?, is_active = 1, updated_at = CURRENT_TIMESTAMP WHERE name = ?",
|
|
(prompt_text, name)
|
|
)
|
|
else:
|
|
cursor.execute(
|
|
"INSERT INTO system_prompts (name, prompt_text, is_active) VALUES (?, ?, 1)",
|
|
(name, prompt_text)
|
|
)
|
|
|
|
conn.commit()
|
|
|
|
logger.info("Системный промпт успешно сохранен и применен в базе данных.")
|
|
return {"status": "success", "message": "Системный промпт успешно обновлен"}
|
|
except Exception as e:
|
|
logger.error(f"Ошибка при сохранении промпта в БД: {e}")
|
|
return {"status": "error", "error": str(e)}
|
|
|
|
def db_get_rules() -> List[Dict[str, Any]]:
|
|
conn = get_db_connection()
|
|
cursor = conn.cursor()
|
|
cursor.execute("SELECT id, rule_text FROM ai_knowledge_base ORDER BY id ASC")
|
|
rows = cursor.fetchall()
|
|
conn.close()
|
|
return [dict(r) for r in rows]
|
|
|
|
def db_set_session_state(session_id: str, state_type: str, data: str):
|
|
conn = get_db_connection()
|
|
cursor = conn.cursor()
|
|
cursor.execute("""
|
|
INSERT INTO session_states (session_id, state_type, pending_data, updated_at)
|
|
VALUES (?, ?, ?, CURRENT_TIMESTAMP)
|
|
ON CONFLICT(session_id) DO UPDATE SET
|
|
state_type = excluded.state_type,
|
|
pending_data = excluded.pending_data,
|
|
updated_at = CURRENT_TIMESTAMP
|
|
""", (session_id, state_type, data))
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def db_get_session_state(session_id: str) -> Optional[Dict[str, Any]]:
|
|
conn = get_db_connection()
|
|
cursor = conn.cursor()
|
|
cursor.execute("SELECT state_type, pending_data FROM session_states WHERE session_id = ?", (session_id,))
|
|
row = cursor.fetchone()
|
|
conn.close()
|
|
return dict(row) if row else None
|
|
|
|
def db_clear_session_state(session_id: str):
|
|
conn = get_db_connection()
|
|
cursor = conn.cursor()
|
|
cursor.execute("DELETE FROM session_states WHERE session_id = ?", (session_id,))
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def db_get_stats() -> Dict[str, Any]:
|
|
"""Возвращает общую статистику по количеству записей во всех таблицах БД."""
|
|
conn = get_db_connection()
|
|
cursor = conn.cursor()
|
|
tables = ['scud_logs', 'zup_staff', 'zup_absences', 'anomalies_history', 'ai_knowledge_base', 'system_prompts', 'session_states', 'tasks']
|
|
stats = {}
|
|
for t in tables:
|
|
try:
|
|
cursor.execute(f"SELECT COUNT(*) FROM {t}")
|
|
stats[t] = cursor.fetchone()[0]
|
|
except Exception:
|
|
stats[t] = 0
|
|
conn.close()
|
|
return {"status": "success", "tables_stats": stats}
|
|
|
|
def db_get_anomalies(limit: int = 100, date_str: Optional[str] = None) -> Dict[str, Any]:
|
|
"""Возвращает историю аномалий СКУД с опциональной фильтрацией по дате."""
|
|
conn = get_db_connection()
|
|
cursor = conn.cursor()
|
|
|
|
query = "SELECT anomaly_date, fio, anomaly_type, details FROM anomalies_history"
|
|
params = []
|
|
|
|
if date_str:
|
|
query += " WHERE anomaly_date = ?"
|
|
params.append(date_str)
|
|
|
|
query += " ORDER BY id DESC LIMIT ?"
|
|
params.append(limit)
|
|
|
|
cursor.execute(query, params)
|
|
rows = cursor.fetchall()
|
|
conn.close()
|
|
|
|
anomalies_list = [dict(r) for r in rows]
|
|
return {
|
|
"status": "success",
|
|
"count": len(anomalies_list),
|
|
"anomalies": anomalies_list
|
|
}
|
|
|
|
def db_get_session_states() -> Dict[str, Any]:
|
|
"""Возвращает список всех активных сессий и состояний превью."""
|
|
conn = get_db_connection()
|
|
cursor = conn.cursor()
|
|
cursor.execute("SELECT session_id, state_type, updated_at FROM session_states")
|
|
rows = cursor.fetchall()
|
|
conn.close()
|
|
return {"status": "success", "active_sessions": [dict(r) for r in rows]}
|
|
|
|
def db_delete_snapshots(snapshot_id: Optional[str] = None, day_str: Optional[str] = None) -> Dict[str, Any]:
|
|
"""Удаляет снапшот по ID или за конкретную дату."""
|
|
if not snapshot_id and not day_str:
|
|
return {"status": "error", "message": "Необходимо указать snapshot_id или day_str (ДД.ММ.ГГГГ)."}
|
|
|
|
conn = get_db_connection()
|
|
cursor = conn.cursor()
|
|
|
|
if snapshot_id:
|
|
cursor.execute("DELETE FROM scud_logs WHERE snapshot_id = ?", (snapshot_id,))
|
|
deleted = cursor.rowcount
|
|
else:
|
|
cursor.execute("DELETE FROM scud_logs WHERE log_date = ? OR snapshot_id LIKE ?", (day_str, f"%{day_str.replace('.', '')}%"))
|
|
deleted = cursor.rowcount
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
return {"status": "success", "message": f"Успешно удалено записей: {deleted}"}
|
|
|
|
def db_get_reference(category: Optional[str] = None) -> Dict[str, Any]:
|
|
"""Возвращает системные справочники и примеры команд для оператора."""
|
|
conn = get_db_connection()
|
|
cursor = conn.cursor()
|
|
|
|
query = "SELECT category, title, example_prompt, description FROM system_reference"
|
|
params = []
|
|
|
|
if category:
|
|
query += " WHERE category = ?"
|
|
params.append(category)
|
|
|
|
query += " ORDER BY id ASC"
|
|
cursor.execute(query, params)
|
|
rows = cursor.fetchall()
|
|
conn.close()
|
|
|
|
return {
|
|
"status": "success",
|
|
"count": len(rows),
|
|
"reference_items": [dict(r) for r in rows]
|
|
} |