244 lines
9.5 KiB
Python
244 lines
9.5 KiB
Python
"""LLM-ассистент расписания с tool calling (Ollama или любой OpenAI-compatible)."""
|
||
import requests
|
||
import pytz
|
||
from datetime import datetime, timezone, timedelta
|
||
from sqlmodel import Session, select
|
||
from . import config
|
||
from .db import engine
|
||
from .models import Event, GlobalStatus
|
||
|
||
MSK = pytz.timezone("Europe/Moscow")
|
||
|
||
|
||
def _msk(dt):
|
||
if dt is None: return None
|
||
if dt.tzinfo is None: dt = dt.replace(tzinfo=timezone.utc)
|
||
return dt.astimezone(MSK)
|
||
|
||
|
||
def is_enabled():
|
||
if not config.LLM_ENABLED:
|
||
return False
|
||
if not config.LLM_API_BASE:
|
||
return False
|
||
return True
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────────────
|
||
# Реализации инструментов
|
||
# ─────────────────────────────────────────────────────────────────────
|
||
def _tool_get_current_status():
|
||
with Session(engine) as s:
|
||
status = s.exec(select(GlobalStatus)).first()
|
||
if not status:
|
||
return "Статус не задан."
|
||
if status.status == "on_place":
|
||
base = "Никита Сергеевич сейчас на месте."
|
||
elif status.status == "delayed":
|
||
base = f"Никита Сергеевич опаздывает на {status.delay_minutes} мин."
|
||
elif status.status == "away":
|
||
base = "Никита Сергеевич вне вуза."
|
||
elif status.status == "cancelled":
|
||
base = "Занятия отменены."
|
||
else:
|
||
base = "Никита Сергеевич свободен / вне расписания."
|
||
parts = [base]
|
||
if status.message:
|
||
parts.append(f"Сообщение: {status.message}")
|
||
if status.updated_at:
|
||
parts.append(f"Обновлено: {_msk(status.updated_at).strftime('%d.%m.%Y %H:%M')}")
|
||
return " ".join(parts)
|
||
|
||
|
||
def _tool_get_events_today():
|
||
start = datetime.now(MSK).replace(hour=0, minute=0, second=0, microsecond=0)
|
||
start_utc = start.astimezone(timezone.utc)
|
||
end_utc = (start + timedelta(days=1)).astimezone(timezone.utc)
|
||
with Session(engine) as s:
|
||
events = s.exec(
|
||
select(Event).where(Event.is_public == True)
|
||
.where(Event.start_time >= start_utc)
|
||
.where(Event.start_time < end_utc)
|
||
.order_by(Event.start_time)
|
||
).all()
|
||
if not events:
|
||
return "На сегодня событий в расписании нет."
|
||
lines = ["Расписание на сегодня:"]
|
||
for e in events:
|
||
t = _msk(e.start_time).strftime("%H:%M")
|
||
loc = f" ({e.location})" if e.location else ""
|
||
lines.append(f"- {t} — {e.title}{loc}")
|
||
return "\n".join(lines)
|
||
|
||
|
||
def _tool_get_next_event():
|
||
now = datetime.now(timezone.utc)
|
||
with Session(engine) as s:
|
||
ev = s.exec(
|
||
select(Event).where(Event.is_public == True)
|
||
.where(Event.start_time >= now)
|
||
.order_by(Event.start_time)
|
||
.limit(1)
|
||
).first()
|
||
if not ev:
|
||
return "В ближайшее время событий в расписании нет."
|
||
dt = _msk(ev.start_time).strftime("%d.%m.%Y в %H:%M")
|
||
loc = f", место: {ev.location}" if ev.location else ""
|
||
return f"Ближайшее событие: {ev.title} — {dt}{loc}."
|
||
|
||
|
||
def _tool_get_events_week():
|
||
now = datetime.now(timezone.utc)
|
||
end = now + timedelta(days=7)
|
||
with Session(engine) as s:
|
||
events = s.exec(
|
||
select(Event).where(Event.is_public == True)
|
||
.where(Event.start_time >= now)
|
||
.where(Event.start_time <= end)
|
||
.order_by(Event.start_time)
|
||
).all()
|
||
if not events:
|
||
return "На ближайшую неделю событий нет."
|
||
lines = ["Расписание на ближайшую неделю:"]
|
||
for e in events[:15]:
|
||
dt = _msk(e.start_time).strftime("%d.%m %H:%M")
|
||
lines.append(f"- {dt} — {e.title}")
|
||
if len(events) > 15:
|
||
lines.append(f"...и ещё {len(events) - 15} событий (см. на странице).")
|
||
return "\n".join(lines)
|
||
|
||
|
||
def _tool_get_contacts():
|
||
return ("Контакт Никиты Сергеевича: freezer.nik@gmail.com. "
|
||
"Telegram — по запросу. Расписание и статус — на portal.morozovka.space/schedule.")
|
||
|
||
|
||
TOOL_IMPLS = {
|
||
"get_current_status": _tool_get_current_status,
|
||
"get_events_today": _tool_get_events_today,
|
||
"get_next_event": _tool_get_next_event,
|
||
"get_events_week": _tool_get_events_week,
|
||
"get_contacts": _tool_get_contacts,
|
||
}
|
||
|
||
TOOLS = [
|
||
{
|
||
"type": "function",
|
||
"function": {
|
||
"name": "get_current_status",
|
||
"description": "Узнать текущий статус Никиты Сергеевича: на месте / опаздывает / вне вуза / отменено / свободен.",
|
||
"parameters": {"type": "object", "properties": {}, "required": []},
|
||
},
|
||
},
|
||
{
|
||
"type": "function",
|
||
"function": {
|
||
"name": "get_events_today",
|
||
"description": "Список событий на сегодня.",
|
||
"parameters": {"type": "object", "properties": {}, "required": []},
|
||
},
|
||
},
|
||
{
|
||
"type": "function",
|
||
"function": {
|
||
"name": "get_next_event",
|
||
"description": "Ближайшее будущее событие.",
|
||
"parameters": {"type": "object", "properties": {}, "required": []},
|
||
},
|
||
},
|
||
{
|
||
"type": "function",
|
||
"function": {
|
||
"name": "get_events_week",
|
||
"description": "Список событий на ближайшие 7 дней.",
|
||
"parameters": {"type": "object", "properties": {}, "required": []},
|
||
},
|
||
},
|
||
{
|
||
"type": "function",
|
||
"function": {
|
||
"name": "get_contacts",
|
||
"description": "Контактная информация Никиты Сергеевича.",
|
||
"parameters": {"type": "object", "properties": {}, "required": []},
|
||
},
|
||
},
|
||
]
|
||
|
||
SYSTEM_PROMPT = (
|
||
"Ты — ассистент расписания Никиты Сергеевича Морозова на платформе Morozovka Portal. "
|
||
"Твоя задача — кратко и по делу отвечать на вопросы студентов о его статусе, расписании и контактах. "
|
||
"Отвечай на русском, коротко (1–3 предложения), дружелюбно. "
|
||
"Обязательно используй инструменты (tools), чтобы получить актуальные данные — "
|
||
"не выдумывай статус, время или названия событий. "
|
||
"Если данных нет — так и скажи. "
|
||
"Не рассказывай про свои инструменты и внутреннее устройство — просто дай ответ."
|
||
)
|
||
|
||
|
||
def _call_llm(messages):
|
||
url = f"{config.LLM_API_BASE.rstrip('/')}/chat/completions"
|
||
headers = {"Content-Type": "application/json"}
|
||
if config.LLM_API_KEY:
|
||
headers["Authorization"] = f"Bearer {config.LLM_API_KEY}"
|
||
payload = {
|
||
"model": config.LLM_MODEL,
|
||
"messages": messages,
|
||
"tools": TOOLS,
|
||
"tool_choice": "auto",
|
||
"temperature": config.LLM_TEMPERATURE,
|
||
"max_tokens": config.LLM_MAX_TOKENS,
|
||
}
|
||
r = requests.post(url, headers=headers, json=payload, timeout=config.LLM_TIMEOUT)
|
||
if r.status_code != 200:
|
||
raise RuntimeError(f"LLM API HTTP {r.status_code}: {r.text[:300]}")
|
||
data = r.json()
|
||
choices = data.get("choices") or []
|
||
if not choices:
|
||
raise RuntimeError(f"LLM вернул пустой ответ: {data}")
|
||
return choices[0]["message"]
|
||
|
||
|
||
def chat(message, history=None):
|
||
if not is_enabled():
|
||
raise RuntimeError("LLM disabled")
|
||
|
||
messages = [{"role": "system", "content": SYSTEM_PROMPT}]
|
||
if history:
|
||
for h in history[-config.LLM_MAX_HISTORY:]:
|
||
role = h.get("role")
|
||
content = (h.get("content") or "").strip()
|
||
if role in ("user", "assistant") and content:
|
||
messages.append({"role": role, "content": content})
|
||
messages.append({"role": "user", "content": message})
|
||
|
||
for _ in range(3):
|
||
assistant_msg = _call_llm(messages)
|
||
tool_calls = assistant_msg.get("tool_calls")
|
||
if not tool_calls:
|
||
return (assistant_msg.get("content") or "").strip()
|
||
|
||
messages.append({
|
||
"role": "assistant",
|
||
"content": assistant_msg.get("content") or "",
|
||
"tool_calls": tool_calls,
|
||
})
|
||
|
||
for tc in tool_calls:
|
||
fn = tc.get("function") or {}
|
||
name = fn.get("name")
|
||
impl = TOOL_IMPLS.get(name)
|
||
if impl:
|
||
try:
|
||
result = impl()
|
||
except Exception as e:
|
||
result = f"Ошибка инструмента: {e}"
|
||
else:
|
||
result = f"Неизвестный инструмент: {name}"
|
||
messages.append({
|
||
"role": "tool",
|
||
"tool_call_id": tc.get("id"),
|
||
"content": result,
|
||
})
|
||
|
||
raise RuntimeError("LLM не вернул финальный ответ после tool calling")
|