Implement Phase 19: Query Rewriting via LangGraph node
- query_rewrite 노드 추가 (agent → query_rewrite → tools 순서) - route_after_agent: search_documents 호출 시에만 query_rewrite 라우팅, 그 외 직접 tools - tools_condition(prebuilt) 제거 → 커스텀 라우팅 함수로 대체 - query_rewrite_node: 구어체 쿼리를 키워드 중심 문장으로 변환 - 이전 대화 2턴 컨텍스트로 대명사·지시어 해소 - enable_thinking=False 바인딩으로 불필요한 사고 과정 제거 - __query_rewrite 커스텀 이벤트 emit → RAG_VERBOSE 시 변환 결과 출력 - QUERY_REWRITE_ENABLED=true 로 활성화 (기본값 false) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -7,8 +7,8 @@ from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage, Sys
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.memory import MemorySaver
|
||||
from langgraph.config import get_stream_writer
|
||||
from langgraph.graph import START, MessagesState, StateGraph
|
||||
from langgraph.prebuilt import ToolNode, tools_condition
|
||||
from langgraph.graph import END, START, MessagesState, StateGraph
|
||||
from langgraph.prebuilt import ToolNode
|
||||
|
||||
from services.agent.tools import get_current_date, make_memory_tools, make_retriever_tool, make_search_tool, web_search
|
||||
|
||||
@@ -28,6 +28,7 @@ class AgentService:
|
||||
rag_show_sources: bool = False,
|
||||
langgraph_verbose: bool = False,
|
||||
think_verbose: bool = False,
|
||||
query_rewrite_enabled: bool = False,
|
||||
user_profile_repository=None,
|
||||
conversation_repository=None,
|
||||
user_id: str = "default",
|
||||
@@ -37,6 +38,7 @@ class AgentService:
|
||||
self._rag_show_sources = rag_show_sources
|
||||
self._langgraph_verbose = langgraph_verbose
|
||||
self._think_verbose = think_verbose
|
||||
self._query_rewrite_enabled = query_rewrite_enabled
|
||||
self._source_buffer: list[dict] = []
|
||||
self._thread_id = "default"
|
||||
self._profile_repo = user_profile_repository
|
||||
@@ -133,11 +135,76 @@ class AgentService:
|
||||
additional_kwargs=extra,
|
||||
)]}
|
||||
|
||||
async def query_rewrite_node(state: MessagesState, config: RunnableConfig) -> dict:
|
||||
last_msg = state["messages"][-1]
|
||||
if not (hasattr(last_msg, "tool_calls") and last_msg.tool_calls):
|
||||
return {}
|
||||
|
||||
# 최근 사용자 메시지 2개를 컨텍스트로 활용 (대명사·지시어 해소)
|
||||
recent_human = [m.content for m in state["messages"][:-1]
|
||||
if isinstance(m, HumanMessage)][-2:]
|
||||
ctx = ("\n\n이전 대화 컨텍스트:\n" + "\n".join(f"- {m}" for m in recent_human)
|
||||
if recent_human else "")
|
||||
|
||||
try:
|
||||
writer = get_stream_writer()
|
||||
except Exception:
|
||||
writer = None
|
||||
|
||||
_rewrite_llm = chat_model.bind(enable_thinking=False)
|
||||
new_tool_calls = []
|
||||
for tc in last_msg.tool_calls:
|
||||
if tc["name"] == "search_documents":
|
||||
original = tc["args"].get("query", "")
|
||||
prompt = (
|
||||
f"다음 구어체 질문을 문서 검색에 최적화된 키워드 중심 문장으로 변환하세요.{ctx}\n\n"
|
||||
f"규칙:\n"
|
||||
f"- 핵심 개념과 전문용어를 포함하세요\n"
|
||||
f"- 대명사(이것, 그것, 그 논문 등)는 구체적인 명칭으로 교체하세요\n"
|
||||
f"- 변환된 질문만 한 문장으로 출력하세요. 부가 설명 없이 질문만 출력하세요\n\n"
|
||||
f"원본 질문: {original}\n최적화된 질문:"
|
||||
)
|
||||
try:
|
||||
result = await _rewrite_llm.ainvoke([HumanMessage(content=prompt)])
|
||||
rewritten = result.content.strip()
|
||||
except Exception as e:
|
||||
print(f"[QueryRewrite] 실패: {e}")
|
||||
rewritten = original
|
||||
if rewritten and rewritten != original:
|
||||
new_tool_calls.append({**tc, "args": {**tc["args"], "query": rewritten}})
|
||||
if writer:
|
||||
writer({"__query_rewrite": {"original": original, "rewritten": rewritten}})
|
||||
else:
|
||||
new_tool_calls.append(tc)
|
||||
else:
|
||||
new_tool_calls.append(tc)
|
||||
|
||||
if not last_msg.id:
|
||||
return {}
|
||||
new_msg = AIMessage(
|
||||
id=last_msg.id,
|
||||
content=last_msg.content,
|
||||
tool_calls=new_tool_calls,
|
||||
additional_kwargs=last_msg.additional_kwargs,
|
||||
)
|
||||
return {"messages": [new_msg]}
|
||||
|
||||
def route_after_agent(state: MessagesState) -> str:
|
||||
last_msg = state["messages"][-1]
|
||||
if not (hasattr(last_msg, "tool_calls") and last_msg.tool_calls):
|
||||
return END
|
||||
if self._query_rewrite_enabled:
|
||||
if any(tc["name"] == "search_documents" for tc in last_msg.tool_calls):
|
||||
return "query_rewrite"
|
||||
return "tools"
|
||||
|
||||
builder = StateGraph(MessagesState)
|
||||
builder.add_node("agent", call_model)
|
||||
builder.add_node("query_rewrite", query_rewrite_node)
|
||||
builder.add_node("tools", ToolNode(tools))
|
||||
builder.add_edge(START, "agent")
|
||||
builder.add_conditional_edges("agent", tools_condition)
|
||||
builder.add_conditional_edges("agent", route_after_agent)
|
||||
builder.add_edge("query_rewrite", "tools")
|
||||
builder.add_edge("tools", "agent")
|
||||
|
||||
self._agent = builder.compile(checkpointer=MemorySaver())
|
||||
@@ -176,8 +243,13 @@ class AgentService:
|
||||
):
|
||||
mode, data = stream_event
|
||||
|
||||
# ── custom 이벤트 — call_model writer가 emit한 thinking 토큰 ──
|
||||
# ── custom 이벤트 ────────────────────────────────────────────
|
||||
if mode == "custom":
|
||||
if isinstance(data, dict) and "__query_rewrite" in data:
|
||||
info = data["__query_rewrite"]
|
||||
if lg or self._rag_verbose:
|
||||
yield f'\n쿼리 최적화: "{info["original"]}" → "{info["rewritten"]}"\n'
|
||||
continue
|
||||
if isinstance(data, dict) and "__thinking" in data:
|
||||
# thinking 첫 토큰 도착 시 agent 레이블 + prev_node 갱신
|
||||
if "agent" != prev_node:
|
||||
@@ -209,12 +281,13 @@ class AgentService:
|
||||
thinking_open = False
|
||||
content_started = False
|
||||
if lg:
|
||||
elapsed = time.perf_counter() - start_time
|
||||
if node == "agent":
|
||||
elapsed = time.perf_counter() - start_time
|
||||
label = "agent: 검색 결과 반영 중" if prev_node == "tools" else "agent: 질문 분석 중"
|
||||
yield f"\n[LangGraph → {label}] ({elapsed:.2f}s)\n"
|
||||
elif node == "query_rewrite":
|
||||
yield f"\n[LangGraph → query_rewrite: 쿼리 최적화 중] ({elapsed:.2f}s)\n"
|
||||
elif node == "tools":
|
||||
elapsed = time.perf_counter() - start_time
|
||||
yield f"\n[LangGraph → tools: 도구 실행 중] ({elapsed:.2f}s)\n"
|
||||
prev_node = node
|
||||
|
||||
|
||||
Reference in New Issue
Block a user