refactor: review fixes and code improvements
- Add shared proxy module (bot/utils/proxy.py) to eliminate duplicate SOCKS5 parsing - Fix AI client: escape HTML before Markdown→HTML conversion, unify timeouts, make health check optional and disabled by default, handle 429 retries - Fix ai.py: correct forwarded message handling (aiogram 3.x forward_origin), pass relevant summaries via extra_system_content - Fix dialogue.py: only respond to Astra's messages, use system context instead of prompt injection, answer on the limit message before phase transition - Fix voice.py: load Whisper model in thread pool, safe WAV path generation - Improve database.py: composite indexes, Boolean is_active, upsert file_id cache, add context cleanup helper - Update weather.py and yadisk_download.py to use shared proxy connector - Update yadisk.py: validate URL before cache clear, add download size limit, wrap sync file ops in to_thread - Reuse S3 client via lru_cache - Update setup_commands with /aiclear and /aiuser - Update README, Dockerfile (Python 3.11), docker-compose (mount models) - Pin dependency versions, remove unused httpx[socks] - Add basic pytest tests for layout converter and voice normalization
This commit is contained in:
Галингер Р.С.
@@ -5,80 +5,64 @@ from pathlib import Path
|
||||
from urllib.parse import urlencode, unquote
|
||||
|
||||
import aiohttp
|
||||
from aiohttp_socks import ProxyConnector, ProxyType
|
||||
|
||||
from config import PROXY_URL, PROXY_ENABLED
|
||||
from bot.utils.proxy import get_proxy_connector
|
||||
|
||||
DOWNLOAD_DIR = Path("bot/data/downloads")
|
||||
DOWNLOAD_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def _get_connector():
|
||||
if PROXY_ENABLED and PROXY_URL:
|
||||
parsed = PROXY_URL.replace("socks5://", "").replace("socks5h://", "")
|
||||
if "@" in parsed:
|
||||
auth, host_port = parsed.split("@", 1)
|
||||
username, password = auth.split(":", 1)
|
||||
else:
|
||||
username = None
|
||||
password = None
|
||||
host_port = parsed
|
||||
|
||||
host, port = host_port.rsplit(":", 1)
|
||||
port = int(port)
|
||||
|
||||
return ProxyConnector(
|
||||
proxy_type=ProxyType.SOCKS5,
|
||||
host=host,
|
||||
port=port,
|
||||
username=username,
|
||||
password=password,
|
||||
)
|
||||
return None
|
||||
|
||||
def sanitize_filename(filename: str) -> str:
|
||||
filename = os.path.basename(filename).strip()
|
||||
filename = re.sub(r"[\\/:*?\"<>|]+", "_", filename)
|
||||
return filename or "downloaded_file"
|
||||
|
||||
|
||||
async def download_yandex_file(public_url: str, progress_callback=None) -> str:
|
||||
base_url = "https://cloud-api.yandex.net/v1/disk/public/resources/download?"
|
||||
final_url = base_url + urlencode({"public_key": public_url})
|
||||
|
||||
timeout = aiohttp.ClientTimeout(total=None, sock_connect=30, sock_read=300)
|
||||
connector = _get_connector()
|
||||
connector = get_proxy_connector()
|
||||
|
||||
async with aiohttp.ClientSession(timeout=timeout, connector=connector) as session:
|
||||
async with session.get(final_url) as response:
|
||||
response.raise_for_status()
|
||||
payload = await response.json()
|
||||
download_url = payload["href"]
|
||||
try:
|
||||
async with aiohttp.ClientSession(timeout=timeout, connector=connector) as session:
|
||||
async with session.get(final_url) as response:
|
||||
response.raise_for_status()
|
||||
payload = await response.json()
|
||||
download_url = payload["href"]
|
||||
|
||||
async with session.get(download_url) as response:
|
||||
response.raise_for_status()
|
||||
async with session.get(download_url) as response:
|
||||
response.raise_for_status()
|
||||
|
||||
content_disposition = response.headers.get("Content-Disposition", "")
|
||||
filename = download_url.split("/")[-1]
|
||||
content_disposition = response.headers.get("Content-Disposition", "")
|
||||
filename = download_url.split("/")[-1]
|
||||
|
||||
if "filename*" in content_disposition:
|
||||
try:
|
||||
encoded = content_disposition.split("filename*=")[1].strip()
|
||||
parts = encoded.split("''", 1)
|
||||
if len(parts) == 2:
|
||||
filename = unquote(parts[1], encoding=parts[0] or "utf-8")
|
||||
except Exception:
|
||||
pass
|
||||
if "filename*" in content_disposition:
|
||||
try:
|
||||
encoded = content_disposition.split("filename*=")[1].strip()
|
||||
parts = encoded.split("''", 1)
|
||||
if len(parts) == 2:
|
||||
filename = unquote(parts[1], encoding=parts[0] or "utf-8")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
filename = sanitize_filename(filename)
|
||||
suffix = Path(filename).suffix or ".bin"
|
||||
filename = sanitize_filename(filename)
|
||||
suffix = Path(filename).suffix or ".bin"
|
||||
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix, dir=str(DOWNLOAD_DIR)) as tmp_file:
|
||||
total_size = int(response.headers.get("Content-Length", 0))
|
||||
downloaded_size = 0
|
||||
with tempfile.NamedTemporaryFile(
|
||||
delete=False, suffix=suffix, dir=str(DOWNLOAD_DIR)
|
||||
) as tmp_file:
|
||||
total_size = int(response.headers.get("Content-Length", 0))
|
||||
downloaded_size = 0
|
||||
|
||||
async for chunk in response.content.iter_chunked(1024 * 64):
|
||||
tmp_file.write(chunk)
|
||||
downloaded_size += len(chunk)
|
||||
if progress_callback:
|
||||
await progress_callback(downloaded_size, total_size)
|
||||
async for chunk in response.content.iter_chunked(1024 * 64):
|
||||
tmp_file.write(chunk)
|
||||
downloaded_size += len(chunk)
|
||||
if progress_callback:
|
||||
await progress_callback(downloaded_size, total_size)
|
||||
|
||||
return tmp_file.name
|
||||
return tmp_file.name
|
||||
finally:
|
||||
if connector:
|
||||
await connector.close()
|
||||
|
||||
Reference in New Issue
Block a user