import asyncio import json import os import argparse from datetime import datetime from itertools import combinations, permutations from agents import Agent, Runner from lib.my_trace import set_trace from typing import Literal from pydantic import BaseModel, Field from lib.utils import read_file_as_string from script.search_recommendations.xiaohongshu_search_recommendations import XiaohongshuSearchRecommendations # ============================================================================ # 并发控制配置 # ============================================================================ # API请求并发度(小红书接口) API_CONCURRENCY_LIMIT = 5 # 模型评估并发度(GPT评估) MODEL_CONCURRENCY_LIMIT = 10 class RunContext(BaseModel): version: str = Field(..., description="当前运行的脚本版本(文件名)") input_files: dict[str, str] = Field(..., description="输入文件路径映射") q_with_context: str q_context: str q: str log_url: str log_dir: str # 问题标注 question_annotation: str | None = Field(default=None, description="问题的标注结果(三层)") # 分词和组合 keywords: list[str] | None = Field(default=None, description="提取的关键词") query_combinations: dict[str, list[str]] = Field(default_factory=dict, description="各层级的query组合") # 探索结果 all_sug_queries: list[dict] = Field(default_factory=list, description="所有获取到的推荐词") # 评估结果 evaluation_results: list[dict] = Field(default_factory=list, description="所有推荐词的评估结果") optimization_result: dict | None = Field(default=None, description="最终优化结果对象") final_output: str | None = Field(default=None, description="最终输出结果(格式化文本)") # ============================================================================ # Agent 1: 问题标注专家 # ============================================================================ question_annotation_instructions = """ 你是搜索需求分析专家。给定问题(含需求背景),在原文上标注三层:本质、硬性、软性。 ## 判断标准 **[本质]** - 问题的核心意图 - 如何获取、教程、推荐、作品、测评等 **[硬]** - 客观事实性约束(可明确验证、非主观判断) - 能明确区分类别的:地域、时间、对象、工具、操作类型 - 特征:改变后得到完全不同类别的结果 **[软]** - 主观判断性修饰(因人而异、程度性的) - 需要主观评价的:质量、速度、美观、特色、程度 - 特征:改变后仍是同类结果,只是满足程度不同 ## 输出格式 词语[本质-描述]、词语[硬-描述]、词语[软-描述] ## 注意 - 只输出标注后的字符串 - 结合需求背景判断意图 """.strip() question_annotator = Agent[None]( name="问题标注专家", instructions=question_annotation_instructions, ) # ============================================================================ # Agent 2: 分词专家 # ============================================================================ segmentation_instructions = """ 你是中文分词专家。给定一个句子,将其分词。 ## 分词原则 1. 去掉标点符号 2. 拆分成最小的有意义单元 3. 去掉助词、语气词、助动词 4. 保留疑问词 5. 保留实词:名词、动词、形容词、副词 ## 输出要求 输出分词列表。 """.strip() class SegmentationResult(BaseModel): """分词结果""" words: list[str] = Field(..., description="分词列表") reasoning: str = Field(..., description="分词说明") segmenter = Agent[None]( name="分词专家", instructions=segmentation_instructions, output_type=SegmentationResult, ) # ============================================================================ # Agent 3: 评估专家(意图匹配 + 相关性评分) # ============================================================================ eval_instructions = """ 你是搜索query评估专家。给定原始问题、问题标注和推荐query,评估两个维度。 ## 输入信息 你会收到: 1. 原始问题:用户的原始表述 2. 问题标注:对原始问题的三层标注(本质、硬性、软性) 3. 推荐query:待评估的推荐词 ## 评估目标 用这个推荐query搜索,能否找到满足原始需求的内容? ## 两层评分 ### 1. intent_match(意图匹配)= true/false 推荐query的**使用意图**是否与原问题的**本质**一致? **核心:只关注[本质]标注** - 问题标注中的 `[本质-XXX]` 标记明确了用户的核心意图 - 判断推荐词是否能达成这个核心意图 **常见本质类型:** - 找方法/如何获取 → 推荐词应包含方法、途径、网站、渠道等 - 找教程 → 推荐词应是教程、教学相关 - 找资源/素材 → 推荐词应是资源、素材本身 - 找工具 → 推荐词应是工具推荐 - 看作品 → 推荐词应是作品展示 **评分:** - true = 推荐词的意图与 `[本质]` 一致 - false = 推荐词的意图与 `[本质]` 不一致 ### 2. relevance_score(相关性)= 0-1 连续分数 在意图匹配的前提下,推荐query在**主题、要素、属性**上与原问题的相关程度? **评估维度:** - 主题相关:核心主题是否匹配?(如:摄影、旅游、美食) - 要素覆盖:`[硬-XXX]` 标记的硬性约束保留了多少?(地域、时间、对象、工具等) - 属性匹配:`[软-XXX]` 标记的软性修饰保留了多少?(质量、速度、美观等) **评分参考:** - 0.9-1.0 = 几乎完美匹配,[硬]和[软]标注的要素都保留 - 0.7-0.8 = 高度相关,[硬]标注的要素都保留,[软]标注少数缺失 - 0.5-0.6 = 中度相关,[硬]标注的要素保留大部分,[软]标注多数缺失 - 0.3-0.4 = 低度相关,[硬]标注的要素部分缺失 - 0-0.2 = 基本不相关,[硬]标注的要素大量缺失 ## 评估策略 1. **先看[本质]判断 intent_match**:意图不匹配直接 false 2. **再看[硬][软]评估 relevance_score**:计算要素和属性的保留程度 ## 输出要求 - intent_match: true/false - relevance_score: 0-1 的浮点数 - reason: 详细的评估理由,需要说明: - 原问题的[本质]是什么,推荐词是否匹配这个本质 - [硬]约束哪些保留/缺失 - [软]修饰哪些保留/缺失 - 最终相关性分数的依据 """.strip() class RelevanceEvaluation(BaseModel): """评估反馈模型 - 意图匹配 + 相关性""" intent_match: bool = Field(..., description="意图是否匹配") relevance_score: float = Field(..., description="相关性分数 0-1,分数越高越相关") reason: str = Field(..., description="评估理由,需说明意图判断和相关性依据") evaluator = Agent[None]( name="评估专家", instructions=eval_instructions, output_type=RelevanceEvaluation, ) # ============================================================================ # 核心函数 # ============================================================================ async def annotate_question(q_with_context: str) -> str: """标注问题(三层)""" print("\n正在标注问题...") result = await Runner.run(question_annotator, q_with_context) annotation = str(result.final_output) print(f"问题标注完成:{annotation}") return annotation async def segment_text(q: str) -> SegmentationResult: """分词""" print("\n正在分词...") result = await Runner.run(segmenter, q) seg_result: SegmentationResult = result.final_output print(f"分词结果:{seg_result.words}") print(f"分词说明:{seg_result.reasoning}") return seg_result def generate_query_combinations(keywords: list[str], max_combination_size: int) -> dict[str, list[str]]: """ 生成query组合(考虑词的顺序) Args: keywords: 关键词列表 max_combination_size: 最大组合词数(N) Returns: { "1-word": [...], "2-word": [...], "3-word": [...], ... "N-word": [...] } """ result = {} for size in range(1, max_combination_size + 1): if size > len(keywords): break # 1-word组合:不需要考虑顺序 if size == 1: queries = keywords.copy() else: # 多词组合:先选择size个词(combinations),再排列(permutations) all_queries = [] combs = list(combinations(keywords, size)) for comb in combs: # 对每个组合生成所有排列 perms = list(permutations(comb)) for perm in perms: query = ''.join(perm) # 直接拼接,无空格 all_queries.append(query) # 去重(虽然理论上不会重复,但保险起见) queries = list(dict.fromkeys(all_queries)) result[f"{size}-word"] = queries print(f"\n{size}词组合:共 {len(queries)} 个") # 打印所有query,带序号 for i, q in enumerate(queries, 1): print(f" {i}. {q}") return result async def fetch_and_evaluate_streaming( queries: list[str], original_question: str, question_annotation: str, context: RunContext ) -> tuple[list[dict], list[dict]]: """ 流式处理:边获取推荐词边评估 Returns: (sug_results, evaluations) """ xiaohongshu_api = XiaohongshuSearchRecommendations() # 创建信号量 api_semaphore = asyncio.Semaphore(API_CONCURRENCY_LIMIT) model_semaphore = asyncio.Semaphore(MODEL_CONCURRENCY_LIMIT) # 结果收集 sug_results = [] all_evaluations = [] # 统计 total_queries = len(queries) completed_queries = 0 total_sugs = 0 completed_evals = 0 async def get_and_evaluate_single_query(query: str): nonlocal completed_queries, total_sugs, completed_evals # 步骤1:获取推荐词 async with api_semaphore: suggestions = xiaohongshu_api.get_recommendations(keyword=query) sug_count = len(suggestions) if suggestions else 0 completed_queries += 1 total_sugs += sug_count print(f" [{completed_queries}/{total_queries}] {query} → {sug_count} 个推荐词") sug_result = { "query": query, "suggestions": suggestions or [], "timestamp": datetime.now().isoformat() } sug_results.append(sug_result) # 步骤2:立即评估这些推荐词 if suggestions: eval_tasks = [] for sug in suggestions: eval_tasks.append(evaluate_single_sug_with_semaphore( query, sug, original_question, question_annotation, model_semaphore )) if eval_tasks: evals = await asyncio.gather(*eval_tasks) all_evaluations.extend(evals) completed_evals += len(evals) print(f" ↳ 已评估 {len(evals)} 个,累计评估 {completed_evals} 个") # 并发处理所有query await asyncio.gather(*[get_and_evaluate_single_query(q) for q in queries]) # 保存到context context.all_sug_queries = sug_results context.evaluation_results = all_evaluations print(f"\n总计:获取 {total_sugs} 个推荐词,完成 {completed_evals} 个评估") return sug_results, all_evaluations async def evaluate_single_sug_with_semaphore( source_query: str, sug_query: str, original_question: str, question_annotation: str, semaphore: asyncio.Semaphore ) -> dict: """带信号量的单个推荐词评估""" async with semaphore: eval_input = f""" <原始问题> {original_question} <问题标注(三层)> {question_annotation} <待评估的推荐query> {sug_query} 请评估该推荐query: 1. intent_match: 意图是否匹配(true/false) 2. relevance_score: 相关性分数(0-1) 3. reason: 详细的评估理由 评估时请参考问题标注中的[本质]、[硬]、[软]标记。 """ result = await Runner.run(evaluator, eval_input) evaluation: RelevanceEvaluation = result.final_output return { "source_query": source_query, "sug_query": sug_query, "intent_match": evaluation.intent_match, "relevance_score": evaluation.relevance_score, "reason": evaluation.reason, } async def fetch_suggestions_for_queries(queries: list[str], context: RunContext) -> list[dict]: """ 并发获取所有query的推荐词(带并发控制) Returns: [ { "query": "川西", "suggestions": ["川西旅游", "川西攻略", ...], "timestamp": "..." }, ... ] """ print(f"\n{'='*60}") print(f"获取推荐词:{len(queries)} 个query(并发度:{API_CONCURRENCY_LIMIT})") print(f"{'='*60}") xiaohongshu_api = XiaohongshuSearchRecommendations() # 创建信号量控制并发 semaphore = asyncio.Semaphore(API_CONCURRENCY_LIMIT) async def get_single_sug(query: str): async with semaphore: print(f" 查询: {query}") suggestions = xiaohongshu_api.get_recommendations(keyword=query) print(f" → {len(suggestions) if suggestions else 0} 个推荐词") return { "query": query, "suggestions": suggestions or [], "timestamp": datetime.now().isoformat() } results = await asyncio.gather(*[get_single_sug(q) for q in queries]) return results async def evaluate_all_suggestions( sug_results: list[dict], original_question: str, question_annotation: str, context: RunContext ) -> list[dict]: """ 评估所有推荐词(带并发控制) Args: sug_results: 所有query的推荐词结果 original_question: 原始问题 question_annotation: 问题标注(三层) Returns: [ { "source_query": "川西秋季", "sug_query": "川西秋季旅游", "intent_match": True, "relevance_score": 0.8, "reason": "..." }, ... ] """ print(f"\n{'='*60}") print(f"评估推荐词(并发度:{MODEL_CONCURRENCY_LIMIT})") print(f"{'='*60}") # 创建信号量控制并发 semaphore = asyncio.Semaphore(MODEL_CONCURRENCY_LIMIT) # 收集所有推荐词 all_evaluations = [] async def evaluate_single_sug(source_query: str, sug_query: str): async with semaphore: eval_input = f""" <原始问题> {original_question} <问题标注(三层)> {question_annotation} <待评估的推荐query> {sug_query} 请评估该推荐query: 1. intent_match: 意图是否匹配(true/false) 2. relevance_score: 相关性分数(0-1) 3. reason: 详细的评估理由 评估时请参考问题标注中的[本质]、[硬]、[软]标记。 """ result = await Runner.run(evaluator, eval_input) evaluation: RelevanceEvaluation = result.final_output return { "source_query": source_query, "sug_query": sug_query, "intent_match": evaluation.intent_match, "relevance_score": evaluation.relevance_score, "reason": evaluation.reason, } # 并发评估所有推荐词 tasks = [] for sug_result in sug_results: source_query = sug_result["query"] for sug in sug_result["suggestions"]: tasks.append(evaluate_single_sug(source_query, sug)) if tasks: print(f" 总共需要评估 {len(tasks)} 个推荐词...") all_evaluations = await asyncio.gather(*tasks) context.evaluation_results = all_evaluations return all_evaluations def find_qualified_queries(evaluations: list[dict], min_relevance_score: float = 0.7) -> list[dict]: """ 查找所有合格的query 筛选标准: 1. intent_match = True(必须满足) 2. relevance_score >= min_relevance_score 返回:按 relevance_score 降序排列 """ qualified = [ e for e in evaluations if e['intent_match'] is True and e['relevance_score'] >= min_relevance_score ] # 按relevance_score降序排列 return sorted(qualified, key=lambda x: x['relevance_score'], reverse=True) # ============================================================================ # 主流程 # ============================================================================ async def combinatorial_search(context: RunContext, max_combination_size: int = 1) -> dict: """ 组合式搜索流程(带问题标注) Args: context: 运行上下文 max_combination_size: 最大组合词数(N),默认1 返回格式: { "success": True/False, "results": [...], "message": "..." } """ # 步骤1:标注问题(三层) annotation = await annotate_question(context.q_with_context) context.question_annotation = annotation # 步骤2:分词 seg_result = await segment_text(context.q) context.keywords = seg_result.words # 步骤3:生成query组合 print(f"\n{'='*60}") print(f"生成query组合(最大组合数:{max_combination_size})") print(f"{'='*60}") query_combinations = generate_query_combinations(context.keywords, max_combination_size) context.query_combinations = query_combinations # 步骤4:获取所有query的推荐词,并流式评估 all_queries = [] for level, queries in query_combinations.items(): all_queries.extend(queries) print(f"\n{'='*60}") print(f"流式处理:边获取推荐词边评估(API并发度:{API_CONCURRENCY_LIMIT},模型并发度:{MODEL_CONCURRENCY_LIMIT})") print(f"{'='*60}") # 流式处理:边获取边评估 sug_results, evaluations = await fetch_and_evaluate_streaming( all_queries, context.q, annotation, context ) # 步骤6:筛选合格query qualified = find_qualified_queries(evaluations, min_relevance_score=0.7) if qualified: return { "success": True, "results": qualified, "message": f"找到 {len(qualified)} 个合格query(intent_match=True 且 relevance>=0.7)" } # 降低标准 acceptable = find_qualified_queries(evaluations, min_relevance_score=0.5) if acceptable: return { "success": True, "results": acceptable, "message": f"找到 {len(acceptable)} 个可接受query(intent_match=True 且 relevance>=0.5)" } # 完全失败:返回所有intent_match=True的 intent_matched = [e for e in evaluations if e['intent_match'] is True] if intent_matched: intent_matched_sorted = sorted(intent_matched, key=lambda x: x['relevance_score'], reverse=True) return { "success": False, "results": intent_matched_sorted[:10], # 只返回前10个 "message": f"未找到高相关性query,但有 {len(intent_matched)} 个意图匹配的推荐词" } return { "success": False, "results": [], "message": "未找到任何意图匹配的推荐词" } # ============================================================================ # 输出格式化 # ============================================================================ def format_output(optimization_result: dict, context: RunContext) -> str: """格式化输出结果""" results = optimization_result.get("results", []) output = f"原始问题:{context.q}\n" output += f"问题标注:{context.question_annotation}\n" output += f"提取的关键词:{', '.join(context.keywords or [])}\n" output += f"关键词数量:{len(context.keywords or [])}\n" output += f"\nquery组合统计:\n" for level, queries in context.query_combinations.items(): output += f" - {level}: {len(queries)} 个\n" # 统计信息 total_queries = sum(len(q) for q in context.query_combinations.values()) total_sugs = sum(len(r["suggestions"]) for r in context.all_sug_queries) total_evals = len(context.evaluation_results) output += f"\n探索统计:\n" output += f" - 总query数:{total_queries}\n" output += f" - 总推荐词数:{total_sugs}\n" output += f" - 总评估数:{total_evals}\n" output += f"\n状态:{optimization_result['message']}\n\n" if optimization_result["success"] and results: output += "=" * 60 + "\n" output += "合格的推荐query(按relevance_score降序):\n" output += "=" * 60 + "\n" for i, result in enumerate(results[:20], 1): # 只显示前20个 output += f"\n{i}. [{result['relevance_score']:.2f}] {result['sug_query']}\n" output += f" 来源:{result['source_query']}\n" output += f" 意图:{'✓ 匹配' if result['intent_match'] else '✗ 不匹配'}\n" output += f" 理由:{result['reason'][:150]}...\n" if len(result['reason']) > 150 else f" 理由:{result['reason']}\n" else: output += "=" * 60 + "\n" output += "结果:未找到足够相关的推荐query\n" output += "=" * 60 + "\n" if results: output += "\n最接近的推荐词(前10个):\n\n" for i, result in enumerate(results[:10], 1): output += f"{i}. [{result['relevance_score']:.2f}] {result['sug_query']}\n" output += f" 来源:{result['source_query']}\n" output += f" 意图:{'✓ 匹配' if result['intent_match'] else '✗ 不匹配'}\n\n" # 按source_query分组显示 output += "\n" + "=" * 60 + "\n" output += "按查询词分组的推荐词情况:\n" output += "=" * 60 + "\n" for sug_data in context.all_sug_queries: source_q = sug_data["query"] sugs = sug_data["suggestions"] # 找到这个source_query对应的所有评估 related_evals = [e for e in context.evaluation_results if e["source_query"] == source_q] intent_match_count = sum(1 for e in related_evals if e["intent_match"]) avg_relevance = sum(e["relevance_score"] for e in related_evals) / len(related_evals) if related_evals else 0 output += f"\n查询:{source_q}\n" output += f" 推荐词数:{len(sugs)}\n" output += f" 意图匹配数:{intent_match_count}/{len(related_evals)}\n" output += f" 平均相关性:{avg_relevance:.2f}\n" # 显示前3个推荐词 if sugs: output += f" 示例推荐词:\n" for sug in sugs[:3]: eval_item = next((e for e in related_evals if e["sug_query"] == sug), None) if eval_item: output += f" - {sug} [意图:{'✓' if eval_item['intent_match'] else '✗'}, 相关:{eval_item['relevance_score']:.2f}]\n" else: output += f" - {sug}\n" return output.strip() # ============================================================================ # 主函数 # ============================================================================ async def main( input_dir: str, max_combination_size: int = 1, api_concurrency: int = API_CONCURRENCY_LIMIT, model_concurrency: int = MODEL_CONCURRENCY_LIMIT ): # 更新全局并发配置 global API_CONCURRENCY_LIMIT, MODEL_CONCURRENCY_LIMIT API_CONCURRENCY_LIMIT = api_concurrency MODEL_CONCURRENCY_LIMIT = model_concurrency 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, ) print(f"\n{'='*60}") print(f"并发配置") print(f"{'='*60}") print(f"API请求并发度:{API_CONCURRENCY_LIMIT}") print(f"模型评估并发度:{MODEL_CONCURRENCY_LIMIT}") # 执行组合式搜索(带问题标注) optimization_result = await combinatorial_search(run_context, max_combination_size=max_combination_size) # 格式化输出 final_output = format_output(optimization_result, run_context) print(f"\n{'='*60}") print("最终结果") print(f"{'='*60}") print(final_output) # 保存结果 run_context.optimization_result = optimization_result run_context.final_output = 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__": parser = argparse.ArgumentParser( description="搜索query优化工具 - v6.3 组合式搜索+问题标注版", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" 示例: # 默认参数 python sug_v6_3_with_annotation.py # 2词组合,API并发5,模型并发20 python sug_v6_3_with_annotation.py --max-combo 2 --api-concurrency 5 --model-concurrency 20 # 3词组合,降低并发度 python sug_v6_3_with_annotation.py --max-combo 3 --api-concurrency 3 --model-concurrency 10 """ ) parser.add_argument( "--input-dir", type=str, default="input/简单扣图", help="输入目录路径,默认: input/简单扣图" ) parser.add_argument( "--max-combo", type=int, default=1, help="最大组合词数(N),默认: 1" ) parser.add_argument( "--api-concurrency", type=int, default=API_CONCURRENCY_LIMIT, help=f"API请求并发度,默认: {API_CONCURRENCY_LIMIT}" ) parser.add_argument( "--model-concurrency", type=int, default=MODEL_CONCURRENCY_LIMIT, help=f"模型评估并发度,默认: {MODEL_CONCURRENCY_LIMIT}" ) args = parser.parse_args() asyncio.run(main( args.input_dir, max_combination_size=args.max_combo, api_concurrency=args.api_concurrency, model_concurrency=args.model_concurrency ))