92 lines
3.6 KiB
Python
92 lines
3.6 KiB
Python
"""ImpactAgent RAG 上下文注入测试（RAG 迭代 Task 4，异步形态）。
|
||
|
||
验证：
|
||
- use_rag=True 时，run_impact 发送给 LLM 的 prompt 文本包含 RAG 检索命中片段与明确小节标题。
|
||
- use_rag=False（或默认）时，prompt 文本不含 RAG 小节标题（向后兼容）。
|
||
- run_impact 为 async def，真实 await 引擎 chat_structured。
|
||
"""
|
||
import asyncio
|
||
import types
|
||
|
||
import pytest
|
||
|
||
from genesis.impact.impact_agent import ImpactAgent, _RAG_CONTEXT_TITLE
|
||
from genesis.rag.embeddings import FakeEmbedder
|
||
from genesis.rag.impact_rag import ImpactRAG
|
||
from genesis.rag.store import RagStore
|
||
|
||
|
||
class FakeEngine:
|
||
"""捕获真实 LLM 方法（chat_structured，async）收到的 prompt 文本。
|
||
|
||
方法名与签名刻意复用本仓库 InferenceEngine.chat_structured 的形参风格，
|
||
以保证 mock 的是真实接口（key=session_id/prompt/variables/schema/retry_count）。
|
||
"""
|
||
|
||
def __init__(self) -> None:
|
||
self.captured: str | None = None
|
||
self.calls = 0
|
||
|
||
async def chat_structured(self, *, session_id, prompt, variables, schema, retry_count=2):
|
||
self.captured = prompt
|
||
self.calls += 1
|
||
# 返回结构兼容 StructuredResult 的最小占位（含 data/raw_text）
|
||
return types.SimpleNamespace(data={}, raw_text=prompt)
|
||
|
||
|
||
def _make_rag(session_id: str, text: str) -> ImpactRAG:
|
||
store = RagStore(":memory:")
|
||
rag = ImpactRAG(store, FakeEmbedder())
|
||
rag.index(session_id, [("TradeApplication.java", text)])
|
||
return rag
|
||
|
||
|
||
def test_run_impact_with_rag_injects_context():
|
||
rag = _make_rag("p1", "订单创建调用 MyBatis")
|
||
engine = FakeEngine()
|
||
agent = ImpactAgent(engine=engine, rag=rag, use_rag=True)
|
||
asyncio.run(agent.run_impact("p1", "创建订单的影响"))
|
||
assert engine.captured is not None
|
||
# 命中片段（含文件名 TradeApplication.java）被注入
|
||
assert "TradeApplication" in engine.captured
|
||
# 明确小节标题被注入
|
||
assert _RAG_CONTEXT_TITLE in engine.captured
|
||
|
||
|
||
def test_run_impact_without_rag_no_context():
|
||
engine = FakeEngine()
|
||
agent = ImpactAgent(engine=engine, rag=_make_rag("p1", "x"), use_rag=False)
|
||
asyncio.run(agent.run_impact("p1", "创建订单的影响"))
|
||
assert _RAG_CONTEXT_TITLE not in engine.captured
|
||
|
||
|
||
def test_run_impact_rag_enabled_but_no_hits():
|
||
# rag 已建但索引在另一 scope，retrieve 命中为空 → 不应注入 RAG 小节
|
||
engine = FakeEngine()
|
||
agent = ImpactAgent(engine=engine, rag=_make_rag("other", "订单创建调用 MyBatis"), use_rag=True)
|
||
asyncio.run(agent.run_impact("p1", "创建订单的影响"))
|
||
assert _RAG_CONTEXT_TITLE not in engine.captured
|
||
|
||
|
||
def test_run_impact_use_rag_none_falls_back_to_instance_default():
|
||
# use_rag=None 时回退实例级 self.use_rag（此处为 True）→ 应注入 RAG 小节
|
||
engine = FakeEngine()
|
||
agent = ImpactAgent(engine=engine, rag=_make_rag("p1", "订单创建调用 MyBatis"), use_rag=True)
|
||
asyncio.run(agent.run_impact("p1", "创建订单的影响", use_rag=None))
|
||
assert _RAG_CONTEXT_TITLE in engine.captured
|
||
assert "TradeApplication" in engine.captured
|
||
|
||
|
||
def test_run_impact_explicit_false_no_context():
|
||
# 显式 use_rag=False（覆盖 use_rag is None 的 False 分支）→ 不注入 RAG 小节
|
||
engine = FakeEngine()
|
||
agent = ImpactAgent(engine=engine, rag=_make_rag("p1", "x"), use_rag=True)
|
||
asyncio.run(agent.run_impact("p1", "创建订单的影响", use_rag=False))
|
||
assert _RAG_CONTEXT_TITLE not in engine.captured
|
||
|
||
|
||
def test_run_impact_requires_engine():
|
||
agent = ImpactAgent()
|
||
with pytest.raises(RuntimeError):
|
||
asyncio.run(agent.run_impact("p1", "x"))
|