query_filter.py 4.5 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495
  1. """创作 query 筛选器:调用 acquisition/query_filter.txt,批量判一组 query。
  2. 模型对每条 query 打两项:语义合法性 valid(0-10) + 创作相关性 relevant(bool)。
  3. 最终保留 keep = valid ≥ VALID_MIN 且 relevant —— 阈值放代码里(好调,不改提示词)。
  4. · valid:词组合本身说不说得通(机械正交易产生的废话短语在此淘汰)。
  5. · relevant:搜回来是否和创作知识相关(三把尺=设计决策vs工艺/可迁移/业务阶段;再排 制作/题材本身/应试政企/作品)。
  6. 返回与输入等长的 [{keep, valid, relevant, reason}]。供 build_creation_demo、filter_multiaxis 等共用。
  7. """
  8. from __future__ import annotations
  9. import json
  10. import os
  11. from pathlib import Path
  12. import httpx
  13. from core.config import Settings, load_env_file
  14. PROMPT = Path(__file__).resolve().parent / "query_filter.txt"
  15. VALID_MIN = 6 # 语义合法性阈值:valid ≥ 此值才算说得通;调这里即可松紧
  16. def filter_queries(queries: list[str], settings: Settings, *, batch: int = 40) -> list[dict]:
  17. """批量过滤(分批调用,避免一次喂太多)。返回 [{keep, reason}],与 queries 对齐。"""
  18. out: list[dict] = []
  19. for i in range(0, len(queries), batch):
  20. out.extend(_filter_batch(queries[i:i + batch], settings))
  21. return out
  22. def _filter_batch(queries: list[str], settings: Settings) -> list[dict]:
  23. if not queries:
  24. return []
  25. user = json.dumps([{"idx": i, "query": q} for i, q in enumerate(queries)], ensure_ascii=False)
  26. messages = [
  27. {"role": "system", "content": PROMPT.read_text("utf-8")},
  28. {"role": "user", "content": user},
  29. ]
  30. try:
  31. txt = _chat_content(settings, messages)
  32. # query_filter 要求输出数组;有的模型会包一层 {"result":[...]},都兜住
  33. data = json.loads(txt)
  34. arr = data if isinstance(data, list) else next((v for v in data.values() if isinstance(v, list)), [])
  35. by = {d.get("idx"): d for d in arr if isinstance(d, dict)}
  36. out = []
  37. for i in range(len(queries)):
  38. d = by.get(i, {})
  39. try:
  40. valid = int(d.get("valid", 10))
  41. except (TypeError, ValueError):
  42. valid = 10
  43. relevant = bool(d.get("relevant", True))
  44. out.append({"keep": valid >= VALID_MIN and relevant, # 阈值在代码里判
  45. "valid": valid, "relevant": relevant,
  46. "reason": str(d.get("reason", ""))[:50]})
  47. return out
  48. except Exception as exc:
  49. return [{"keep": True, "valid": 10, "relevant": True, "reason": f"筛选失败:{str(exc)[:30]}"} for _ in queries]
  50. def _chat_content(settings: Settings, messages: list[dict]) -> str:
  51. """Return chat content. Prefer OpenRouter; fall back to Ark when OpenRouter is region-blocked locally."""
  52. body = {"model": settings.llm_model, "messages": messages, "response_format": {"type": "json_object"}}
  53. env = load_env_file(os.getenv("CK_ENV_FILE", ".env"))
  54. prefer_ark = os.getenv("QUERY_FILTER_PROVIDER") == "ark" or bool(os.getenv("ARK_CHAT_MODEL"))
  55. openrouter_exc: Exception | None = None
  56. if settings.openrouter_api_key and not prefer_ark:
  57. try:
  58. resp = httpx.post(
  59. settings.openrouter_base_url.rstrip("/") + "/chat/completions",
  60. headers={"Authorization": f"Bearer {settings.openrouter_api_key}", "Content-Type": "application/json"},
  61. json=body,
  62. timeout=120,
  63. )
  64. resp.raise_for_status()
  65. return resp.json()["choices"][0]["message"]["content"]
  66. except Exception as exc:
  67. openrouter_exc = exc
  68. ark_key = os.getenv("ARK_API_KEY") or env.get("ARK_API_KEY")
  69. if ark_key:
  70. ark_model = os.getenv("ARK_CHAT_MODEL") or env.get("ARK_CHAT_MODEL") or "doubao-seed-1-6-flash-250615"
  71. ark_url = os.getenv("ARK_CHAT_URL") or env.get("ARK_CHAT_URL") or "https://ark.cn-beijing.volces.com/api/v3/chat/completions"
  72. resp = httpx.post(
  73. ark_url,
  74. headers={"Authorization": f"Bearer {ark_key}", "Content-Type": "application/json"},
  75. json={"model": ark_model, "messages": messages, "response_format": {"type": "json_object"}},
  76. timeout=120,
  77. )
  78. resp.raise_for_status()
  79. return resp.json()["choices"][0]["message"]["content"]
  80. if openrouter_exc:
  81. raise openrouter_exc
  82. raise RuntimeError("missing OpenRouter/Ark chat credentials")