init
This commit is contained in:
+31
-6
@@ -4,29 +4,40 @@ from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage, AI
|
||||
from .state import AgentState
|
||||
from services.prompt_manager import get_prompt_manager
|
||||
from services.template_matcher import get_template_matcher
|
||||
from services.sql_prompt_manager import SqlPromptManager
|
||||
from services.sql_prompt_manager import get_sql_prompt_manager
|
||||
from tools.sr_api_tool import SrApiQueryTool
|
||||
|
||||
|
||||
def _short(value, max_len: int = 500) -> str:
|
||||
text = str(value)
|
||||
return text if len(text) <= max_len else text[:max_len] + "..."
|
||||
|
||||
|
||||
def process_input(state: AgentState) -> AgentState:
|
||||
"""处理用户输入"""
|
||||
print("[process_input][in] messages=", _short(state.messages))
|
||||
state.current_step = "processed"
|
||||
print("[process_input][out] current_step=", state.current_step)
|
||||
return state
|
||||
|
||||
|
||||
def generate_response(state: AgentState, model) -> AgentState:
|
||||
"""使用 LLM 生成回复"""
|
||||
print("[generate_response][in] context_keys=", list((state.context or {}).keys()))
|
||||
sr_api_result = state.context.get("sr_api_result")
|
||||
if sr_api_result:
|
||||
state.messages.append(AIMessage(content=str(sr_api_result)))
|
||||
print("[generate_response][out] source=sr_api_result")
|
||||
return state
|
||||
final_sql = state.context.get("final_sql")
|
||||
if final_sql:
|
||||
state.messages.append(AIMessage(content=final_sql))
|
||||
print("[generate_response][out] source=final_sql")
|
||||
return state
|
||||
if state.messages:
|
||||
response = model.invoke(state.messages)
|
||||
state.messages.append(response)
|
||||
print("[generate_response][out] source=model_invoke")
|
||||
return state
|
||||
|
||||
|
||||
@@ -39,19 +50,25 @@ def normalize_input(state: AgentState, model) -> AgentState:
|
||||
if not isinstance(last_message, HumanMessage):
|
||||
return state
|
||||
|
||||
print("[normalize_input][in] user_input=", _short(last_message.content))
|
||||
|
||||
prompt_manager = get_prompt_manager()
|
||||
system_prompt = SystemMessage(
|
||||
content=prompt_manager.get("system", "english_normalizer")
|
||||
normalizer_prompt = (
|
||||
prompt_manager.get("system", "english_normalizer")
|
||||
or prompt_manager.get("user", "english_normalizer")
|
||||
)
|
||||
system_prompt = SystemMessage(content=normalizer_prompt)
|
||||
|
||||
response = model.invoke([system_prompt, HumanMessage(content=last_message.content)])
|
||||
normalized = response.content if hasattr(response, "content") else str(response)
|
||||
print("[normalize_input][out] normalized=", _short(normalized))
|
||||
|
||||
state.context["original_input"] = last_message.content
|
||||
state.context["normalized_input"] = normalized
|
||||
|
||||
matcher = get_template_matcher()
|
||||
state.context["table_match"] = matcher.match(normalized)
|
||||
print("[normalize_input][out] table_match=", _short(state.context.get("table_match")))
|
||||
return state
|
||||
|
||||
|
||||
@@ -62,11 +79,16 @@ def generate_sql(state: AgentState, model) -> AgentState:
|
||||
normalized = state.context.get("normalized_input")
|
||||
|
||||
if not table_name or not normalized:
|
||||
print("[generate_sql][skip] missing table_name or normalized")
|
||||
return state
|
||||
|
||||
prompt_manager = SqlPromptManager()
|
||||
print("[generate_sql][in] table_name=", table_name)
|
||||
print("[generate_sql][in] normalized=", _short(normalized))
|
||||
|
||||
prompt_manager = get_sql_prompt_manager()
|
||||
prompt_data = prompt_manager.get_prompt(table_name)
|
||||
if not prompt_data:
|
||||
print("[generate_sql][skip] prompt not found for table=", table_name)
|
||||
return state
|
||||
|
||||
prompt_text = json.dumps(prompt_data, ensure_ascii=False, indent=2)
|
||||
@@ -75,8 +97,11 @@ def generate_sql(state: AgentState, model) -> AgentState:
|
||||
user_content = f"User question (normalized English): {normalized}"
|
||||
response = model.invoke([SystemMessage(content=system_content), HumanMessage(content=user_content)])
|
||||
sql_text = response.content if hasattr(response, "content") else str(response)
|
||||
print("[generate_sql][out] sql=", _short(sql_text))
|
||||
|
||||
state.context["final_sql"] = sql_text
|
||||
tool = SrApiQueryTool()
|
||||
state.context["sr_api_result"] = tool.run(json.dumps({"sql": sql_text}, ensure_ascii=False))
|
||||
if not state.context.get("skip_sr_api"):
|
||||
tool = SrApiQueryTool()
|
||||
state.context["sr_api_result"] = tool.run(json.dumps({"sql": sql_text}, ensure_ascii=False))
|
||||
print("[generate_sql][out] sr_api_result=", _short(state.context.get("sr_api_result")))
|
||||
return state
|
||||
|
||||
Reference in New Issue
Block a user