Spaces:
Running
Running
File size: 6,196 Bytes
26d7b76 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | """
Unit and integration tests for FrugalGPT cascading router and prompt adaptation.
"""
import json
import pytest
from fastapi.testclient import TestClient
from inferroute.main import app
from inferroute.prompt_adapter import adapt_prompt, compress_few_shot_examples
from inferroute.validator import ReliabilityScorer
client = TestClient(app)
HEADERS = {"Authorization": "Bearer sk-inferroute-demo"}
# ── Test 1: Prompt Adaptation Trimming ─────────────────────────────────────────
def test_prompt_adaptation_trimming():
"""
Ensure prompt adapter trims few-shot examples for cheap local backends (ollama, vllm)
but preserves them for premium cloud models (openai, gemini).
"""
few_shot_prompt = (
"Classify the sentiment of these headlines:\n\n"
"Example 1:\nHeadline: Gold prices surge amid inflation fears.\nSentiment: positive\n\n"
"Example 2:\nHeadline: Retail sales plummet in Q3.\nSentiment: negative\n\n"
"Example 3:\nHeadline: Central bank holds interest rates steady.\nSentiment: neutral\n\n"
"Headline: Stock index hits record high today.\nSentiment:"
)
messages = [{"role": "user", "content": few_shot_prompt}]
# Cheap backend -> should compress few-shot examples to keep max 1 example
compressed = adapt_prompt(messages, "ollama")
assert len(compressed[0]["content"]) < len(few_shot_prompt)
assert "Example 1:" in compressed[0]["content"]
assert "Example 2:" not in compressed[0]["content"]
assert "Stock index hits record high today." in compressed[0]["content"]
# Premium backend -> should preserve full few-shots
preserved = adapt_prompt(messages, "openai")
assert len(preserved[0]["content"]) == len(few_shot_prompt)
assert "Example 2:" in preserved[0]["content"]
assert "Example 3:" in preserved[0]["content"]
# ── Test 2: Reliability Scorer Math Classification ───────────────────────────
def test_reliability_scorer_math():
"""
Test that ReliabilityScorer correctly matches known prompt templates
and evaluates math answer correctness.
"""
scorer = ReliabilityScorer()
req = {
"messages": [{"role": "user", "content": "Solve for x: 5x - 15 = 20. Output ONLY the final numeric value of x as an integer."}]
}
# Correct response (7) should score 1.0
assert scorer.evaluate_reliability(req, "7") == 1.0
assert scorer.evaluate_reliability(req, "The answer is 7") == 1.0
# Incorrect response (5) should score 0.0
assert scorer.evaluate_reliability(req, "5") == 0.0
assert scorer.evaluate_reliability(req, "x = 10") == 0.0
# ── Test 3: Server-side Cascade Blocking Integration ─────────────────────────
def test_cascade_blocking_integration():
"""
Test cascade routing logic on a math prompt.
Since Ollama returns incorrect math answers, the router should escalate
and eventually accept OpenAI/Gemini response.
"""
payload = {
"model": "edge/auto",
"messages": [{"role": "user", "content": "Solve for x: 5x - 15 = 20. Output ONLY the final numeric value of x as an integer."}],
"stream": False,
"routing": {
"policy": "cascade",
"acceptance_threshold": 0.8,
"cascade_chain": ["ollama", "vllm", "openai"]
}
}
response = client.post("/v1/chat/completions", headers=HEADERS, json=payload)
assert response.status_code == 200
data = response.json()
# The final accepted model should be openai (since ollama and vllm return wrong math results)
assert data["model"] == "openai"
assert data["choices"][0]["message"]["content"].strip() == "7"
# Verify cascade routing trace metadata
route = data.get("route", {})
assert route["policy"] == "cascade"
assert route["fallback_count"] > 0
assert len(route["cascade_steps"]) > 0
# Confirm the first step failed quality but last step accepted
assert route["cascade_steps"][0]["backend"] == "ollama"
assert route["cascade_steps"][0]["accepted"] is False
assert route["cascade_steps"][-1]["backend"] == "openai"
assert route["cascade_steps"][-1]["accepted"] is True
# ── Test 4: Server-side Cascade Streaming Integration ────────────────────────
def test_cascade_streaming_integration():
"""
Test cascade streaming execution flow.
Ensures that stream chunks are correctly buffered and pumped only when accepted,
and returns a final end-of-stream stats chunk.
"""
payload = {
"model": "edge/auto",
"messages": [{"role": "user", "content": "Solve for x: 5x - 15 = 20. Output ONLY the final numeric value of x as an integer."}],
"stream": True,
"routing": {
"policy": "cascade",
"acceptance_threshold": 0.8,
"cascade_chain": ["ollama", "vllm", "openai"]
}
}
response = client.post("/v1/chat/completions", headers=HEADERS, json=payload)
assert response.status_code == 200
# Parse SSE stream chunks
lines = response.text.split("\n")
events = []
for line in lines:
if line.startswith("data:"):
data_str = line[5:].strip()
if data_str != "[DONE]":
events.append(json.loads(data_str))
assert len(events) > 0
# Find the final stats chunk
stats_chunk = next((e for e in events if e.get("id") == "inferroute-stream-end"), None)
assert stats_chunk is not None
route = stats_chunk.get("route", {})
assert route["policy"] == "cascade"
assert route["selected_backend"] == "openai"
assert route["fallback_count"] > 0
cascade_steps = route.get("cascade_steps", [])
assert cascade_steps[0]["backend"] == "ollama"
assert cascade_steps[0]["accepted"] is False
assert cascade_steps[-1]["backend"] == "openai"
assert cascade_steps[-1]["accepted"] is True
|