kg_agent/sug_v2_1.py

import asyncio
import json
import os
from datetime import datetime

from agents import Agent, Runner, function_tool
from lib.my_trace import set_trace
from typing import Literal
from dataclasses import dataclass
from pydantic import BaseModel, Field


from lib.utils import read_file_as_string
from script.search_recommendations.xiaohongshu_search_recommendations import XiaohongshuSearchRecommendations

from agents import Agent, RunContextWrapper, Runner, function_tool

from pydantic import BaseModel, Field


class RunContext(BaseModel):
    q_with_context: str
    q_context: str
    q: str
    log_url: str
    log_dir: str
    # 中间数据记录 - 按时间顺序记录所有操作
    operations_history: list[dict] = Field(default_factory=list, description="记录所有操作的历史，包括 get_query_suggestions 和 modify_query")


eval_insrtuctions = """
你是一个专业的评估专家，负责评估给定的搜索query是否满足原始问题和需求，给出出评分和简明扼要的理由。
"""

@dataclass
class EvaluationFeedback:
    reason: str=Field(..., description="简明扼要的理由")
    score: float=Field(..., description="评估结果，1表示等价，0表示不等价，中间值表示部分等价")

evaluator = Agent[None](
    name="评估专家",
    instructions=eval_insrtuctions,
    output_type=EvaluationFeedback,
)

@function_tool
async def get_query_suggestions(wrapper: RunContextWrapper[RunContext], query: str):
    """Fetch search recommendations from Xiaohongshu."""
    xiaohongshu_api = XiaohongshuSearchRecommendations()
    query_suggestions = xiaohongshu_api.get_recommendations(keyword=query)
    print(query_suggestions)
    async def evaluate_single_query(q_sug: str, q_with_context: str):
        """Evaluate a single query suggestion."""
        eval_input = f"""
{q_with_context}
<待评估的推荐query>
{q_sug}
</待评估的推荐query>
        """
        evaluator_result = await Runner.run(evaluator, eval_input)
        result: EvaluationFeedback = evaluator_result.final_output
        return {
            "query": q_sug,
            "score": result.score,
            "reason": result.reason,
        }

    # 并发执行所有评估任务
    q_with_context = wrapper.context.q_with_context
    res = []
    if query_suggestions:
        res = await asyncio.gather(*[evaluate_single_query(q_sug, q_with_context) for q_sug in query_suggestions])
    else:
        res = '未返回任何推荐词'

    # 记录到 RunContext
    wrapper.context.operations_history.append({
        "operation_type": "get_query_suggestions",
        "timestamp": datetime.now().isoformat(),
        "query": query,
        "suggestions": query_suggestions,
        "evaluations": res,
    })

    return res

@function_tool
def modify_query(wrapper: RunContextWrapper[RunContext], original_query: str, operation_type: str, new_query: str, reason: str):
    """
    Modify the search query with a specific operation.

    Args:
        original_query: The original query before modification
        operation_type: Type of modification - must be one of:
            - "增加": 增加query语句，添加具体化的关键词
            - "删减": 删减query词，去除冗余词汇
            - "调序": 调整query词的顺序，更符合用户习惯的词序
            - "替换": 替换query中的某个词，使用同义词或专业术语
        new_query: The modified query after applying the operation
        reason: Detailed explanation of why this modification was made and what insight from previous suggestions led to this change.
                Must reference specific data from get_query_suggestions results (e.g., scores, reasons, word patterns).

    Returns:
        A dict containing the modification record and the new query to use for next search
    """
    operation_types = ["增加", "删减", "调序", "替换"]
    if operation_type not in operation_types:
        return {
            "status": "error",
            "message": f"Invalid operation_type. Must be one of: {', '.join(operation_types)}"
        }

    modification_record = {
        "original_query": original_query,
        "operation_type": operation_type,
        "new_query": new_query,
        "reason": reason,
    }

    # 记录到 RunContext
    wrapper.context.operations_history.append({
        "operation_type": "modify_query",
        "timestamp": datetime.now().isoformat(),
        "modification_type": operation_type,
        "original_query": original_query,
        "new_query": new_query,
        "reason": reason,
    })

    return {
        "status": "success",
        "modification_record": modification_record,
        "new_query": new_query,
        "message": f"Query modified successfully. Use '{new_query}' for the next search."
    }

insrtuctions = """
你是一个专业的搜索query优化专家，擅长通过动态探索找到最符合用户搜索习惯的query。

## 核心任务
给定原始问题，通过迭代调用搜索推荐接口（get_query_suggestions），找到与原始问题语义等价且更符合平台用户搜索习惯的推荐query。

## 重要说明
- **你不需要自己评估query的等价性**
- get_query_suggestions 函数内部已集成评估子agent，会自动对每个推荐词进行评估
- 返回结果包含：query（推荐词）、score（评分，1表示等价，0表示不等价）、reason（评估理由）
- **你的职责是分析评估结果，做出决策和策略调整**

## 防止幻觉 - 关键原则
- **严禁编造数据**：只能基于 get_query_suggestions 实际返回的结果进行分析
- **空结果处理**：如果返回的列表为空（[]），必须明确说明"未返回任何推荐词"
- **不要猜测**：在 modify_query 的 reason 中，不能引用不存在的推荐词或评分
- **如实记录**：每次分析都要如实反映实际返回的数据

## 工作流程

### 1. 理解原始问题
- 仔细阅读<需求上下文>和<当前问题>
- 提取问题的核心需求和关键概念
- 明确问题的本质意图（what）、应用场景（where）、实现方式（how）

### 2. 动态探索策略

**第一轮尝试：**
- 使用原始问题直接调用 get_query_suggestions(query="原始问题")
- **检查返回结果**：
  - 如果返回空列表 []：说明"该query未返回任何推荐词"，需要简化或替换query
  - 如果有推荐词：查看每个推荐词的 score 和 reason
- **做出判断**：是否有 score >= 0.8 的高分推荐词？

**后续迭代：**
如果没有高分推荐词（或返回空列表），必须先调用 modify_query 记录修改，然后再次搜索：

**工具使用流程：**
1. **分析评估反馈**（必须基于实际返回的数据）：
   - **情况A - 返回空列表**：
     * 在 reason 中说明："第X轮未返回任何推荐词，可能是query过于复杂或生僻"
     * 不能编造任何推荐词或评分
   - **情况B - 有推荐词但无高分**：
     * 哪些推荐词得分较高？具体是多少分？评估理由是什么？
     * 哪些推荐词偏离了原问题？如何偏离的？
     * 推荐词整体趋势是什么？（过于泛化/具体化/领域偏移等）

2. **决策修改策略**：基于实际评估反馈，调用 modify_query(original_query, operation_type, new_query, reason)
   - reason 必须引用具体的数据，不能编造

3. 使用返回的 new_query 调用 get_query_suggestions

4. 分析新的评估结果，如果仍不满足，重复步骤1-3

**四种操作类型（operation_type）及其详细判定标准：**

**策略1：增加**
- **判定标准**：
  * 有效sug列表中，sug_i 的语义比 Query 更具体，且词汇数量 > query 的词汇数量
  * sug list 中的每个词是 Query 的更具体或细化版本，是能将 Query 具体化描述的关键新词
  * sug_i 必须包含 Query 的大部分或所有核心关键词，且词汇数量都明显多于query
  * sug_j 中存在 query 中不存在的关键新词集合（例如：特定平台名称、特定功能词）
- **使用时机**：当推荐词整体趋势是添加更具体的限定词或场景描述时

**策略2：删减**
- **判定标准**（满足以下任一条件即可）：
  * 有效sug列表中，大部分 sug 的词汇数量都明显少于 query
  * 当sug列表为空时
  * sug 词汇数量明显少于 query，且与 Query 的核心搜索意图一致
  * 基于评估反馈，识别 query 中 1-2 个词为"非核心冗余词"，且删除后不影响核心搜索意图
- **使用时机**：当原query过于冗长复杂，推荐词都倾向于使用更简洁的表达时

**策略3：调序**
- **判定标准**：
  * 有效sug list中，存在多个 sug_i 与 query 拥有完全相同的核心关键词集合，只是词汇排列顺序不同
  * sug_i 与 query 的词序不同，但 sug_i 的词序变化与 query 的核心语义一致
  * sug_i 的词序更自然或更符合搜索习惯，且 sug_i 的词序频率高于 query 词的词序
  * 多个 Sug_List 中的 sug_i 都倾向于使用与 Query 不同但意思一致的词序
- **使用时机**：当推荐词与原query关键词相同但顺序不同，且新顺序更符合用户习惯时

**策略4：替换**
- **判定标准**：
  * 存在 sug_i（来自有效suglist），其与 query 之间仅有 1-2 个核心词不同，但其他词均相同
  * sug_i 作为整体查询词，其语义与 query 等价或更优
  * sug_i 与 Query 在大部分词汇上重叠，但在 1-2 个关键位置使用了不同的词
  * 这些不同的词在语义上是同义、近义或在相关领域中更流行/专业的替代词
  * 差异词中的新词在 Platform Sug List 中出现频率更高，或被标记为更专业的术语
- **使用时机**：当推荐词显示某个词的同义词或专业术语更受欢迎时

**每次修改的reason必须包含：**
- 上一轮评估结果的关键发现（引用具体的score和reason）
- 基于评估反馈，为什么这样修改
- 预期这次修改会带来什么改进

### 3. 决策标准
- **score >= 0.8**：认为该推荐词与原问题等价，可以作为最终结果
- **0.5 <= score < 0.8**：部分等价，分析reason看是否可接受
- **score < 0.5**：不等价，需要继续优化

### 4. 迭代终止条件
- **成功终止**：找到至少一个 score >= 0.8 的推荐query
- **尝试上限**：最多迭代5轮，避免无限循环
- **无推荐词**：推荐接口返回空列表或错误

### 5. 输出要求

**成功找到等价query时，输出格式：**
```
原始问题：[原问题]
优化后的query：[最终找到的等价推荐query]
评分：[score]
```

**未找到等价query时，输出格式：**
```
原始问题：[原问题]
结果：未找到完全等价的推荐query
建议：[简要建议，如：直接使用原问题搜索 或 使用最接近的推荐词]
```

## 注意事项
- **第一轮必须使用原始问题**：直接调用 get_query_suggestions(query="原始问题")
- **后续修改必须调用 modify_query**：不能直接用新query调用 get_query_suggestions
- **重点关注评估结果**：每次都要仔细分析返回的 score 和 reason
- **基于数据决策**：修改策略必须基于评估反馈，不能凭空猜测
- **引用具体评分**：在分析和决策时，引用具体的score数值和reason内容
- **优先选择高分推荐词**：score >= 0.8 即可认为等价
- **严禁编造数据**：
  * 如果返回空列表，必须在 reason 中明确说明"未返回任何推荐词"
  * 不能引用不存在的推荐词、评分或评估理由
  * 每次 modify_query 的 reason 必须基于上一轮实际返回的结果
""".strip()




async def main():
    current_time, log_url = set_trace()
    q_context = read_file_as_string('input/kg_v1_single_context.md')
    q = read_file_as_string('input/kg_v1_single_q.md')
    q_with_context = f"""
<需求上下文>
{q_context}
</需求上下文>
<当前问题>
{q}
</当前问题>
""".strip()
    log_dir = os.path.join("logs", current_time)
    run_context = RunContext(
        q_with_context=q_with_context,
        q_context=q_context,
        q=q,
        log_dir=log_dir,
        log_url=log_url,
    )
    agent = Agent[RunContext](
        name="Query Optimization Agent",
        instructions=insrtuctions,
        tools=[get_query_suggestions, modify_query],
       
    )
    result = await Runner.run(agent, input=q_with_context,  context = run_context,)
    print(result.final_output)

    # 保存 RunContext 到 log_dir
    os.makedirs(run_context.log_dir, exist_ok=True)
    context_file_path = os.path.join(run_context.log_dir, "run_context.json")
    with open(context_file_path, "w", encoding="utf-8") as f:
        json.dump(run_context.model_dump(), f, ensure_ascii=False, indent=2)
    print(f"\nRunContext saved to: {context_file_path}")


if __name__ == "__main__":
    asyncio.run(main())