102 lines
3.3 KiB
Python
102 lines
3.3 KiB
Python
from __future__ import annotations
|
|
|
|
import inspect
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from src.rag.qa_v1_1_73 import api as api_module
|
|
from src.rag.qa_v1_1_73.pipeline import QAV1Pipeline
|
|
from src.rag.qa_v1_1_73.subquestion_orchestrator import SubquestionRetrievalOrchestrator
|
|
|
|
|
|
def test_pipeline_accepts_optional_progress_callback():
|
|
assert "progress_callback" in inspect.signature(QAV1Pipeline.answer).parameters
|
|
|
|
|
|
def test_subquestion_orchestrator_reports_review_and_bounded_retry():
|
|
events = []
|
|
review_calls = []
|
|
|
|
def retrieve(task, *, attempt, refinement):
|
|
return {
|
|
"evidence_top30": [{"evidence_id": f"{task['subquestion_id']}.A{attempt}.E1"}],
|
|
"chains": [],
|
|
}
|
|
|
|
def review(_question, groups, *, attempt):
|
|
review_calls.append(attempt)
|
|
return {
|
|
"groups": [
|
|
{
|
|
"subquestion_id": group["subquestion_id"],
|
|
"status": "PARTIAL" if attempt == 1 else "SATISFIED",
|
|
"accepted_evidence_ids": [group["evidence"][0]["evidence_id"]],
|
|
"accepted_chain_ids": [],
|
|
"refinement": {"vector_queries": ["refined"]} if attempt == 1 else {},
|
|
}
|
|
for group in groups
|
|
]
|
|
}
|
|
|
|
result = SubquestionRetrievalOrchestrator().run(
|
|
"test question",
|
|
{
|
|
"subquestions": [{"subquestion_id": "Q0", "question": "test question"}],
|
|
"synthesis_order": ["Q0"],
|
|
},
|
|
retrieve=retrieve,
|
|
review=review,
|
|
progress_callback=events.append,
|
|
)
|
|
|
|
assert review_calls == [1, 2]
|
|
assert result["audit_groups"][0]["attempt_count"] == 2
|
|
assert [event["stage"] for event in events] == [
|
|
"subquestion_retrieved",
|
|
"agent_review_started",
|
|
"subquestion_retry_started",
|
|
"subquestion_retrieved",
|
|
"agent_review_started",
|
|
"evidence_review_completed",
|
|
]
|
|
|
|
|
|
def test_api_exposes_sse_chat_stream_without_replacing_chat():
|
|
source = (
|
|
Path(__file__).resolve().parents[1] / "src/rag/qa_v1_1_73/api.py"
|
|
).read_text(encoding="utf-8")
|
|
|
|
assert '@router.post("/chat")' in source
|
|
assert '@router.post("/chat/stream")' in source
|
|
assert "StreamingResponse" in source
|
|
assert "progress_callback" in source
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sse_chat_stream_emits_progress_before_complete(monkeypatch):
|
|
class FakeLLM:
|
|
available = True
|
|
|
|
class FakePipeline:
|
|
llm = FakeLLM()
|
|
|
|
def answer(self, question, **kwargs):
|
|
kwargs["progress_callback"]({
|
|
"stage": "entities_identified",
|
|
"details": {"entities": ["WRKY"]},
|
|
})
|
|
return {"success": True, "question": question, "answer": "answer"}
|
|
|
|
monkeypatch.setattr(api_module, "get_pipeline", lambda: FakePipeline())
|
|
response = await api_module.chat_stream(api_module.QARequest(question="WRKY?"))
|
|
body = b""
|
|
async for chunk in response.body_iterator:
|
|
body += chunk.encode("utf-8") if isinstance(chunk, str) else chunk
|
|
text = body.decode("utf-8")
|
|
|
|
assert text.index("event: progress") < text.index("event: complete")
|
|
complete_data = text.split("event: complete", 1)[1].split("data:", 1)[1].strip()
|
|
assert json.loads(complete_data)["answer"] == "answer"
|