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:
Галингер Р.С.
+62
-44
@@ -11,83 +11,101 @@ from config import PROXY_ENABLED, PROXY_URL
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_model = None
|
||||
_model: WhisperModel | None = None
|
||||
_model_lock = asyncio.Lock()
|
||||
|
||||
BASE_DIR = Path(__file__).resolve().parent.parent.parent
|
||||
MODELS_DIR = BASE_DIR / "models" / "faster-whisper"
|
||||
VOICE_DIR = BASE_DIR / "bot" / "data" / "voice"
|
||||
|
||||
|
||||
def _setup_proxy():
|
||||
def _setup_proxy() -> None:
|
||||
if PROXY_ENABLED and PROXY_URL:
|
||||
os.environ["ALL_PROXY"] = PROXY_URL
|
||||
logger.info("Proxy set for model download: %s", PROXY_URL[:30] + "...")
|
||||
logger.info("Proxy set for model download: %s...", PROXY_URL[:30])
|
||||
|
||||
|
||||
def _get_model() -> WhisperModel:
|
||||
def _load_model() -> WhisperModel:
|
||||
"""Synchronous model load/download. Must run in a thread."""
|
||||
MODELS_DIR.mkdir(parents=True, exist_ok=True)
|
||||
VOICE_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
model_path = MODELS_DIR / "base"
|
||||
model_bin = model_path / "model.bin"
|
||||
|
||||
if not model_path.is_dir() or not model_bin.is_file():
|
||||
_setup_proxy()
|
||||
logger.info("Model not found locally, downloading to %s...", MODELS_DIR)
|
||||
return WhisperModel(
|
||||
"base",
|
||||
device="cpu",
|
||||
cpu_threads=4,
|
||||
compute_type="int8",
|
||||
download_root=str(MODELS_DIR),
|
||||
)
|
||||
|
||||
logger.info("Loading model from %s...", model_path)
|
||||
return WhisperModel(
|
||||
str(model_path),
|
||||
device="cpu",
|
||||
cpu_threads=4,
|
||||
compute_type="int8",
|
||||
)
|
||||
|
||||
|
||||
async def _get_model() -> WhisperModel:
|
||||
"""Thread-safe lazy initializer for the Whisper model."""
|
||||
global _model
|
||||
if _model is None:
|
||||
MODELS_DIR.mkdir(parents=True, exist_ok=True)
|
||||
VOICE_DIR.mkdir(parents=True, exist_ok=True)
|
||||
model_path = str(MODELS_DIR / "base")
|
||||
if not os.path.isdir(model_path) or not os.path.isfile(os.path.join(model_path, "model.bin")):
|
||||
_setup_proxy()
|
||||
logger.info("Model not found locally, downloading to %s...", MODELS_DIR)
|
||||
_model = WhisperModel(
|
||||
"base",
|
||||
device="cpu",
|
||||
cpu_threads=4,
|
||||
compute_type="int8",
|
||||
download_root=str(MODELS_DIR),
|
||||
)
|
||||
else:
|
||||
logger.info("Loading model from %s...", model_path)
|
||||
_model = WhisperModel(
|
||||
model_path,
|
||||
device="cpu",
|
||||
cpu_threads=4,
|
||||
compute_type="int8",
|
||||
)
|
||||
logger.info("Whisper model loaded")
|
||||
async with _model_lock:
|
||||
if _model is None:
|
||||
loop = asyncio.get_running_loop()
|
||||
_model = await loop.run_in_executor(None, _load_model)
|
||||
logger.info("Whisper model loaded")
|
||||
return _model
|
||||
|
||||
|
||||
async def download_voice(bot: Bot, file_id: str) -> str:
|
||||
VOICE_DIR.mkdir(parents=True, exist_ok=True)
|
||||
file = await bot.get_file(file_id)
|
||||
path = str(VOICE_DIR / f"{file_id}.ogg")
|
||||
safe_id = re.sub(r"[^A-Za-z0-9_-]", "_", file_id)
|
||||
path = str(VOICE_DIR / f"{safe_id}.ogg")
|
||||
await bot.download_file(file.file_path, destination=path)
|
||||
logger.info("Voice downloaded: %s -> %s", file_id, path)
|
||||
logger.info("Voice downloaded: %s", path)
|
||||
return path
|
||||
|
||||
|
||||
async def convert_to_wav(input_path: str) -> str:
|
||||
output_path = input_path.replace(".ogg", ".wav")
|
||||
output_path = str(Path(input_path).with_suffix(".wav"))
|
||||
proc = await asyncio.create_subprocess_exec(
|
||||
"ffmpeg", "-i", input_path, "-ar", "16000", "-ac", "1",
|
||||
"-c:a", "pcm_s16le", output_path, "-y",
|
||||
"ffmpeg",
|
||||
"-i", input_path,
|
||||
"-ar", "16000",
|
||||
"-ac", "1",
|
||||
"-c:a", "pcm_s16le",
|
||||
output_path,
|
||||
"-y",
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
)
|
||||
_, stderr = await proc.communicate()
|
||||
if proc.returncode != 0:
|
||||
logger.error(f"ffmpeg error: {stderr.decode(errors='replace')}")
|
||||
logger.error("ffmpeg error: %s", stderr.decode(errors="replace"))
|
||||
raise RuntimeError("ffmpeg conversion failed")
|
||||
logger.info(f"Converted to WAV: {output_path}")
|
||||
logger.info("Converted to WAV: %s", output_path)
|
||||
return output_path
|
||||
|
||||
|
||||
async def transcribe_audio(file_path: str) -> str:
|
||||
model = _get_model()
|
||||
model = await _get_model()
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
def _transcribe():
|
||||
segments, info = model.transcribe(file_path, language="ru", beam_size=5)
|
||||
text = " ".join(seg.text for seg in segments)
|
||||
return text.strip()
|
||||
def _transcribe() -> str:
|
||||
segments, _info = model.transcribe(file_path, language="ru", beam_size=5)
|
||||
return " ".join(seg.text for seg in segments).strip()
|
||||
|
||||
text = await loop.run_in_executor(None, _transcribe)
|
||||
logger.info(f"Transcription result ({len(text)} chars): {text[:100]}...")
|
||||
logger.info("Transcription result (%d chars): %s...", len(text), text[:100])
|
||||
return text
|
||||
|
||||
|
||||
@@ -102,11 +120,11 @@ def normalize_text(text: str) -> str:
|
||||
return text
|
||||
|
||||
|
||||
async def cleanup_files(*paths: str):
|
||||
async def cleanup_files(*paths: str | None) -> None:
|
||||
for path in paths:
|
||||
if path and os.path.exists(path):
|
||||
try:
|
||||
os.unlink(path)
|
||||
logger.debug(f"Cleaned up: {path}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Cleanup failed for {path}: {e}")
|
||||
await asyncio.to_thread(os.unlink, path)
|
||||
logger.debug("Cleaned up: %s", path)
|
||||
except Exception as exc:
|
||||
logger.warning("Cleanup failed for %s: %s", path, exc)
|
||||
|
||||
Reference in New Issue
Block a user