Update 4 files

- /rag/rag_server.py
- /bots/rag_client.py
- /bots/workers/indexing_worker.py
- /bots/handlers/file_handler.py
This commit is contained in:
Markov Andrey
2026-06-30 12:43:27 +00:00
parent c82cb17a49
commit 8a2d5ba2ff
4 changed files with 207 additions and 30 deletions

View File

@@ -7,6 +7,10 @@
- Если файл пришёл в групповой чат (room_jid не None), он автоматически индексируется
в базу знаний комнаты, без необходимости команды !learn.
- Для личных чатов поведение зависит от режима обучения (learning_mode) и флага --global.
ДОБАВЛЕНО:
- Удалён прямой вызов GigaChat для изображений (теперь это делает RAG-сервер)
- Удалён прямой вызов транскрибации (теперь это делает RAG-сервер)
"""
import os

View File

@@ -10,8 +10,13 @@ HTTP-клиент для взаимодействия с RAG-сервером.
ИСТОРИЯ НЕ ПЕРЕДАЁТСЯ сервер получает её из БД.
Улучшена обработка сетевых ошибок с повторными попытками.
ДОБАВЛЕНО:
- Метод vision() для распознавания текста на изображениях
- Метод transcribe() для транскрибации аудио
"""
import os
import logging
from typing import Optional, Dict, List, Any
import aiohttp
@@ -281,4 +286,68 @@ class RAGClient:
return result
except Exception as e:
logger.error(f"Ошибка в RAGClient.index_document: {e}", exc_info=True)
return {"doc_id": None, "chunk_count": 0, "error": str(e)}
return {"doc_id": None, "chunk_count": 0, "error": str(e)}
# ==========================================================
# НОВЫЕ МЕТОДЫ: OCR и транскрибация
# ==========================================================
async def vision(self, file_path: str) -> Dict[str, Any]:
"""
Отправляет изображение на RAG-сервер для распознавания текста (OCR).
Аргументы:
file_path: путь к файлу изображения
Возвращает:
Dict с ключами: text (str), status (str), error (опционально)
"""
if not os.path.exists(file_path):
return {"text": "", "status": "file_not_found", "error": "Файл не найден"}
session = await self._get_session()
url = f"{self.server_url}/rag/vision"
try:
with open(file_path, 'rb') as f:
files = {'file': (os.path.basename(file_path), f, 'image/jpeg')}
async with session.post(url, data=files) as resp:
if resp.status == 200:
return await resp.json()
else:
error_text = await resp.text()
logger.error(f"OCR ошибка: {resp.status} - {error_text[:200]}")
return {"text": "", "status": "error", "error": error_text}
except Exception as e:
logger.error(f"Ошибка вызова OCR: {e}", exc_info=True)
return {"text": "", "status": "error", "error": str(e)}
async def transcribe(self, file_path: str) -> Dict[str, Any]:
"""
Отправляет аудиофайл на RAG-сервер для транскрибации (SaluteSpeech).
Аргументы:
file_path: путь к аудиофайлу
Возвращает:
Dict с ключами: text (str), status (str), error (опционально)
"""
if not os.path.exists(file_path):
return {"text": "", "status": "file_not_found", "error": "Файл не найден"}
session = await self._get_session()
url = f"{self.server_url}/rag/transcribe"
try:
with open(file_path, 'rb') as f:
files = {'file': (os.path.basename(file_path), f, 'audio/ogg')}
async with session.post(url, data=files) as resp:
if resp.status == 200:
return await resp.json()
else:
error_text = await resp.text()
logger.error(f"Транскрибация ошибка: {resp.status} - {error_text[:200]}")
return {"text": "", "status": "error", "error": error_text}
except Exception as e:
logger.error(f"Ошибка вызова транскрибации: {e}", exc_info=True)
return {"text": "", "status": "error", "error": str(e)}

View File

@@ -6,6 +6,10 @@
- Для файлов: извлечение текста, обработка изображений/аудио, индексация через HTTP-запрос.
Все операции с БД и Qdrant выполняются на сервере. Воркер только отправляет данные.
ДОБАВЛЕНО:
- Обработка изображений через RAG-сервер (OCR)
- Обработка аудио через RAG-сервер (транскрибация)
"""
import asyncio
@@ -149,7 +153,7 @@ class IndexingWorker:
"""
Обрабатывает задачу с файлом:
1. Извлекает текст (синхронно в потоке).
2. Обрабатывает изображения/аудио (Vision/распознавание) если есть.
2. Обрабатывает изображения/аудио (Vision/распознавание) через RAG-сервер.
3. Сохраняет постоянную копию в data/ (локально, для истории и отладки).
4. Отправляет текст и метаданные на RAG-сервер для индексации.
5. Удаляет временный файл.
@@ -210,40 +214,40 @@ class IndexingWorker:
key = (user_jid, room_jid) if room_jid is not None else user_jid
self.bot._last_file_path[key] = str(dest_path)
# --- 3. Специальная обработка: изображения (Vision API) ---
# Если это изображение и включена опция vision, отправляем на анализ.
# Это можно делать отдельно, не блокируя индексацию.
# --- 3. Специальная обработка: изображения (OCR через RAG-сервер) ---
if ext in ('.jpg', '.jpeg', '.png', '.bmp', '.tiff') and self.bot.config.vision:
file_id = await self.bot.giga.upload_file(file_path) # giga всё ещё локально? нужно будет перенести и это.
# Это временно, пока мы не вынесли всё на сервер. Пока оставим как есть.
if file_id:
answer = await self.bot.giga.chat(
history=[],
query="Проанализируй это изображение. Если есть текст прочитай его.",
system_prompt=None,
file_id=file_id,
temperature=self.bot.config.ai_temperature
result = await self.bot.rag_client.vision(file_path)
if result.get('text'):
await self._send_message(
user_jid,
f"🖼️ **Текст на изображении:**\n{result['text']}",
room_jid
)
await self._send_message(user_jid, answer, room_jid)
else:
await self._send_message(user_jid, "⚠️ Не удалось передать изображение в GigaChat.", room_jid)
# После обработки изображения мы не индексируем его как документ (обычно OCR делает это),
# но можно и проиндексировать текст. Здесь мы просто возвращаемся.
error = result.get('error', 'неизвестная ошибка')
await self._send_message(
user_jid,
f"⚠️ Не удалось распознать текст: {error}",
room_jid
)
return
# --- 4. Специальная обработка: аудио (голос → текст) ---
# Аналогично, используем локальный SaluteSpeech, пока не вынесли.
# --- 4. Специальная обработка: аудио (транскрибация через RAG-сервер) ---
if ext in ('.ogg', '.wav', '.mp3', '.amr', '.m4a') and self.bot.config.voice_recognition:
first_text = files_list[0][1] if files_list else ""
if "--- РАСПОЗНАННЫЙ ГОЛОС ---" in first_text:
voice_text = first_text.split("--- РАСПОЗНАННЫЙ ГОЛОС ---")[1].strip()
await self._send_message(user_jid, f"🎤 Распознано: {voice_text}", room_jid)
# Можно также ответить на голос, но это отдельный сценарий.
# Мы не индексируем аудио как документ (разве что текст распознавания).
# Если хотим индексировать текст голоса, нужно передать его как документ.
# Например, можно создать виртуальный документ с текстом.
result = await self.bot.rag_client.transcribe(file_path)
if result.get('text'):
await self._send_message(
user_jid,
f"🎤 **Распознано:**\n{result['text']}",
room_jid
)
else:
await self._send_message(user_jid, f"⚠️ {first_text}", room_jid)
error = result.get('error', 'неизвестная ошибка')
await self._send_message(
user_jid,
f"⚠️ Не удалось распознать речь: {error}",
room_jid
)
return
# --- 5. Индексация документов через RAG-сервер ---

View File

@@ -9,18 +9,25 @@ RAG-сервер отдельный HTTP-сервис для обработ
ИСПОЛЬЗУЕТ: FastAPI для создания REST API и Uvicorn для запуска.
ИСТОРИЯ ХРАНИТСЯ НА СЕРВЕРЕ (в PostgreSQL).
ДОБАВЛЕНО:
- Эндпоинт /rag/vision для распознавания текста на изображениях (OCR)
- Эндпоинт /rag/transcribe для транскрибации аудио (SaluteSpeech)
"""
import logging
import os
import sys
import tempfile
import shutil
import asyncio
from pathlib import Path
from typing import Optional, Dict, List, Any
# Добавляем путь к core, если запускаем из корня проекта
sys.path.insert(0, str(Path(__file__).parent.parent))
from fastapi import FastAPI, HTTPException, Depends
from fastapi import FastAPI, HTTPException, Depends, UploadFile, File as FastAPIFile
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
import uvicorn
@@ -94,6 +101,18 @@ class IndexResponse(BaseModel):
error: Optional[str] = Field(None, description="Сообщение об ошибке, если есть")
class VisionResponse(BaseModel):
"""Модель ответа на распознавание изображения."""
text: str = Field(..., description="Распознанный текст")
status: str = Field(..., description="Статус: ok, no_text_found, error")
class TranscribeResponse(BaseModel):
"""Модель ответа на транскрибацию аудио."""
text: str = Field(..., description="Распознанный текст")
status: str = Field(..., description="Статус: ok, no_speech_detected, error")
# ============================================================
# ГЛОБАЛЬНЫЙ ОРКЕСТРАТОР (инициализируется один раз при старте)
# ============================================================
@@ -311,6 +330,87 @@ async def health_check():
return {"status": "healthy"}
# ============================================================
# НОВЫЕ ЭНДПОИНТЫ: OCR и транскрибация
# ============================================================
@app.post("/rag/vision", response_model=VisionResponse)
async def vision_endpoint(
file: UploadFile = FastAPIFile(...),
user_jid: str = None,
room_jid: str = None
):
"""
Распознаёт текст на изображении через OCR (Tesseract).
Возвращает JSON: {"text": str, "status": "ok" | "no_text_found" | "error"}
"""
orchestrator = get_orchestrator()
# Проверяем расширение
ext = os.path.splitext(file.filename)[1].lower()
if ext not in ('.jpg', '.jpeg', '.png', '.bmp', '.tiff'):
raise HTTPException(status_code=400, detail="Поддерживаются только изображения (jpg, png, bmp, tiff)")
# Сохраняем во временный файл
with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as tmp:
tmp_path = tmp.name
shutil.copyfileobj(file.file, tmp)
try:
# Извлекаем текст через существующий FileService (синхронный метод)
loop = asyncio.get_running_loop()
text = await loop.run_in_executor(
None,
orchestrator.files.extract_text_from_image,
tmp_path
)
if not text:
return VisionResponse(text="", status="no_text_found")
return VisionResponse(text=text, status="ok")
except Exception as e:
logger.error(f"Ошибка OCR: {e}", exc_info=True)
raise HTTPException(status_code=500, detail=str(e))
finally:
if os.path.exists(tmp_path):
os.unlink(tmp_path)
@app.post("/rag/transcribe", response_model=TranscribeResponse)
async def transcribe_endpoint(
file: UploadFile = FastAPIFile(...),
user_jid: str = None,
room_jid: str = None
):
"""
Транскрибирует аудиофайл через SaluteSpeech.
Возвращает JSON: {"text": str, "status": "ok" | "no_speech_detected" | "error"}
"""
orchestrator = get_orchestrator()
# Проверяем расширение
ext = os.path.splitext(file.filename)[1].lower()
if ext not in ('.ogg', '.wav', '.mp3', '.amr', '.m4a'):
raise HTTPException(status_code=400, detail="Поддерживаются только аудио (ogg, wav, mp3, amr, m4a)")
# Сохраняем во временный файл
with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as tmp:
tmp_path = tmp.name
shutil.copyfileobj(file.file, tmp)
try:
# Транскрибируем через существующий FileService (асинхронный метод)
text = await orchestrator.files.transcribe_audio(tmp_path)
if not text:
return TranscribeResponse(text="", status="no_speech_detected")
return TranscribeResponse(text=text, status="ok")
except Exception as e:
logger.error(f"Ошибка транскрибации: {e}", exc_info=True)
raise HTTPException(status_code=500, detail=str(e))
finally:
if os.path.exists(tmp_path):
os.unlink(tmp_path)
# ============================================================
# ТОЧКА ВХОДА
# ============================================================