Files
WiseGrade/backend/tests/test_chat_sessions.py
T

340 lines
20 KiB
Python

"""Session lifecycle and recoverable background runs against the isolated test database."""
import time
from concurrent.futures import ThreadPoolExecutor
from threading import Event
from types import SimpleNamespace
from uuid import uuid4
from fastapi.testclient import TestClient
from sqlalchemy import func, select
from wisegrade import agent_tools, chat_api, services
from wisegrade.db import get_db
from wisegrade.main import app
from wisegrade.models import (AgentRun, AuditLog, ChatMessage, Conversation, CourseVersion, CriterionScore,
KnowledgeChunk, KnowledgeDocument, PendingAction, Project, ScoreEvidence,
ScoringBatch, Student, SubmissionFile, now)
def client_for(session_factory):
def override():
with session_factory() as db:
yield db
app.dependency_overrides[get_db] = override
return TestClient(app)
def test_first_message_is_atomic_idempotent_and_retryable(session_factory, monkeypatch):
monkeypatch.setattr(chat_api, "_launch", lambda *_: None)
client = client_for(session_factory)
try:
assert client.get('/api/conversations').json() == []
key = str(uuid4())
payload = {"request_id": key, "content": "列出课程"}
assert client.post('/api/conversations/first-message', json={**payload, "content": " "}).status_code == 422
assert client.get('/api/conversations').json() == []
with ThreadPoolExecutor(max_workers=2) as pool:
replies = list(pool.map(lambda _: client.post('/api/conversations/first-message', json=payload), range(2)))
assert all(reply.status_code == 200 for reply in replies)
assert replies[0].json() == replies[1].json()
conversation_id = replies[0].json()['conversation_id']
with session_factory() as db:
assert db.scalar(select(func.count(Conversation.id))) == 1
assert db.scalar(select(func.count(ChatMessage.id))) == 1
assert db.scalar(select(func.count(AgentRun.id))) == 1
detail = client.get(f'/api/conversations/{conversation_id}').json()
assert detail['scope_locked'] and detail['messages'][0]['content'] == '列出课程'
assert client.get('/api/conversations').json()[0]['running'] is True
finally:
app.dependency_overrides.clear()
def test_legacy_empty_scope_can_change_then_locks(session_factory, monkeypatch):
monkeypatch.setattr(chat_api, "_launch", lambda *_: None)
client = client_for(session_factory)
try:
legacy = client.post('/api/conversations', json={"title": "旧空会话"}).json()
assert legacy['scope_locked'] is False
with session_factory() as db:
version = CourseVersion(course_id=1, version=2, title="第二版")
db.add(version)
db.flush()
student = Student(student_no="987654321", name="测试学生")
db.add(student)
db.flush()
project = Project(course_version_id=version.id, student_id=student.id, topic="测试", source_dir="test/scope")
db.add(project)
db.commit()
version_id, project_id = version.id, project.id
changed = client.put(f"/api/conversations/{legacy['id']}/scope", json={"project_id": project_id})
assert changed.status_code == 200
assert changed.json()['course_version_id'] == version_id
assert changed.json()['project_id'] == project_id
assert client.put(f"/api/conversations/{legacy['id']}/scope", json={"course_version_id": 1, "project_id": project_id}).status_code == 400
cleared = client.put(f"/api/conversations/{legacy['id']}/scope", json={})
assert cleared.json()['course_version_id'] is None and cleared.json()['project_id'] is None
started = client.post(f"/api/conversations/{legacy['id']}/messages/start", json={"content": "开始"})
assert started.status_code == 200
assert client.put(f"/api/conversations/{legacy['id']}/scope", json={"course_version_id": version_id}).status_code == 409
finally:
app.dependency_overrides.clear()
def test_delete_related_rows_preserves_audit_and_rejects_running(session_factory, monkeypatch):
monkeypatch.setattr(chat_api, "_launch", lambda *_: None)
client = client_for(session_factory)
try:
first = client.post('/api/conversations/first-message', json={"request_id": str(uuid4()), "content": "运行中"}).json()
second = client.post('/api/conversations', json={"title": "可删除"}).json()
second_id = second['id']
with session_factory() as db:
db.add(ChatMessage(conversation_id=second_id, role="user", content="历史"))
db.add(AgentRun(conversation_id=second_id, status="completed"))
db.add(PendingAction(conversation_id=second_id, kind="export", payload={}, preview={}, expected={},
idempotency_key=uuid4().hex, status="pending", expires_at=now()))
db.add(AuditLog(action="chat_export", entity_type="conversation", entity_id=second_id, details={"result": "kept"}))
db.commit()
assert client.delete(f"/api/conversations/{first['conversation_id']}").status_code == 409
assert client.delete(f'/api/conversations/{second_id}').status_code == 200
assert client.get(f'/api/conversations/{second_id}').status_code == 404
assert client.post(f'/api/conversations/{second_id}/actions/1/confirm').status_code == 404
with session_factory() as db:
assert db.scalar(select(func.count(ChatMessage.id)).where(ChatMessage.conversation_id == second_id)) == 0
assert db.scalar(select(func.count(AgentRun.id)).where(AgentRun.conversation_id == second_id)) == 0
assert db.scalar(select(func.count(PendingAction.id)).where(PendingAction.conversation_id == second_id)) == 0
assert db.scalar(select(AuditLog).where(AuditLog.entity_id == second_id)).details == {"result": "kept"}
finally:
app.dependency_overrides.clear()
def test_background_runs_survive_no_stream_and_isolate_parallel_events(session_factory, monkeypatch):
gates = {"甲": Event(), "乙": Event()}
class FakeClient:
def tool_call(self, messages, _tools):
label = next(message['content'] for message in messages if message['role'] == 'user')
if not any(message['role'] == 'tool' for message in messages):
return {"role": "assistant", "content": None, "tool_calls": [{"id": f"call-{label}",
"type": "function", "function": {"name": "list_courses", "arguments": "{}"}}]}
return {"role": "assistant", "content": f"回答{label}", "tool_calls": []}
def execute(_db, _conversation, _name, _arguments, user_text):
assert gates[user_text].wait(8)
return {"courses": [{"id": 1}]}
monkeypatch.setattr(chat_api, "client_for_chat", lambda _db: FakeClient())
monkeypatch.setattr(chat_api, "execute", execute)
client = client_for(session_factory)
try:
ids = [client.post('/api/conversations/first-message', json={"request_id": str(uuid4()), "content": label}).json()['conversation_id'] for label in gates]
deadline = time.monotonic() + 8
while time.monotonic() < deadline:
details = [client.get(f'/api/conversations/{id}').json() for id in ids]
if all(d['runs'][0]['tool_calls'] for d in details): break
time.sleep(.1)
assert all(d['runs'][0]['tool_calls'][0]['id'].endswith(f'call-{label}') for d, label in zip(details, gates))
assert all(d['runs'][0]['status'] == 'running' for d in details)
subscription = chat_api._stream_run(session_factory, details[0]['runs'][0]['id'])
assert 'event: run' in next(subscription)
subscription.close() # Closing the SSE reader must not cancel the worker.
gates['甲'].set()
deadline = time.monotonic() + 8
while time.monotonic() < deadline:
first = client.get(f'/api/conversations/{ids[0]}').json()
if first['runs'][0]['status'] == 'completed': break
time.sleep(.1)
other = client.get(f'/api/conversations/{ids[1]}').json()
assert first['messages'][-1]['content'] == '回答甲'
assert first['messages'][-1]['tool_calls'][0]['id'].endswith('call-甲')
assert other['runs'][0]['status'] == 'running'
assert all('回答甲' not in m['content'] for m in other['messages'])
gates['乙'].set()
deadline = time.monotonic() + 8
while time.monotonic() < deadline:
other = client.get(f'/api/conversations/{ids[1]}').json()
if other['runs'][0]['status'] == 'completed': break
time.sleep(.1)
assert other['messages'][-1]['content'] == '回答乙'
assert other['messages'][-1]['tool_calls'][0]['id'].endswith('call-乙')
assert other['messages'][-1]['tool_calls'][0]['summary'] == '返回 1 条courses记录'
assert other['messages'][-1]['tool_calls'][0]['result'] == {'courses': {'count': 1}}
first_events = client.get(f"/api/conversations/{ids[0]}/runs/{first['runs'][0]['id']}/events").text
second_events = client.get(f"/api/conversations/{ids[1]}/runs/{other['runs'][0]['id']}/events").text
assert 'call-甲' in first_events and 'call-乙' not in first_events
assert 'call-乙' in second_events and 'call-甲' not in second_events
assert client.get(f"/api/conversations/{ids[0]}/runs/{other['runs'][0]['id']}/events").status_code == 404
finally:
for gate in gates.values(): gate.set()
app.dependency_overrides.clear()
def test_tool_record_keeps_bounded_query_but_redacts_review_text_and_secrets():
saved = chat_api._safe_arguments({"project_id": 2, "query": "电机控制", "api_key": "secret", "comment": "敏感评语"})
assert saved["project_id"] == 2
assert saved["query"] == "电机控制"
assert "secret" not in str(saved) and "敏感评语" not in str(saved)
assert len(chat_api._safe_arguments({"query": "a" * 500})["query"]) == 300
assert "学生原文" not in chat_api._result_summary({"found": True, "citations": [{"quote": "学生原文"}]})
assert "学生原文" not in str(chat_api._safe_result({"citations": [{"quote": "学生原文"}]}))
search = chat_api._safe_result({"found": True, "citations": [{"title": "报告.pdf", "page": 48,
"quote": "可核验片段", "url": "/api/pages/48", "api_key": "secret"}]}, "search_project_material")
assert search['citations'][0]['quote'] == '可核验片段'
assert 'secret' not in str(search)
def test_score_citations_link_text_and_restore_legacy_image_only(session_factory):
with session_factory() as db:
student = Student(student_no="987654320", name="测试学生")
db.add(student)
db.flush()
project = Project(course_version_id=1, student_id=student.id, topic="测试项目", source_dir="test/citation")
db.add(project)
db.flush()
file = SubmissionFile(project_id=project.id, category="final_report", relative_path="test/citation/report.pdf",
mime_type="application/pdf", size=1, mtime=1, sha256="a" * 64,
recognition_status="recognized", parse_status="ready")
db.add(file)
db.flush()
document = KnowledgeDocument(course_version_id=1, project_id=project.id, submission_file_id=file.id,
title="报告.pdf", document_type="final_report", source_hash="a" * 64,
index_version=2, status="ready")
db.add(document)
db.flush()
db.add(KnowledgeChunk(document_id=document.id, course_version_id=1, project_id=project.id,
document_type="final_report", page=2, section="测试", text="这里有原文证据和分析"))
created = services.create_batch(SimpleNamespace(course_version_id=1, rubric_version_id=1), db)
batch = db.get(ScoringBatch, created["id"])
score = CriterionScore(batch_id=batch.id, project_id=project.id,
criterion_id=batch.rubric_snapshot["criteria"][0]["criterion_id"],
model_score=8, status="scored")
db.add(score)
db.flush()
db.add(ScoreEvidence(score_id=score.id, document_id=document.id, submission_file_id=file.id,
page=2, quote="原文证据", image_path="page-2.png"))
conversation = Conversation(title="旧评分引用", course_version_id=1, batch_id=batch.id, project_id=project.id)
db.add(conversation)
db.flush()
image_url = f"/api/assets/{file.id}/2"
db.add(ChatMessage(conversation_id=conversation.id, role="assistant", content="历史回答", citations=[
{"title": "报告.pdf", "page": 2, "quote": "原文证据", "url": image_url},
{"title": "报告.pdf", "page": 2, "quote": "不匹配的引文", "url": image_url}]))
db.commit()
expected = f"/api/knowledge/documents/{document.id}/pages/2"
generated = agent_tools.execute(db, conversation, "get_scores",
{"batch_id": batch.id, "project_id": project.id}, "")
assert generated["citations"][0]["text_url"] == expected
assert chat_api._answer_citations(generated["citations"])[0]["text_url"] == expected
assert generated["citations"][0]["image_url"] == image_url
conversation_id = conversation.id
client = client_for(session_factory)
try:
detail = client.get(f"/api/conversations/{conversation_id}").json()
citations = detail["messages"][0]["citations"]
assert citations[0]["text_url"] == expected
assert "text_url" not in citations[1]
assert "原文证据" in client.get(expected).json()["text"]
with session_factory() as db:
stored = db.scalar(select(ChatMessage).where(ChatMessage.conversation_id == conversation_id))
assert "text_url" not in stored.citations[0]
finally:
app.dependency_overrides.clear()
def test_late_citations_and_search_snippets_survive_restoration(session_factory, monkeypatch):
class FakeClient:
def tool_call(self, messages, _tools):
if not any(message['role'] == 'tool' for message in messages):
return {"role": "assistant", "content": None, "tool_calls": [
{"id": f"find-{index}", "type": "function", "function": {"name": "search_project_material",
"arguments": f'{{"project_id": 1, "query": "检索原文{index}"}}'}} for index in range(5)]}
return {"role": "assistant", "content": "综合报告第48页给出改进方向。", "tool_calls": []}
count = 0
def execute(_db, _conversation, _name, _arguments, _text):
nonlocal count
count += 1
page = 48 if count == 5 else count
return {"found": True, "citations": [{"document_id": 10, "chunk_id": count * 10 + index,
"title": "综合报告.pdf", "page": page, "quote": f"第{page}页原文片段{index}",
"url": f"/api/knowledge/documents/10/pages/{page}"} for index in range(6)]}
monkeypatch.setattr(chat_api, "client_for_chat", lambda _db: FakeClient())
monkeypatch.setattr(chat_api, "execute", execute)
client = client_for(session_factory)
try:
created = client.post('/api/conversations', json={"title": "引用恢复"}).json()
response = client.post(f"/api/conversations/{created['id']}/messages", json={"content": "总结"})
assert response.status_code == 200
detail = client.get(f"/api/conversations/{created['id']}").json()
answer = detail['messages'][-1]
assert len(answer['citations']) == 30
assert [item['page'] for item in answer['citations'][-6:]] == [48] * 6
assert answer['tool_calls'][-1]['result']['citations'][0]['quote'] == '第48页原文片段0'
assert detail['runs'][0]['tool_calls'][-1]['result']['citations'][0]['page'] == 48
finally:
app.dependency_overrides.clear()
def test_large_score_result_keeps_structured_rows_for_model_and_history(session_factory, monkeypatch):
seen = []
class FakeClient:
def tool_call(self, messages, _tools):
if not any(message['role'] == 'tool' for message in messages):
return {"role": "assistant", "content": None, "tool_calls": [{"id": "score-call", "type": "function",
"function": {"name": "get_scores", "arguments": '{"batch_id": 3}'}}]}
seen.append(messages[-1]['content'])
return {"role": "assistant", "content": "已读取评分项目。", "tool_calls": []}
def execute(_db, _conversation, _name, _arguments, _text):
return {"batch_id": 3, "count": 1, "project_offset": 0, "criterion_offset": 0,
"next_project_offset": None, "next_criterion_offset": 6,
"results": [{"project_id": 1, "student_no": "1001", "name": "学生甲", "totals": {},
"url": "/?page=scores&batch_id=3&project_id=1", "total_rows": 22,
"rows": [{"criterion_id": "C1", "criterion_name": "功能", "score_id": 7,
"model_score": 8, "teacher_score": None, "final_score": 8,
"max_score": 10, "status": "scored", "error": "", "evidence": []}]}],
"citations": [], "unrelated_large_field": "x" * 13000}
monkeypatch.setattr(chat_api, "client_for_chat", lambda _db: FakeClient())
monkeypatch.setattr(chat_api, "execute", execute)
client = client_for(session_factory)
try:
created = client.post('/api/conversations', json={"title": "评分恢复"}).json()
assert client.post(f"/api/conversations/{created['id']}/messages", json={"content": "查分"}).status_code == 200
assert seen and 'C1' in seen[0] and 'next_criterion_offset' in seen[0]
assert '结果过长' not in seen[0]
result = client.get(f"/api/conversations/{created['id']}").json()['messages'][-1]['tool_calls'][0]['result']
assert result['results'][0]['rows'][0]['criterion_id'] == 'C1'
assert result['next_criterion_offset'] == 6
finally:
app.dependency_overrides.clear()
def test_upgrade_from_legacy_0003_preserves_history(session_factory):
from pathlib import Path
from alembic import command
from alembic.config import Config
from sqlalchemy import text
with session_factory() as db:
row = Conversation(title="旧记录")
db.add(row)
db.flush()
db.add(ChatMessage(conversation_id=row.id, role="user", content="旧问题"))
db.add(AgentRun(conversation_id=row.id, status="completed"))
db.commit()
conversation_id = row.id
engine = session_factory.kw['bind']
with engine.begin() as connection:
connection.exec_driver_sql('ALTER TABLE conversations DROP COLUMN first_request_id')
connection.exec_driver_sql('ALTER TABLE agent_runs DROP COLUMN tool_calls')
connection.exec_driver_sql('ALTER TABLE agent_runs DROP COLUMN events')
connection.exec_driver_sql("UPDATE alembic_version SET version_num = '0003'")
root = Path(__file__).resolve().parents[1]
config = Config(str(root / 'alembic.ini'))
config.set_main_option('script_location', str(root / 'migrations'))
command.upgrade(config, 'head')
with session_factory() as db:
assert db.scalar(text('SELECT version_num FROM alembic_version')) == '0004'
assert db.get(Conversation, conversation_id).title == '旧记录'
assert db.scalar(select(ChatMessage.content).where(ChatMessage.conversation_id == conversation_id)) == '旧问题'
run = db.scalar(select(AgentRun).where(AgentRun.conversation_id == conversation_id))
assert run.tool_calls == [] and run.events == []