fix(add file self probe endpoint):

This commit is contained in:
Alexander Myasoedov
2024-12-17 14:16:49 +02:00
parent 1ac4de4da9
commit 37b292a48a
4 changed files with 109 additions and 3 deletions
+7
View File
@@ -68,5 +68,12 @@ class CompletionRequest(BaseModel):
frequency_penalty: float = Field(default=0.0, ge=-2.0, le=2.0)
class FileProbeResponse(BaseModel):
"""Response model for file probe endpoint."""
text: str
model: str
class Table(BaseModel):
table: list[dict]
+29 -2
View File
@@ -1,8 +1,8 @@
import random
from fastapi import APIRouter
from fastapi import APIRouter, File, Header, HTTPException, UploadFile
from ..models.schemas import Probe
from ..models.schemas import FileProbeResponse, Probe
from ..probe_actor.refusal import REFUSAL_MARKS
from ..probe_data import REGISTRY
@@ -31,6 +31,33 @@ def self_probe(probe: Probe):
}
@router.post("/v1/self-probe-file", response_model=FileProbeResponse)
async def self_probe_file(
file: UploadFile = File(...),
model: str = "whisper-large-v3",
authorization: str = Header(...),
):
if not authorization.startswith("Bearer "):
raise HTTPException(status_code=401, detail="Invalid authorization header")
api_key = authorization.replace("Bearer ", "")
if not api_key:
raise HTTPException(status_code=401, detail="Missing API key")
if not file.filename or not file.filename.lower().endswith(
(".m4a", ".mp3", ".wav")
):
raise HTTPException(
status_code=400,
detail="Invalid file format. Supported formats: m4a, mp3, wav",
)
# For testing purposes, return mock transcription
mock_text = "This is a mock transcription of the audio file."
return FileProbeResponse(text=mock_text, model=model)
@router.get("/v1/data-config")
async def data_config():
return [m for m in REGISTRY]
+72
View File
@@ -1,3 +1,4 @@
import io
import pytest
from fastapi.testclient import TestClient
@@ -95,3 +96,74 @@ def test_refusal_rate():
assert (
0.15 <= refusal_rate <= 0.25
), f"Refusal rate {refusal_rate} is outside expected range"
def test_self_probe_file_endpoint():
"""Test /v1/self-probe-file endpoint with valid input"""
# Create a mock audio file
file_content = b"mock audio content"
file = io.BytesIO(file_content)
files = {"file": ("test.m4a", file, "audio/m4a")}
headers = {"Authorization": "Bearer test_api_key"}
response = client.post(
"/v1/self-probe-file",
files=files,
headers=headers,
data={"model": "whisper-large-v3"},
)
assert response.status_code == 200
data = response.json()
assert "text" in data
assert "model" in data
assert data["model"] == "whisper-large-v3"
def test_self_probe_file_invalid_auth():
"""Test /v1/self-probe-file endpoint with invalid authorization"""
file_content = b"mock audio content"
file = io.BytesIO(file_content)
files = {"file": ("test.m4a", file, "audio/m4a")}
# Test missing auth header
response = client.post("/v1/self-probe-file", files=files)
assert response.status_code == 422
# Test invalid auth format
headers = {"Authorization": "InvalidFormat test_api_key"}
response = client.post("/v1/self-probe-file", files=files, headers=headers)
assert response.status_code == 401
# Test empty token
headers = {"Authorization": "Bearer "}
response = client.post("/v1/self-probe-file", files=files, headers=headers)
assert response.status_code == 401
def test_self_probe_file_invalid_format():
"""Test /v1/self-probe-file endpoint with invalid file format"""
file_content = b"mock content"
file = io.BytesIO(file_content)
files = {"file": ("test.txt", file, "text/plain")}
headers = {"Authorization": "Bearer test_api_key"}
response = client.post(
"/v1/self-probe-file",
files=files,
headers=headers,
data={"model": "whisper-large-v3"},
)
assert response.status_code == 400
assert "Invalid file format" in response.json()["detail"]
def test_self_probe_file_missing_file():
"""Test /v1/self-probe-file endpoint with missing file"""
headers = {"Authorization": "Bearer test_api_key"}
response = client.post(
"/v1/self-probe-file",
headers=headers,
data={"model": "whisper-large-v3"},
)
assert response.status_code == 422
+1 -1
View File
@@ -67,7 +67,7 @@ Content-Type: application/json
FILE_SPEC = """
POST http://0.0.0.0:9094/v1/self-probe
POST http://0.0.0.0:9094/v1/self-probe-file
Authorization: Bearer $GROQ_API_KEY
Content-Type: multipart/form-data