kg_agent/sug_v6_1_2_5.py

import asyncio
import json
import os
import sys
import argparse
from datetime import datetime
from typing import Literal

from agents import Agent, Runner
from lib.my_trace import set_trace
from pydantic import BaseModel, Field

from lib.utils import read_file_as_string
from lib.client import get_model
MODEL_NAME = "google/gemini-2.5-flash"
from script.search_recommendations.xiaohongshu_search_recommendations import XiaohongshuSearchRecommendations
from script.search.xiaohongshu_search import XiaohongshuSearch


# ============================================================================
# 数据模型
# ============================================================================

class QueryState(BaseModel):
    """Query状态跟踪"""
    query: str
    level: int  # 当前所在层级
    no_suggestion_rounds: int = 0  # 连续没有suggestion的轮数
    relevance_score: float = 0.0  # 与原始需求的相关度
    parent_query: str | None = None  # 父query
    strategy: str | None = None  # 生成策略：direct_sug, rewrite, add_word


class WordLibrary(BaseModel):
    """动态分词库"""
    words: set[str] = Field(default_factory=set)

    def add_word(self, word: str):
        """添加单词到分词库"""
        if word and word.strip():
            self.words.add(word.strip())

    def add_words(self, words: list[str]):
        """批量添加单词"""
        for word in words:
            self.add_word(word)

    def get_unused_word(self, current_query: str) -> str | None:
        """获取一个当前query中没有的词"""
        for word in self.words:
            if word not in current_query:
                return word
        return None

    def model_dump(self):
        """序列化为dict"""
        return {"words": list(self.words)}


class RunContext(BaseModel):
    """运行上下文"""
    version: str
    input_files: dict[str, str]
    q_with_context: str
    q_context: str
    q: str
    log_url: str
    log_dir: str

    # 新增字段
    word_library: dict = Field(default_factory=dict)  # 使用dict存储，因为set不能直接序列化
    query_states: list[dict] = Field(default_factory=list)
    steps: list[dict] = Field(default_factory=list)

    # 最终结果
    satisfied_notes: list[dict] = Field(default_factory=list)
    final_output: str | None = None


# ============================================================================
# Agent 定义
# ============================================================================

# Agent 1: 分词专家
class WordSegmentation(BaseModel):
    """分词结果"""
    words: list[str] = Field(..., description="分词结果列表")
    reasoning: str = Field(..., description="分词理由")

word_segmentation_instructions = """
你是分词专家。给定一个query，将其拆分成有意义的最小单元。

## 分词原则
1. 保留有搜索意义的词汇
2. 拆分成独立的概念
3. 保留专业术语的完整性
4. 去除虚词（的、吗、呢等）

## 输出要求
返回分词列表和分词理由。
""".strip()

word_segmenter = Agent[None](
    name="分词专家",
    instructions=word_segmentation_instructions,
    model=get_model(MODEL_NAME),
    output_type=WordSegmentation,
)


# Agent 2: Query相关度评估专家
class RelevanceEvaluation(BaseModel):
    """相关度评估"""
    relevance_score: float = Field(..., description="相关性分数 0-1")
    is_improved: bool = Field(..., description="是否比之前更好")
    reason: str = Field(..., description="评估理由")

relevance_evaluation_instructions = """
你是Query相关度评估专家。

## 任务
评估当前query与原始需求的匹配程度。

## 评估标准
- 主题相关性
- 要素覆盖度
- 意图匹配度

## 输出
- relevance_score: 0-1的相关性分数
- is_improved: 如果提供了previous_score，判断是否有提升
- reason: 详细理由
""".strip()

relevance_evaluator = Agent[None](
    name="Query相关度评估专家",
    instructions=relevance_evaluation_instructions,
    model=get_model(MODEL_NAME),
    output_type=RelevanceEvaluation,
)


# Agent 3: Query改写专家
class QueryRewrite(BaseModel):
    """Query改写结果"""
    rewritten_query: str = Field(..., description="改写后的query")
    rewrite_type: str = Field(..., description="改写类型：abstract或synonym")
    reasoning: str = Field(..., description="改写理由")

query_rewrite_instructions = """
你是Query改写专家。

## 改写策略
1. **向上抽象**：将具体概念泛化到更高层次
   - 例：iPhone 13 → 智能手机
2. **同义改写**：使用同义词或相关表达
   - 例：购买 → 入手、获取

## 输出要求
返回改写后的query、改写类型和理由。
""".strip()

query_rewriter = Agent[None](
    name="Query改写专家",
    instructions=query_rewrite_instructions,
    model=get_model(MODEL_NAME),
    output_type=QueryRewrite,
)


# Agent 4: 加词位置评估专家
class WordInsertion(BaseModel):
    """加词结果"""
    new_query: str = Field(..., description="加词后的新query")
    insertion_position: str = Field(..., description="插入位置描述")
    reasoning: str = Field(..., description="插入理由")

word_insertion_instructions = """
你是加词位置评估专家。

## 任务
将新词加到当前query的最合适位置，保持语义通顺。

## 原则
1. 保持语法正确
2. 语义连贯
3. 符合搜索习惯

## 输出
返回新query、插入位置描述和理由。
""".strip()

word_inserter = Agent[None](
    name="加词位置评估专家",
    instructions=word_insertion_instructions,
    model=get_model(MODEL_NAME),
    output_type=WordInsertion,
)


# Agent 5: Result匹配度评估专家
class ResultEvaluation(BaseModel):
    """Result评估结果"""
    match_level: str = Field(..., description="匹配等级：satisfied, partial, unsatisfied")
    relevance_score: float = Field(..., description="相关性分数 0-1")
    missing_aspects: list[str] = Field(default_factory=list, description="缺失的方面")
    reason: str = Field(..., description="评估理由")

result_evaluation_instructions = """
你是Result匹配度评估专家。

## 任务
评估搜索结果（帖子）与原始需求的匹配程度。

## 评估等级
1. **satisfied**: 完全满足需求
2. **partial**: 部分满足，但有缺失
3. **unsatisfied**: 基本不满足

## 输出要求
- match_level: 匹配等级
- relevance_score: 相关性分数
- missing_aspects: 如果是partial，列出缺失的方面
- reason: 详细理由
""".strip()

result_evaluator = Agent[None](
    name="Result匹配度评估专家",
    instructions=result_evaluation_instructions,
    model=get_model(MODEL_NAME),
    output_type=ResultEvaluation,
)


# Agent 6: Query改造专家（基于缺失部分）
class QueryImprovement(BaseModel):
    """Query改造结果"""
    improved_query: str = Field(..., description="改造后的query")
    added_aspects: list[str] = Field(..., description="添加的方面")
    reasoning: str = Field(..., description="改造理由")

query_improvement_instructions = """
你是Query改造专家。

## 任务
根据搜索结果的缺失部分，改造query使其包含这些内容。

## 原则
1. 针对性补充缺失方面
2. 保持query简洁
3. 符合搜索习惯

## 输出
返回改造后的query、添加的方面和理由。
""".strip()

query_improver = Agent[None](
    name="Query改造专家",
    instructions=query_improvement_instructions,
    model=get_model(MODEL_NAME),
    output_type=QueryImprovement,
)


# Agent 7: 关键词提取专家
class KeywordExtraction(BaseModel):
    """关键词提取结果"""
    keywords: list[str] = Field(..., description="提取的关键词列表")
    reasoning: str = Field(..., description="提取理由")

keyword_extraction_instructions = """
你是关键词提取专家。

## 任务
从帖子标题和描述中提取核心关键词。

## 提取原则
1. 提取有搜索价值的词汇
2. 去除虚词和通用词
3. 保留专业术语
4. 提取3-10个关键词

## 输出
返回关键词列表和提取理由。
""".strip()

keyword_extractor = Agent[None](
    name="关键词提取专家",
    instructions=keyword_extraction_instructions,
    model=get_model(MODEL_NAME),
    output_type=KeywordExtraction,
)


# ============================================================================
# 辅助函数
# ============================================================================

def add_step(context: RunContext, step_name: str, step_type: str, data: dict):
    """添加步骤记录"""
    step = {
        "step_number": len(context.steps) + 1,
        "step_name": step_name,
        "step_type": step_type,
        "timestamp": datetime.now().isoformat(),
        "data": data
    }
    context.steps.append(step)
    return step


def process_note_data(note: dict) -> dict:
    """处理搜索接口返回的帖子数据"""
    note_card = note.get("note_card", {})
    image_list = note_card.get("image_list", [])
    interact_info = note_card.get("interact_info", {})
    user_info = note_card.get("user", {})

    return {
        "note_id": note.get("id", ""),
        "title": note_card.get("display_title", ""),
        "desc": note_card.get("desc", ""),
        "image_list": image_list,
        "interact_info": {
            "liked_count": interact_info.get("liked_count", 0),
            "collected_count": interact_info.get("collected_count", 0),
            "comment_count": interact_info.get("comment_count", 0),
            "shared_count": interact_info.get("shared_count", 0)
        },
        "user": {
            "nickname": user_info.get("nickname", ""),
            "user_id": user_info.get("user_id", "")
        },
        "type": note_card.get("type", "normal"),
        "note_url": f"https://www.xiaohongshu.com/explore/{note.get('id', '')}"
    }


# ============================================================================
# 核心流程函数
# ============================================================================

async def initialize_word_library(original_query: str, context: RunContext) -> WordLibrary:
    """初始化分词库"""
    print("\n[初始化] 创建分词库...")

    # 使用Agent进行分词
    result = await Runner.run(word_segmenter, original_query)
    segmentation: WordSegmentation = result.final_output

    word_lib = WordLibrary()
    word_lib.add_words(segmentation.words)

    print(f"初始分词库: {list(word_lib.words)}")
    print(f"分词理由: {segmentation.reasoning}")

    # 保存到context
    context.word_library = word_lib.model_dump()

    add_step(context, "初始化分词库", "word_library_init", {
        "original_query": original_query,
        "words": list(word_lib.words),
        "reasoning": segmentation.reasoning
    })

    return word_lib


async def evaluate_query_relevance(
    query: str,
    original_need: str,
    previous_score: float | None = None,
    context: RunContext = None
) -> RelevanceEvaluation:
    """评估query与原始需求的相关度"""

    eval_input = f"""
<原始需求>
{original_need}
</原始需求>

<当前Query>
{query}
</当前Query>

{"<之前的相关度分数>" + str(previous_score) + "</之前的相关度分数>" if previous_score is not None else ""}

请评估当前query与原始需求的相关度。
"""

    result = await Runner.run(relevance_evaluator, eval_input)
    evaluation: RelevanceEvaluation = result.final_output

    return evaluation


async def process_suggestions(
    query: str,
    query_state: QueryState,
    original_need: str,
    word_lib: WordLibrary,
    context: RunContext,
    xiaohongshu_api: XiaohongshuSearchRecommendations
) -> list[QueryState]:
    """处理suggestion分支，返回新的query states"""

    print(f"\n  [Suggestion分支] 处理query: {query}")

    # 1. 获取suggestions
    suggestions = xiaohongshu_api.get_recommendations(keyword=query)

    if not suggestions or len(suggestions) == 0:
        print(f"    → 没有获取到suggestion")
        query_state.no_suggestion_rounds += 1

        # 记录步骤
        add_step(context, f"Suggestion分支 - {query}", "suggestion_branch", {
            "query": query,
            "query_level": query_state.level,
            "suggestions_count": 0,
            "no_suggestion_rounds": query_state.no_suggestion_rounds,
            "new_queries_generated": 0
        })

        return []

    print(f"    → 获取到 {len(suggestions)} 个suggestions")
    query_state.no_suggestion_rounds = 0  # 重置计数

    # 2. 评估每个suggestion
    new_queries = []
    suggestion_evaluations = []

    for sug in suggestions[:5]:  # 限制处理数量
        # 评估sug的相关度
        sug_eval = await evaluate_query_relevance(sug, original_need, query_state.relevance_score, context)

        sug_eval_record = {
            "suggestion": sug,
            "relevance_score": sug_eval.relevance_score,
            "is_improved": sug_eval.is_improved,
            "reason": sug_eval.reason
        }
        suggestion_evaluations.append(sug_eval_record)

        # 判断是否比当前query更好
        if sug_eval.is_improved and sug_eval.relevance_score > query_state.relevance_score:
            print(f"      ✓ {sug} (分数: {sug_eval.relevance_score:.2f}, 提升: {sug_eval.is_improved})")

            # 创建新的query state（直接使用suggestion）
            new_state = QueryState(
                query=sug,
                level=query_state.level + 1,
                relevance_score=sug_eval.relevance_score,
                parent_query=query,
                strategy="direct_sug"
            )
            new_queries.append(new_state)
        else:
            print(f"      ✗ {sug} (分数: {sug_eval.relevance_score:.2f}, 未提升)")

    # 3. 改写策略（向上抽象或同义改写）
    if len(new_queries) < 3:  # 如果直接使用sug的数量不够，尝试改写
        rewrite_input = f"""
<当前Query>
{query}
</当前Query>

<改写要求>
类型: abstract (向上抽象)
</改写要求>

请改写这个query。
"""
        result = await Runner.run(query_rewriter, rewrite_input)
        rewrite: QueryRewrite = result.final_output

        # 评估改写后的query
        rewrite_eval = await evaluate_query_relevance(rewrite.rewritten_query, original_need, query_state.relevance_score, context)

        if rewrite_eval.is_improved:
            print(f"      ✓ 改写: {rewrite.rewritten_query} (分数: {rewrite_eval.relevance_score:.2f})")
            new_state = QueryState(
                query=rewrite.rewritten_query,
                level=query_state.level + 1,
                relevance_score=rewrite_eval.relevance_score,
                parent_query=query,
                strategy="rewrite"
            )
            new_queries.append(new_state)

    # 4. 加词策略
    unused_word = word_lib.get_unused_word(query)
    if unused_word and len(new_queries) < 5:
        insertion_input = f"""
<当前Query>
{query}
</当前Query>

<要添加的词>
{unused_word}
</要添加的词>

请将这个词加到query的最合适位置。
"""
        result = await Runner.run(word_inserter, insertion_input)
        insertion: WordInsertion = result.final_output

        # 评估加词后的query
        insertion_eval = await evaluate_query_relevance(insertion.new_query, original_need, query_state.relevance_score, context)

        if insertion_eval.is_improved:
            print(f"      ✓ 加词: {insertion.new_query} (分数: {insertion_eval.relevance_score:.2f})")
            new_state = QueryState(
                query=insertion.new_query,
                level=query_state.level + 1,
                relevance_score=insertion_eval.relevance_score,
                parent_query=query,
                strategy="add_word"
            )
            new_queries.append(new_state)

    # 记录完整的suggestion分支处理结果
    add_step(context, f"Suggestion分支 - {query}", "suggestion_branch", {
        "query": query,
        "query_level": query_state.level,
        "query_relevance": query_state.relevance_score,
        "suggestions_count": len(suggestions),
        "suggestions_evaluated": len(suggestion_evaluations),
        "suggestion_evaluations": suggestion_evaluations[:10],  # 只保存前10个
        "new_queries_generated": len(new_queries),
        "new_queries": [{"query": nq.query, "score": nq.relevance_score, "strategy": nq.strategy} for nq in new_queries],
        "no_suggestion_rounds": query_state.no_suggestion_rounds
    })

    return new_queries


async def process_search_results(
    query: str,
    query_state: QueryState,
    original_need: str,
    word_lib: WordLibrary,
    context: RunContext,
    xiaohongshu_search: XiaohongshuSearch,
    relevance_threshold: float = 0.6
) -> tuple[list[dict], list[QueryState]]:
    """
    处理搜索结果分支
    返回: (满足需求的notes, 需要继续迭代的新queries)
    """

    print(f"\n  [Result分支] 搜索query: {query}")

    # 1. 判断query相关度是否达到门槛
    if query_state.relevance_score < relevance_threshold:
        print(f"    ✗ 相关度 {query_state.relevance_score:.2f} 低于门槛 {relevance_threshold}，跳过搜索")
        return [], []

    print(f"    ✓ 相关度 {query_state.relevance_score:.2f} 达到门槛，执行搜索")

    # 2. 执行搜索
    try:
        search_result = xiaohongshu_search.search(keyword=query)
        result_str = search_result.get("result", "{}")
        if isinstance(result_str, str):
            result_data = json.loads(result_str)
        else:
            result_data = result_str

        notes = result_data.get("data", {}).get("data", [])
        print(f"    → 搜索到 {len(notes)} 个帖子")

    except Exception as e:
        print(f"    ✗ 搜索失败: {e}")
        return [], []

    if not notes:
        return [], []

    # 3. 评估每个帖子
    satisfied_notes = []
    partial_notes = []

    for note in notes[:10]:  # 限制评估数量
        note_data = process_note_data(note)
        title = note_data["title"] or ""
        desc = note_data["desc"] or ""

        # 跳过空标题和描述的帖子
        if not title and not desc:
            continue

        # 评估帖子
        eval_input = f"""
<原始需求>
{original_need}
</原始需求>

<帖子>
标题: {title}
描述: {desc}
</帖子>

请评估这个帖子与原始需求的匹配程度。
"""
        result = await Runner.run(result_evaluator, eval_input)
        evaluation: ResultEvaluation = result.final_output

        note_data["evaluation"] = {
            "match_level": evaluation.match_level,
            "relevance_score": evaluation.relevance_score,
            "missing_aspects": evaluation.missing_aspects,
            "reason": evaluation.reason
        }

        if evaluation.match_level == "satisfied":
            satisfied_notes.append(note_data)
            print(f"      ✓ 满足: {title[:30] if len(title) > 30 else title}... (分数: {evaluation.relevance_score:.2f})")
        elif evaluation.match_level == "partial":
            partial_notes.append(note_data)
            print(f"      ~ 部分: {title[:30] if len(title) > 30 else title}... (缺失: {', '.join(evaluation.missing_aspects[:2])})")

    # 4. 处理满足的帖子：提取关键词并扩充分词库
    new_queries = []

    if satisfied_notes:
        print(f"\n    从 {len(satisfied_notes)} 个满足的帖子中提取关键词...")
        for note in satisfied_notes[:3]:  # 限制处理数量
            extract_input = f"""
<帖子>
标题: {note['title']}
描述: {note['desc']}
</帖子>

请提取核心关键词。
"""
            result = await Runner.run(keyword_extractor, extract_input)
            extraction: KeywordExtraction = result.final_output

            # 添加新词到分词库
            for keyword in extraction.keywords:
                if keyword not in word_lib.words:
                    word_lib.add_word(keyword)
                    print(f"      + 新词入库: {keyword}")

    # 5. 处理部分匹配的帖子：改造query
    if partial_notes and len(satisfied_notes) < 5:  # 如果满足的不够，基于部分匹配改进
        print(f"\n    基于 {len(partial_notes)} 个部分匹配帖子改造query...")
        # 收集所有缺失方面
        all_missing = []
        for note in partial_notes:
            all_missing.extend(note["evaluation"]["missing_aspects"])

        if all_missing:
            improvement_input = f"""
<当前Query>
{query}
</当前Query>

<缺失的方面>
{', '.join(set(all_missing[:5]))}
</缺失的方面>

请改造query使其包含这些缺失的内容。
"""
            result = await Runner.run(query_improver, improvement_input)
            improvement: QueryImprovement = result.final_output

            # 评估改进后的query
            improved_eval = await evaluate_query_relevance(improvement.improved_query, original_need, query_state.relevance_score, context)

            if improved_eval.is_improved:
                print(f"      ✓ 改进: {improvement.improved_query} (添加: {', '.join(improvement.added_aspects[:2])})")
                new_state = QueryState(
                    query=improvement.improved_query,
                    level=query_state.level + 1,
                    relevance_score=improved_eval.relevance_score,
                    parent_query=query,
                    strategy="improve_from_partial"
                )
                new_queries.append(new_state)

    # 记录完整的result分支处理结果
    add_step(context, f"Result分支 - {query}", "result_branch", {
        "query": query,
        "query_level": query_state.level,
        "query_relevance": query_state.relevance_score,
        "relevance_threshold": relevance_threshold,
        "passed_threshold": query_state.relevance_score >= relevance_threshold,
        "notes_count": len(notes) if 'notes' in locals() else 0,
        "satisfied_count": len(satisfied_notes),
        "partial_count": len(partial_notes),
        "satisfied_notes": [
            {
                "note_id": note["note_id"],
                "title": note["title"],
                "score": note["evaluation"]["relevance_score"],
                "match_level": note["evaluation"]["match_level"]
            }
            for note in satisfied_notes[:10]  # 只保存前10个
        ],
        "new_queries_generated": len(new_queries),
        "new_queries": [{"query": nq.query, "score": nq.relevance_score, "strategy": nq.strategy} for nq in new_queries],
        "word_library_expanded": len(new_queries) > 0  # 是否扩充了分词库
    })

    return satisfied_notes, new_queries


async def iterative_search_loop(
    context: RunContext,
    max_iterations: int = 20,
    max_concurrent_queries: int = 5,
    relevance_threshold: float = 0.6
) -> list[dict]:
    """
    主循环：迭代搜索

    Args:
        context: 运行上下文
        max_iterations: 最大迭代次数
        max_concurrent_queries: 最大并发query数量
        relevance_threshold: 相关度门槛

    Returns:
        满足需求的帖子列表
    """

    print(f"\n{'='*60}")
    print(f"开始迭代搜索循环")
    print(f"{'='*60}")

    # 1. 初始化分词库
    word_lib = await initialize_word_library(context.q, context)

    # 2. 初始化query队列 - 智能选择最相关的词
    all_words = list(word_lib.words)
    query_queue = []

    print(f"\n评估所有初始分词的相关度...")
    word_scores = []

    for word in all_words:
        # 评估每个词的相关度
        eval_result = await evaluate_query_relevance(word, context.q, None, context)
        word_scores.append({
            'word': word,
            'score': eval_result.relevance_score,
            'eval': eval_result
        })
        print(f"  {word}: {eval_result.relevance_score:.2f}")

    # 按相关度排序，选择top 3
    word_scores.sort(key=lambda x: x['score'], reverse=True)
    selected_words = word_scores[:3]

    for item in selected_words:
        query_queue.append(QueryState(
            query=item['word'],
            level=1,
            relevance_score=item['score'],
            strategy="initial"
        ))

    print(f"\n初始query队列（按相关度选择）: {[(q.query, f'{q.relevance_score:.2f}') for q in query_queue]}")

    # 3. API实例
    xiaohongshu_api = XiaohongshuSearchRecommendations()
    xiaohongshu_search = XiaohongshuSearch()

    # 4. 主循环
    all_satisfied_notes = []
    iteration = 0

    while query_queue and iteration < max_iterations:
        iteration += 1
        print(f"\n{'='*60}")
        print(f"迭代 {iteration}: 队列中有 {len(query_queue)} 个query")
        print(f"{'='*60}")

        # 限制并发数量
        current_batch = query_queue[:max_concurrent_queries]
        query_queue = query_queue[max_concurrent_queries:]

        # 记录本轮处理的queries
        add_step(context, f"迭代 {iteration}", "iteration", {
            "iteration": iteration,
            "queue_size": len(query_queue) + len(current_batch),
            "processing_queries": [q.query for q in current_batch]
        })

        new_queries_from_sug = []
        new_queries_from_result = []

        # 处理每个query
        for query_state in current_batch:
            print(f"\n处理Query [{query_state.level}]: {query_state.query} (分数: {query_state.relevance_score:.2f})")

            # 检查终止条件
            if query_state.no_suggestion_rounds >= 2:
                print(f"  ✗ 连续2轮无suggestion，终止该分支")
                continue

            # 并行处理两个分支
            sug_task = process_suggestions(
                query_state.query, query_state, context.q, word_lib, context, xiaohongshu_api
            )
            result_task = process_search_results(
                query_state.query, query_state, context.q, word_lib, context,
                xiaohongshu_search, relevance_threshold
            )

            # 等待两个分支完成
            sug_queries, (satisfied_notes, result_queries) = await asyncio.gather(
                sug_task,
                result_task
            )

            new_queries_from_sug.extend(sug_queries)
            new_queries_from_result.extend(result_queries)
            all_satisfied_notes.extend(satisfied_notes)

        # 更新队列
        all_new_queries = new_queries_from_sug + new_queries_from_result
        query_queue.extend(all_new_queries)

        # 去重（基于query文本）
        seen = set()
        unique_queue = []
        for q in query_queue:
            if q.query not in seen:
                seen.add(q.query)
                unique_queue.append(q)
        query_queue = unique_queue

        # 按相关度排序
        query_queue.sort(key=lambda x: x.relevance_score, reverse=True)

        print(f"\n本轮结果:")
        print(f"  新增满足帖子: {len(satisfied_notes)}")
        print(f"  累计满足帖子: {len(all_satisfied_notes)}")
        print(f"  新增queries: {len(all_new_queries)}")
        print(f"  队列剩余: {len(query_queue)}")

        # 更新分词库到context
        context.word_library = word_lib.model_dump()

        # 如果满足条件的帖子足够多，可以提前结束
        if len(all_satisfied_notes) >= 20:
            print(f"\n已找到足够的满足帖子 ({len(all_satisfied_notes)}个)，提前结束")
            break

    print(f"\n{'='*60}")
    print(f"迭代搜索完成")
    print(f"  总迭代次数: {iteration}")
    print(f"  最终满足帖子数: {len(all_satisfied_notes)}")
    print(f"  最终分词库大小: {len(word_lib.words)}")
    print(f"{'='*60}")

    # 保存最终结果
    add_step(context, "迭代搜索完成", "loop_complete", {
        "total_iterations": iteration,
        "total_satisfied_notes": len(all_satisfied_notes),
        "final_word_library_size": len(word_lib.words),
        "final_word_library": list(word_lib.words)
    })

    return all_satisfied_notes


# ============================================================================
# 主函数
# ============================================================================

async def main(input_dir: str, max_iterations: int = 20, visualize: bool = False):
    """主函数"""
    current_time, log_url = set_trace()

    # 读取输入
    input_context_file = os.path.join(input_dir, 'context.md')
    input_q_file = os.path.join(input_dir, 'q.md')

    q_context = read_file_as_string(input_context_file)
    q = read_file_as_string(input_q_file)
    q_with_context = f"""
<需求上下文>
{q_context}
</需求上下文>
<当前问题>
{q}
</当前问题>
""".strip()

    # 版本信息
    version = os.path.basename(__file__)
    version_name = os.path.splitext(version)[0]

    # 日志目录
    log_dir = os.path.join(input_dir, "output", version_name, current_time)

    # 创建运行上下文
    run_context = RunContext(
        version=version,
        input_files={
            "input_dir": input_dir,
            "context_file": input_context_file,
            "q_file": input_q_file,
        },
        q_with_context=q_with_context,
        q_context=q_context,
        q=q,
        log_dir=log_dir,
        log_url=log_url,
    )

    # 执行迭代搜索
    satisfied_notes = await iterative_search_loop(
        run_context,
        max_iterations=max_iterations,
        max_concurrent_queries=3,
        relevance_threshold=0.6
    )

    # 保存结果
    run_context.satisfied_notes = satisfied_notes

    # 格式化输出
    output = f"原始问题：{run_context.q}\n"
    output += f"找到满足需求的帖子：{len(satisfied_notes)} 个\n"
    output += f"分词库大小：{len(run_context.word_library.get('words', []))} 个词\n"
    output += "\n" + "="*60 + "\n"

    if satisfied_notes:
        output += "【满足需求的帖子】\n\n"
        for idx, note in enumerate(satisfied_notes[:10], 1):
            output += f"{idx}. {note['title']}\n"
            output += f"   相关度: {note['evaluation']['relevance_score']:.2f}\n"
            output += f"   URL: {note['note_url']}\n\n"
    else:
        output += "未找到满足需求的帖子\n"

    run_context.final_output = output

    print(f"\n{'='*60}")
    print("最终结果")
    print(f"{'='*60}")
    print(output)

    # 保存日志
    os.makedirs(run_context.log_dir, exist_ok=True)

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

    steps_file_path = os.path.join(run_context.log_dir, "steps.json")
    with open(steps_file_path, "w", encoding="utf-8") as f:
        json.dump(run_context.steps, f, ensure_ascii=False, indent=2)
    print(f"Steps log saved to: {steps_file_path}")

    # 可视化
    if visualize:
        import subprocess
        output_html = os.path.join(run_context.log_dir, "visualization.html")
        print(f"\n🎨 生成可视化HTML...")
        result = subprocess.run([
            "python", "sug_v6_1_2_3.visualize.py",
            steps_file_path,
            "-o", output_html
        ])
        if result.returncode == 0:
            print(f"✅ 可视化已生成: {output_html}")
        else:
            print(f"❌ 可视化生成失败")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="搜索query优化工具 - v6.1.2.5 迭代循环版")
    parser.add_argument(
        "--input-dir",
        type=str,
        default="input/简单扣图",
        help="输入目录路径，默认: input/简单扣图"
    )
    parser.add_argument(
        "--max-iterations",
        type=int,
        default=20,
        help="最大迭代次数，默认: 20"
    )
    parser.add_argument(
        "--visualize",
        action="store_true",
        default=False,
        help="运行完成后自动生成可视化HTML"
    )
    args = parser.parse_args()

    asyncio.run(main(args.input_dir, max_iterations=args.max_iterations, visualize=args.visualize))