388 lines
13 KiB
Python
388 lines
13 KiB
Python
from __future__ import annotations
|
||
|
||
import json
|
||
import time
|
||
from pathlib import Path
|
||
from typing import Any
|
||
from urllib import error, request
|
||
|
||
from config import BotConfig
|
||
from telegram_api import TelegramAPI
|
||
|
||
|
||
STATE_FILE = Path(__file__).resolve().parent.parent / ".new-qwen" / "telegram-state.json"
|
||
|
||
|
||
def load_state() -> dict[str, Any]:
|
||
if not STATE_FILE.exists():
|
||
return {"offset": None, "sessions": {}, "auth_flows": {}}
|
||
state = json.loads(STATE_FILE.read_text(encoding="utf-8"))
|
||
state.setdefault("sessions", {})
|
||
state.setdefault("auth_flows", {})
|
||
return state
|
||
|
||
|
||
def save_state(state: dict[str, Any]) -> None:
|
||
STATE_FILE.parent.mkdir(parents=True, exist_ok=True)
|
||
STATE_FILE.write_text(json.dumps(state, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
|
||
|
||
def post_json(url: str, payload: dict[str, Any]) -> dict[str, Any]:
|
||
data = json.dumps(payload).encode("utf-8")
|
||
req = request.Request(
|
||
url,
|
||
data=data,
|
||
headers={"Content-Type": "application/json"},
|
||
method="POST",
|
||
)
|
||
with request.urlopen(req, timeout=300) as response:
|
||
return json.loads(response.read().decode("utf-8"))
|
||
|
||
|
||
def get_json(url: str) -> dict[str, Any]:
|
||
with request.urlopen(url, timeout=60) as response:
|
||
return json.loads(response.read().decode("utf-8"))
|
||
|
||
|
||
def send_text_chunks(api: TelegramAPI, chat_id: int, text: str) -> None:
|
||
normalized = text or "Пустой ответ."
|
||
chunk_size = 3800
|
||
for start in range(0, len(normalized), chunk_size):
|
||
api.send_message(chat_id, normalized[start : start + chunk_size])
|
||
|
||
|
||
def summarize_event(event: dict[str, Any]) -> str | None:
|
||
event_type = event.get("type")
|
||
if event_type == "job_status":
|
||
return event.get("message")
|
||
if event_type == "model_request":
|
||
return "Думаю над ответом"
|
||
if event_type == "tool_call":
|
||
return f"Вызываю инструмент: {event.get('name')}"
|
||
if event_type == "tool_result":
|
||
result = event.get("result", {})
|
||
if isinstance(result, dict) and "error" in result:
|
||
return f"Инструмент {event.get('name')} завершился с ошибкой"
|
||
return f"Инструмент {event.get('name')} завершён"
|
||
if event_type == "error":
|
||
return f"Ошибка: {event.get('message')}"
|
||
return None
|
||
|
||
|
||
def get_auth_flow(state: dict[str, Any], chat_id: int) -> dict[str, Any] | None:
|
||
return state.setdefault("auth_flows", {}).get(str(chat_id))
|
||
|
||
|
||
def start_auth_flow(
|
||
api: TelegramAPI,
|
||
config: BotConfig,
|
||
state: dict[str, Any],
|
||
chat_id: int,
|
||
*,
|
||
force_new: bool = False,
|
||
) -> dict[str, Any]:
|
||
existing = get_auth_flow(state, chat_id)
|
||
now = time.time()
|
||
if existing and not force_new and existing.get("expires_at", 0) > now:
|
||
send_text_chunks(
|
||
api,
|
||
chat_id,
|
||
"Авторизация Qwen OAuth ещё не завершена.\n"
|
||
f"Открой ссылку:\n{existing['verification_uri_complete']}\n\n"
|
||
"Бот сам проверит завершение и продолжит работу.",
|
||
)
|
||
return existing
|
||
|
||
started = post_json(f"{config.server_url}/api/v1/auth/device/start", {})
|
||
flow = {
|
||
"flow_id": started["flow_id"],
|
||
"user_code": started.get("user_code"),
|
||
"verification_uri": started.get("verification_uri"),
|
||
"verification_uri_complete": started["verification_uri_complete"],
|
||
"expires_at": started["expires_at"],
|
||
"interval_seconds": started.get("interval_seconds", 2),
|
||
"next_poll_at": now + started.get("interval_seconds", 2),
|
||
"pending_messages": existing.get("pending_messages", []) if existing else [],
|
||
}
|
||
state.setdefault("auth_flows", {})[str(chat_id)] = flow
|
||
send_text_chunks(
|
||
api,
|
||
chat_id,
|
||
"Нужна авторизация Qwen OAuth.\n"
|
||
f"Открой ссылку:\n{flow['verification_uri_complete']}\n\n"
|
||
f"user_code: {flow.get('user_code')}\n"
|
||
"После подтверждения бот сам продолжит работу.",
|
||
)
|
||
return flow
|
||
|
||
|
||
def ensure_auth(api: TelegramAPI, config: BotConfig, state: dict[str, Any], chat_id: int) -> bool:
|
||
status = get_json(f"{config.server_url}/api/v1/auth/status")
|
||
if status.get("authenticated"):
|
||
return True
|
||
start_auth_flow(api, config, state, chat_id)
|
||
return False
|
||
|
||
|
||
def enqueue_pending_message(
|
||
state: dict[str, Any],
|
||
chat_id: int,
|
||
user_id: str,
|
||
session_key: str,
|
||
text: str,
|
||
) -> None:
|
||
flow = get_auth_flow(state, chat_id)
|
||
if not flow:
|
||
return
|
||
pending_messages = flow.setdefault("pending_messages", [])
|
||
pending_messages.append(
|
||
{
|
||
"user_id": user_id,
|
||
"session_key": session_key,
|
||
"text": text,
|
||
"created_at": int(time.time()),
|
||
}
|
||
)
|
||
|
||
|
||
def deliver_chat_message(
|
||
api: TelegramAPI,
|
||
config: BotConfig,
|
||
state: dict[str, Any],
|
||
chat_id: int,
|
||
user_id: str,
|
||
session_key: str,
|
||
text: str,
|
||
*,
|
||
delayed: bool = False,
|
||
) -> None:
|
||
session_id = state.setdefault("sessions", {}).get(session_key)
|
||
prefix = "Обрабатываю отложенный запрос..." if delayed else "Обрабатываю запрос..."
|
||
api.send_message(chat_id, prefix)
|
||
start_result = post_json(
|
||
f"{config.server_url}/api/v1/chat/start",
|
||
{
|
||
"session_id": session_id,
|
||
"user_id": user_id,
|
||
"message": text,
|
||
},
|
||
)
|
||
state["sessions"][session_key] = start_result["session_id"]
|
||
job_id = start_result["job_id"]
|
||
seen_seq = 0
|
||
sent_statuses: set[str] = set()
|
||
answer = None
|
||
while True:
|
||
poll_result = post_json(
|
||
f"{config.server_url}/api/v1/chat/poll",
|
||
{"job_id": job_id, "since_seq": seen_seq},
|
||
)
|
||
for event in poll_result.get("events", []):
|
||
seen_seq = max(seen_seq, int(event.get("seq", 0)))
|
||
summary = summarize_event(event)
|
||
if summary and summary not in sent_statuses:
|
||
api.send_message(chat_id, summary[:4000])
|
||
sent_statuses.add(summary)
|
||
if poll_result.get("status") == "completed":
|
||
answer = poll_result.get("answer")
|
||
state["sessions"][session_key] = poll_result["session_id"]
|
||
break
|
||
if poll_result.get("status") == "failed":
|
||
raise RuntimeError(poll_result.get("error") or "Chat job failed")
|
||
time.sleep(1.2)
|
||
|
||
answer = answer or "Пустой ответ от модели."
|
||
send_text_chunks(api, chat_id, answer)
|
||
|
||
|
||
def poll_auth_flow(
|
||
api: TelegramAPI,
|
||
config: BotConfig,
|
||
state: dict[str, Any],
|
||
chat_id: int,
|
||
*,
|
||
force: bool = False,
|
||
) -> bool:
|
||
flow = get_auth_flow(state, chat_id)
|
||
if not flow:
|
||
return False
|
||
now = time.time()
|
||
if flow.get("expires_at", 0) <= now:
|
||
state["auth_flows"].pop(str(chat_id), None)
|
||
api.send_message(chat_id, "OAuth flow истёк. Запусти /auth ещё раз.")
|
||
return False
|
||
if not force and now < flow.get("next_poll_at", 0):
|
||
return False
|
||
|
||
try:
|
||
result = post_json(
|
||
f"{config.server_url}/api/v1/auth/device/poll",
|
||
{"flow_id": flow["flow_id"]},
|
||
)
|
||
except error.HTTPError as exc:
|
||
body = exc.read().decode("utf-8", errors="replace")
|
||
state["auth_flows"].pop(str(chat_id), None)
|
||
send_text_chunks(
|
||
api,
|
||
chat_id,
|
||
"Не удалось завершить OAuth flow на сервере.\n"
|
||
f"Ответ сервера: {body}\n"
|
||
"Запусти /auth ещё раз.",
|
||
)
|
||
return False
|
||
|
||
if not result.get("done"):
|
||
interval = result.get("interval_seconds", flow.get("interval_seconds", 2))
|
||
flow["interval_seconds"] = interval
|
||
flow["next_poll_at"] = now + interval
|
||
return False
|
||
|
||
state["auth_flows"].pop(str(chat_id), None)
|
||
api.send_message(chat_id, "Qwen OAuth успешно настроен.")
|
||
for item in flow.get("pending_messages", []):
|
||
deliver_chat_message(
|
||
api,
|
||
config,
|
||
state,
|
||
chat_id,
|
||
item["user_id"],
|
||
item["session_key"],
|
||
item["text"],
|
||
delayed=True,
|
||
)
|
||
if not flow.get("pending_messages"):
|
||
api.send_message(chat_id, "Можно отправлять обычные сообщения.")
|
||
return True
|
||
|
||
|
||
def process_auth_flows(
|
||
api: TelegramAPI,
|
||
config: BotConfig,
|
||
state: dict[str, Any],
|
||
) -> None:
|
||
for chat_id_raw in list(state.setdefault("auth_flows", {}).keys()):
|
||
try:
|
||
poll_auth_flow(api, config, state, int(chat_id_raw), force=False)
|
||
except Exception as exc:
|
||
print(f"auth flow poll error for chat {chat_id_raw}: {exc}")
|
||
|
||
|
||
def handle_message(api: TelegramAPI, config: BotConfig, state: dict[str, Any], message: dict[str, Any]) -> None:
|
||
chat_id = message["chat"]["id"]
|
||
user_id = str(message.get("from", {}).get("id", chat_id))
|
||
text = (message.get("text") or "").strip()
|
||
if not text:
|
||
api.send_message(chat_id, "Поддерживаются только текстовые сообщения.")
|
||
return
|
||
|
||
session_key = f"{chat_id}:{user_id}"
|
||
session_id = state.setdefault("sessions", {}).get(session_key)
|
||
state.setdefault("auth_flows", {})
|
||
|
||
if text == "/start":
|
||
api.send_message(
|
||
chat_id,
|
||
"new-qwen bot готов.\nКоманды: /help, /auth, /status, /session, /clear.",
|
||
)
|
||
return
|
||
|
||
if text == "/help":
|
||
api.send_message(
|
||
chat_id,
|
||
"Команды:\n"
|
||
"/auth - начать Qwen OAuth\n"
|
||
"/auth_check [flow_id] - проверить авторизацию\n"
|
||
"/status - статус OAuth и сервера\n"
|
||
"/session - показать текущую сессию\n"
|
||
"/clear - очистить контекст",
|
||
)
|
||
return
|
||
|
||
if text == "/auth":
|
||
start_auth_flow(api, config, state, chat_id, force_new=True)
|
||
return
|
||
|
||
if text.startswith("/auth_check"):
|
||
parts = text.split(maxsplit=1)
|
||
if len(parts) == 2:
|
||
flow = get_auth_flow(state, chat_id)
|
||
if flow:
|
||
flow["flow_id"] = parts[1]
|
||
flow = get_auth_flow(state, chat_id)
|
||
if not flow:
|
||
api.send_message(chat_id, "Нет активного flow_id. Сначала вызови /auth.")
|
||
return
|
||
if not poll_auth_flow(api, config, state, chat_id, force=True):
|
||
api.send_message(chat_id, "Авторизация ещё не завершена. Повторите команду через пару секунд.")
|
||
return
|
||
|
||
if text == "/status":
|
||
status = get_json(f"{config.server_url}/api/v1/auth/status")
|
||
send_text_chunks(
|
||
api,
|
||
chat_id,
|
||
"Сервер доступен.\n"
|
||
f"OAuth: {'configured' if status.get('authenticated') else 'not configured'}\n"
|
||
f"resource_url: {status.get('resource_url')}\n"
|
||
f"expires_at: {status.get('expires_at')}",
|
||
)
|
||
return
|
||
|
||
if text == "/session":
|
||
if not session_id:
|
||
api.send_message(chat_id, "У этого чата ещё нет активной сессии.")
|
||
return
|
||
session = post_json(
|
||
f"{config.server_url}/api/v1/session/get",
|
||
{"session_id": session_id},
|
||
)
|
||
send_text_chunks(
|
||
api,
|
||
chat_id,
|
||
"session_id: {session_id}\nmessages: {count}\nlast_answer: {last_answer}".format(
|
||
session_id=session_id,
|
||
count=len(session.get("messages", [])),
|
||
last_answer=session.get("last_answer") or "-",
|
||
),
|
||
)
|
||
return
|
||
|
||
if text == "/clear":
|
||
if session_id:
|
||
post_json(f"{config.server_url}/api/v1/session/clear", {"session_id": session_id})
|
||
state["sessions"].pop(session_key, None)
|
||
api.send_message(chat_id, "Контекст сессии очищен.")
|
||
return
|
||
|
||
if not ensure_auth(api, config, state, chat_id):
|
||
enqueue_pending_message(state, chat_id, user_id, session_key, text)
|
||
api.send_message(chat_id, "Сообщение поставлено в очередь до завершения авторизации.")
|
||
return
|
||
|
||
deliver_chat_message(api, config, state, chat_id, user_id, session_key, text)
|
||
|
||
|
||
def main() -> None:
|
||
config = BotConfig.load()
|
||
api = TelegramAPI(config.token)
|
||
state = load_state()
|
||
print("new-qwen bot polling started")
|
||
while True:
|
||
try:
|
||
process_auth_flows(api, config, state)
|
||
updates = api.get_updates(state.get("offset"), config.poll_timeout)
|
||
for update in updates:
|
||
state["offset"] = update["update_id"] + 1
|
||
message = update.get("message")
|
||
if message:
|
||
handle_message(api, config, state, message)
|
||
save_state(state)
|
||
save_state(state)
|
||
except Exception as exc:
|
||
print(f"bot loop error: {exc}")
|
||
time.sleep(3)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|