mirror of
https://github.com/msoedov/agentic_security.git
synced 2026-06-24 06:09:55 +02:00
fix(add file self probe endpoint):
This commit is contained in:
@@ -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]
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user