feat(transcribe,texts): voice notes and round videos to the bot come with a transcript, beavergram finds them by sha256 and telegram premium transcribes
This commit is contained in:
1 parent
6db6c7b6ec
commit
d7eb17ba94
6 files changed
+346
-1
No files matched your search
@@ -0,0 +1,179 @@
|
||||
import hashlib
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from aiohttp import web
|
||||
from aiohttp.test_utils import TestServer
|
||||
from beaver_gateway.transcribe import Audio
|
||||
|
||||
from beaver_agent.transcribe import BeavergramTranscriber
|
||||
|
||||
OWNER = 777
|
||||
VOICE = b"OggS voice"
|
||||
DIGEST = hashlib.sha256(VOICE).hexdigest()
|
||||
|
||||
|
||||
class Beavergram:
|
||||
"""Поддельный beavergram: аккаунты, поиск по sha256 и расшифровка."""
|
||||
|
||||
def __init__(self, *, empty_lookups=0, kind="voice", text=None, fail=None):
|
||||
self.empty_lookups = empty_lookups
|
||||
self.kind = kind
|
||||
self.text = text
|
||||
self.fail = fail
|
||||
self.lookups = []
|
||||
self.transcribes = []
|
||||
self.auth = set()
|
||||
|
||||
def app(self):
|
||||
app = web.Application(middlewares=[self.gate])
|
||||
app.router.add_get("/api/accounts", self.accounts)
|
||||
app.router.add_get("/api/media/lookup", self.lookup)
|
||||
app.router.add_post("/api/media/transcribe", self.transcribe)
|
||||
return app
|
||||
|
||||
@web.middleware
|
||||
async def gate(self, request, handler):
|
||||
self.auth.add(request.headers.get("Authorization"))
|
||||
if self.fail:
|
||||
return web.json_response({"detail": "nope"}, status=self.fail)
|
||||
return await handler(request)
|
||||
|
||||
async def accounts(self, request):
|
||||
return web.json_response(
|
||||
[
|
||||
{"account_id": 1, "tg_user_id": 5, "is_bot": True},
|
||||
{"account_id": 2, "tg_user_id": OWNER, "is_bot": False},
|
||||
]
|
||||
)
|
||||
|
||||
async def lookup(self, request):
|
||||
self.lookups.append(dict(request.query))
|
||||
if len(self.lookups) <= self.empty_lookups:
|
||||
return web.json_response([])
|
||||
return web.json_response(
|
||||
[
|
||||
{
|
||||
"account_id": 2,
|
||||
"chat_id": 42,
|
||||
"message_id": 9,
|
||||
"kind": self.kind,
|
||||
"extracted_text": self.text,
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
async def transcribe(self, request):
|
||||
self.transcribes.append(await request.json())
|
||||
return web.json_response({"status": "done", "text": "привет из премиума"})
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def voice():
|
||||
path = Path(tempfile.mkdtemp(prefix="beaver-bg-")) / "1-vu1.ogg"
|
||||
path.write_bytes(VOICE)
|
||||
return path
|
||||
|
||||
|
||||
async def serve(fake):
|
||||
server = TestServer(fake.app())
|
||||
await server.start_server()
|
||||
return server
|
||||
|
||||
|
||||
def transcriber(server):
|
||||
return BeavergramTranscriber(
|
||||
origin=str(server.make_url("")).rstrip("/"),
|
||||
token="t0k",
|
||||
tg_user_id=OWNER,
|
||||
find_for=2.0,
|
||||
find_every=0.01,
|
||||
)
|
||||
|
||||
|
||||
async def test_finds_the_voice_after_a_retry_and_asks_for_a_transcript(voice):
|
||||
fake = Beavergram(empty_lookups=2)
|
||||
server = await serve(fake)
|
||||
try:
|
||||
stt = transcriber(server)
|
||||
assert await stt(Audio(voice, "audio/ogg", 3)) == "привет из премиума"
|
||||
finally:
|
||||
await server.close()
|
||||
assert len(fake.lookups) == 3
|
||||
assert fake.lookups[-1] == {"sha256": DIGEST, "account_id": "2"}
|
||||
assert fake.transcribes == [
|
||||
{"account_id": 2, "chat_id": 42, "message_id": 9, "wait": 60}
|
||||
]
|
||||
assert fake.auth == {"Bearer t0k"}
|
||||
|
||||
|
||||
async def test_an_existing_transcript_comes_back_without_a_transcribe(voice):
|
||||
fake = Beavergram(kind="video_note", text="уже расшифровано")
|
||||
server = await serve(fake)
|
||||
try:
|
||||
assert (
|
||||
await transcriber(server)(Audio(voice, "video/mp4")) == "уже расшифровано"
|
||||
)
|
||||
finally:
|
||||
await server.close()
|
||||
assert fake.transcribes == []
|
||||
|
||||
|
||||
async def test_music_and_other_files_are_not_touched(voice):
|
||||
fake = Beavergram()
|
||||
server = await serve(fake)
|
||||
try:
|
||||
stt = transcriber(server)
|
||||
for media_type in ("audio/mpeg", "audio/mp4", None):
|
||||
assert await stt(Audio(voice, media_type)) is None
|
||||
assert fake.lookups == []
|
||||
fake.kind = "audio"
|
||||
assert await stt(Audio(voice, "audio/ogg")) is None
|
||||
finally:
|
||||
await server.close()
|
||||
assert len(fake.lookups) == 1 and fake.transcribes == []
|
||||
|
||||
|
||||
async def test_nothing_found_in_time_is_none(voice):
|
||||
fake = Beavergram(empty_lookups=1000)
|
||||
server = await serve(fake)
|
||||
try:
|
||||
stt = transcriber(server)
|
||||
stt.find_for = 0.05
|
||||
assert await stt(Audio(voice, "audio/ogg")) is None
|
||||
finally:
|
||||
await server.close()
|
||||
assert len(fake.lookups) >= 2 and fake.transcribes == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status", [401, 500])
|
||||
async def test_http_errors_are_none_not_raised(voice, status):
|
||||
fake = Beavergram(fail=status)
|
||||
server = await serve(fake)
|
||||
try:
|
||||
assert await transcriber(server)(Audio(voice, "audio/ogg")) is None
|
||||
finally:
|
||||
await server.close()
|
||||
assert fake.auth == {"Bearer t0k"}
|
||||
|
||||
|
||||
def test_from_env_needs_the_token_and_the_owner():
|
||||
url = "https://tg.example.com/mcp?token=" + "a" * 64
|
||||
stt = BeavergramTranscriber.from_env(
|
||||
{"BEAVERGRAM_MCP": url, "TELEGRAM_USER_ID": "777"}
|
||||
)
|
||||
assert stt is not None
|
||||
assert (stt.origin, stt.token, stt.tg_user_id) == (
|
||||
"https://tg.example.com",
|
||||
"a" * 64,
|
||||
777,
|
||||
)
|
||||
assert BeavergramTranscriber.from_env({"TELEGRAM_USER_ID": "777"}) is None
|
||||
assert (
|
||||
BeavergramTranscriber.from_env(
|
||||
{"BEAVERGRAM_MCP": "https://tg.example.com/mcp", "TELEGRAM_USER_ID": "777"}
|
||||
)
|
||||
is None
|
||||
)
|
||||
assert BeavergramTranscriber.from_env({"BEAVERGRAM_MCP": url}) is None
|
||||
Reference in new issue
Block a user