277 lines
9.8 KiB
Python
277 lines
9.8 KiB
Python
#!/usr/bin/env python3
|
||
"""AI agent — supports OpenRouter and Cloudflare Workers AI."""
|
||
|
||
import datetime
|
||
import json
|
||
import logging
|
||
import os
|
||
import time
|
||
from typing import List, Dict
|
||
|
||
import requests
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
SESSIONS_FILE = os.path.expanduser("~/.config/deltabot/ai_sessions.json")
|
||
MAX_HISTORY = 50
|
||
|
||
def _system_prompt() -> str:
|
||
now = datetime.datetime.now().strftime("%d.%m.%Y %H:%M:%S %Z")
|
||
return (
|
||
f"Ты — полезный AI-ассистент в Delta Chat. "
|
||
f"Текущее серверное время: {now}. "
|
||
f"Отвечай на русском, если не просят иное. "
|
||
f"Будь точным, лаконичным и полезным."
|
||
)
|
||
|
||
# --- Models ---
|
||
|
||
OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions"
|
||
|
||
DEFAULT_MODEL = "@cf/qwen/qwen3-30b-a3b-fp8"
|
||
FALLBACK_MODEL = "openrouter/free"
|
||
|
||
MODELS = {
|
||
# Cloudflare Workers AI
|
||
"@cf/qwen/qwen3-30b-a3b-fp8": "Qwen 3 30B (Cloudflare, по умолчанию)",
|
||
"@cf/meta/llama-3.3-70b-instruct-fp8-fast": "Llama 3.3 70B (Cloudflare)",
|
||
"@cf/meta/llama-3.1-8b-instruct-fast": "Llama 3.1 8B (Cloudflare, быстрая)",
|
||
"@cf/meta/llama-4-scout-17b-16e-instruct": "Llama 4 Scout 17B (Cloudflare)",
|
||
"@cf/deepseek-ai/deepseek-r1-distill-qwen-32b": "DeepSeek R1 Distill Qwen 32B (Cloudflare)",
|
||
"@cf/moonshotai/kimi-k2-instruct": "Kimi K2 (Cloudflare)",
|
||
"@cf/aisingapore/gemma-sea-lion-v4-27b-it": "Gemma Sea Lion v4 27B (Cloudflare)",
|
||
# OpenRouter (fallback)
|
||
"openrouter/free": "OpenRouter Free (auto-route)",
|
||
"deepseek/deepseek-v4-flash:free": "DeepSeek V4 Flash (1M ctx, free)",
|
||
"moonshotai/kimi-k2.6:free": "Kimi K2.6 (262K ctx, free)",
|
||
"minimax/minimax-m2.5:free": "MiniMax M2.5 (262K ctx, free)",
|
||
}
|
||
|
||
# --- Session persistence ---
|
||
|
||
|
||
def _load():
|
||
if os.path.exists(SESSIONS_FILE):
|
||
with open(SESSIONS_FILE, "r", encoding="utf-8") as f:
|
||
return json.load(f)
|
||
return {}
|
||
|
||
|
||
def _save(data):
|
||
os.makedirs(os.path.dirname(SESSIONS_FILE), exist_ok=True)
|
||
with open(SESSIONS_FILE, "w", encoding="utf-8") as f:
|
||
json.dump(data, f, ensure_ascii=False, indent=2)
|
||
|
||
|
||
def load_session(chat_id: str) -> Dict:
|
||
sessions = _load()
|
||
if chat_id not in sessions:
|
||
sessions[chat_id] = {
|
||
"model": DEFAULT_MODEL,
|
||
"enabled": False,
|
||
"messages": [],
|
||
"api_key": "",
|
||
}
|
||
_save(sessions)
|
||
return sessions[chat_id]
|
||
|
||
|
||
def save_session(chat_id: str, session: Dict):
|
||
sessions = _load()
|
||
sessions[chat_id] = session
|
||
_save(sessions)
|
||
|
||
|
||
def is_ai_enabled(chat_id: str) -> bool:
|
||
return load_session(chat_id).get("enabled", False)
|
||
|
||
|
||
def get_api_key(chat_id: str) -> str:
|
||
session = load_session(chat_id)
|
||
if session.get("api_key"):
|
||
return session["api_key"]
|
||
return os.environ.get("OPENROUTER_API_KEY", "")
|
||
|
||
|
||
# --- Provider detection ---
|
||
|
||
|
||
def _is_cloudflare(model: str) -> bool:
|
||
return model.startswith("@cf/")
|
||
|
||
|
||
# --- Cloudflare Workers AI ---
|
||
|
||
|
||
def _call_cloudflare(model: str, messages: List[Dict]) -> str:
|
||
account_id = os.environ.get("CLOUDFLARE_ACCOUNT_ID")
|
||
api_token = os.environ.get("CLOUDFLARE_API_TOKEN")
|
||
if not account_id:
|
||
return "❌ CLOUDFLARE_ACCOUNT_ID не установлен в .env"
|
||
if not api_token:
|
||
return "❌ CLOUDFLARE_API_TOKEN не установлен в .env"
|
||
|
||
url = f"https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/run/{model}"
|
||
headers = {
|
||
"Authorization": f"Bearer {api_token}",
|
||
"Content-Type": "application/json",
|
||
}
|
||
body = {"messages": messages, "max_tokens": 8192}
|
||
|
||
timeout = 180 if "deepseek" in model else 60
|
||
|
||
for attempt in range(3):
|
||
try:
|
||
resp = requests.post(url, json=body, headers=headers, timeout=timeout)
|
||
if resp.status_code == 429:
|
||
logger.warning(
|
||
"Cloudflare 429 (attempt %d/3): %s", attempt + 1, resp.text[:300]
|
||
)
|
||
if attempt < 2:
|
||
time.sleep(3)
|
||
continue
|
||
return "⚠️ Лимит Cloudflare AI (429). Попробуй позже."
|
||
if resp.status_code != 200:
|
||
return f"❌ Cloudflare API ошибка: {resp.status_code}"
|
||
|
||
try:
|
||
data = resp.json()
|
||
if not data.get("success"):
|
||
err = data.get("errors", [{}])[0].get("message", resp.text[:200])
|
||
return f"❌ Cloudflare: {err}"
|
||
raw = data["result"]["response"]
|
||
if isinstance(raw, dict):
|
||
raw = raw.get("content", str(raw))
|
||
return raw
|
||
except Exception:
|
||
return f"❌ Не удалось распарсить ответ Cloudflare: {resp.text[:200]}"
|
||
|
||
except requests.Timeout:
|
||
if attempt < 2:
|
||
logger.warning("Cloudflare timeout (attempt %d/3), retrying...", attempt + 1)
|
||
time.sleep(3)
|
||
continue
|
||
return "⚠️ Таймаут Cloudflare AI. Попробуй позже."
|
||
except Exception as e:
|
||
return f"❌ Ошибка Cloudflare: {e}"
|
||
|
||
return "⚠️ Не удалось получить ответ Cloudflare."
|
||
|
||
|
||
# --- OpenRouter ---
|
||
|
||
|
||
def _call_openrouter(model: str, messages: List[Dict], api_key: str) -> str:
|
||
if not api_key:
|
||
return "❌ API ключ не установлен. Используй /apikey [ключ]"
|
||
|
||
for attempt in range(3):
|
||
try:
|
||
resp = requests.post(
|
||
OPENROUTER_URL,
|
||
json={
|
||
"model": model,
|
||
"messages": messages,
|
||
"temperature": 0.7,
|
||
"max_tokens": 2048,
|
||
},
|
||
headers={
|
||
"Authorization": f"Bearer {api_key}",
|
||
"Content-Type": "application/json",
|
||
},
|
||
timeout=60,
|
||
)
|
||
if resp.status_code == 429:
|
||
body = resp.text[:300]
|
||
retry = resp.headers.get("Retry-After", "?")
|
||
limit = resp.headers.get("X-RateLimit-Limit", "?")
|
||
remaining = resp.headers.get("X-RateLimit-Remaining", "?")
|
||
reset = resp.headers.get("X-RateLimit-Reset", "?")
|
||
logger.warning(
|
||
"OpenRouter 429 (attempt %d/3): retry=%s limit=%s remaining=%s reset=%s body=%s",
|
||
attempt + 1, retry, limit, remaining, reset, body,
|
||
)
|
||
if attempt < 2:
|
||
time.sleep(3)
|
||
continue
|
||
return "⚠️ Лимит запросов (429). Попробуй позже или смени модель (/model)."
|
||
if resp.status_code != 200:
|
||
return f"❌ API ошибка: {resp.status_code}"
|
||
break
|
||
except requests.Timeout:
|
||
if attempt < 2:
|
||
logger.warning("OpenRouter timeout (attempt %d/3), retrying...", attempt + 1)
|
||
time.sleep(3)
|
||
continue
|
||
return "⚠️ Таймаут API. Попробуй позже."
|
||
except Exception as e:
|
||
return f"❌ Ошибка: {e}"
|
||
|
||
try:
|
||
data = resp.json()
|
||
return data["choices"][0]["message"]["content"]
|
||
except Exception:
|
||
return "❌ Не удалось распарсить ответ API."
|
||
|
||
|
||
# --- Main ---
|
||
|
||
|
||
def summarize_session(chat_id: str) -> str:
|
||
session = load_session(chat_id)
|
||
messages = session.get("messages", [])
|
||
if not messages:
|
||
return "История пуста — нечего резюмировать."
|
||
model = session.get("model", DEFAULT_MODEL)
|
||
context = messages[-20:]
|
||
dialog = "\n".join(
|
||
f"{'Пользователь' if m['role'] == 'user' else 'AI'}: {m['content']}"
|
||
for m in context
|
||
)
|
||
summary_prompt = (
|
||
"Кратко (3–7 предложений) резюмируй следующий диалог на русском языке. "
|
||
"Выдели главные темы и итоги:\n\n" + dialog
|
||
)
|
||
api_messages = [
|
||
{"role": "system", "content": _system_prompt()},
|
||
{"role": "user", "content": summary_prompt},
|
||
]
|
||
if _is_cloudflare(model):
|
||
reply = _call_cloudflare(model, api_messages)
|
||
if (reply.startswith("❌") or reply.startswith("⚠️")) and get_api_key(chat_id):
|
||
logger.warning("Cloudflare summarization failed, fallback to OpenRouter: %s", reply)
|
||
reply = _call_openrouter(FALLBACK_MODEL, api_messages, get_api_key(chat_id))
|
||
return reply
|
||
return _call_openrouter(model, api_messages, get_api_key(chat_id))
|
||
|
||
|
||
def process_message(chat_id: str, user_message: str) -> str:
|
||
session = load_session(chat_id)
|
||
model = session.get("model", DEFAULT_MODEL)
|
||
|
||
messages = session.get("messages", [])
|
||
messages.append({"role": "user", "content": user_message})
|
||
|
||
if len(messages) > MAX_HISTORY:
|
||
messages = messages[-MAX_HISTORY:]
|
||
|
||
api_messages = [{"role": "system", "content": _system_prompt()}] + messages
|
||
|
||
if _is_cloudflare(model):
|
||
reply = _call_cloudflare(model, api_messages)
|
||
if (reply.startswith("❌") or reply.startswith("⚠️")) and get_api_key(chat_id):
|
||
cf_error = reply
|
||
logger.warning("Cloudflare failed, fallback to OpenRouter: %s", cf_error)
|
||
reply = _call_openrouter(FALLBACK_MODEL, api_messages, get_api_key(chat_id))
|
||
else:
|
||
api_key = get_api_key(chat_id)
|
||
reply = _call_openrouter(model, api_messages, api_key)
|
||
|
||
if reply.startswith("❌") or reply.startswith("⚠️"):
|
||
return reply
|
||
|
||
messages.append({"role": "assistant", "content": reply})
|
||
session["messages"] = messages
|
||
save_session(chat_id, session)
|
||
|
||
return reply
|