86 lines
3.3 KiB
Python
86 lines
3.3 KiB
Python
"""ImpactAgent RAG 上下文注入测试（RAG 迭代 Task 4）。
|
|
|
|
验证：
|
|
- use_rag=True 时，run_impact 发送给 LLM 的 prompt 文本包含 RAG 检索命中片段与明确小节标题。
|
|
- use_rag=False（或默认）时，prompt 文本不含 RAG 小节标题（向后兼容）。
|
|
"""
|
|
from genesis.impact.impact_agent import ImpactAgent
|
|
from genesis.rag.embeddings import FakeEmbedder
|
|
from genesis.rag.impact_rag import ImpactRAG
|
|
from genesis.rag.store import RagStore
|
|
|
|
|
|
_RAG_SECTION_TITLE = "# 既有系统关联上下文（RAG 检索，辅助判断影响范围）"
|
|
|
|
|
|
class FakeEngine:
|
|
"""捕获真实 LLM 方法（chat_structured）收到的 prompt 文本。
|
|
|
|
方法名与签名刻意复用本仓库 InferenceEngine.chat_structured 的形参风格，
|
|
以保证 mock 的是真实接口（key=session_id/prompt/variables/schema/retry_count）。
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self.last_prompt: str | None = None
|
|
self.calls = 0
|
|
|
|
def chat_structured(self, *, session_id, prompt, variables, schema, retry_count=2):
|
|
self.last_prompt = prompt
|
|
self.calls += 1
|
|
# 返回结构兼容 ChatResult 的最小占位（测试仅校验 prompt 注入）
|
|
return {"session_id": session_id, "prompt": prompt}
|
|
|
|
|
|
def _make_rag(session_id: str, sources):
|
|
store = RagStore(":memory:")
|
|
rag = ImpactRAG(store, FakeEmbedder())
|
|
rag.index(session_id, sources)
|
|
return store, rag
|
|
|
|
|
|
def test_run_impact_with_rag_injects_context():
|
|
session_id = "sess-rag"
|
|
store, rag = _make_rag(session_id, [("TradeApplication.java", "订单创建调用 MyBatis")])
|
|
try:
|
|
engine = FakeEngine()
|
|
agent = ImpactAgent(engine=engine, rag=rag, use_rag=True)
|
|
agent.run_impact(session_id, requirements_text="创建订单的影响", k=5)
|
|
prompt = engine.last_prompt
|
|
assert prompt is not None
|
|
# 命中片段（含文件名 TradeApplication.java）被注入
|
|
assert "TradeApplication" in prompt
|
|
# 明确小节标题被注入
|
|
assert _RAG_SECTION_TITLE in prompt
|
|
finally:
|
|
store.close()
|
|
|
|
|
|
def test_run_impact_without_rag_no_context():
|
|
session_id = "sess-no-rag"
|
|
store, rag = _make_rag(session_id, [("TradeApplication.java", "订单创建调用 MyBatis")])
|
|
try:
|
|
engine = FakeEngine()
|
|
agent = ImpactAgent(engine=engine, rag=rag, use_rag=False)
|
|
agent.run_impact(session_id, requirements_text="创建订单的影响", k=5)
|
|
prompt = engine.last_prompt
|
|
# 显式关闭 RAG：不含小节标题，也不含检索片段
|
|
assert _RAG_SECTION_TITLE not in prompt
|
|
assert "TradeApplication" not in prompt
|
|
finally:
|
|
store.close()
|
|
|
|
|
|
def test_run_impact_default_no_rag_no_context():
|
|
# 默认 use_rag 为 False（未显式开启），行为与关闭一致
|
|
session_id = "sess-default"
|
|
store, rag = _make_rag(session_id, [("TradeApplication.java", "订单创建调用 MyBatis")])
|
|
try:
|
|
engine = FakeEngine()
|
|
agent = ImpactAgent(engine=engine, rag=rag) # 不传 use_rag
|
|
agent.run_impact(session_id, requirements_text="创建订单的影响", k=5)
|
|
prompt = engine.last_prompt
|
|
assert _RAG_SECTION_TITLE not in prompt
|
|
assert "TradeApplication" not in prompt
|
|
finally:
|
|
store.close()
|