98 lines
3.1 KiB
Python
98 lines
3.1 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import re
|
|
|
|
import httpx
|
|
|
|
from .config import Settings
|
|
from .parsers import extract_json_object
|
|
from .prompts import (
|
|
PLANNER_DECOMPOSE_PROMPT,
|
|
PLANNER_NEXT_QUERY_PROMPT,
|
|
PLANNER_TRANSLATE_PROMPT,
|
|
)
|
|
|
|
_CJK_RE = re.compile(r"[\u4e00-\u9fff\u3040-\u30ff\uac00-\ud7af]")
|
|
|
|
|
|
class QueryPlanner:
|
|
"""LLM-driven query decomposition, multilingual expansion, and next-round planning."""
|
|
|
|
def __init__(self, settings: Settings):
|
|
self.settings = settings
|
|
|
|
async def decompose(self, query: str) -> list[str]:
|
|
prompt = PLANNER_DECOMPOSE_PROMPT.format(query=query)
|
|
result = await asyncio.to_thread(self._call_sync, prompt, max_tokens=300)
|
|
raw = extract_json_object(result)
|
|
if raw:
|
|
try:
|
|
subs = json.loads(raw).get("sub_queries", [])
|
|
cleaned = [s.strip() for s in subs if s.strip()]
|
|
if cleaned:
|
|
return cleaned
|
|
except Exception:
|
|
pass
|
|
return [query]
|
|
|
|
async def expand_multilingual(self, query: str) -> list[str]:
|
|
if not _CJK_RE.search(query):
|
|
return [query]
|
|
prompt = PLANNER_TRANSLATE_PROMPT.format(query=query)
|
|
result = await asyncio.to_thread(self._call_sync, prompt, max_tokens=100)
|
|
raw = extract_json_object(result)
|
|
if raw:
|
|
try:
|
|
en = json.loads(raw).get("en", "").strip()
|
|
if en and en.lower() != query.lower():
|
|
return [query, en]
|
|
except Exception:
|
|
pass
|
|
return [query]
|
|
|
|
async def next_query(
|
|
self,
|
|
original_query: str,
|
|
round_index: int,
|
|
search_digest: str,
|
|
usable_count: int,
|
|
) -> str:
|
|
prompt = PLANNER_NEXT_QUERY_PROMPT.format(
|
|
query=original_query,
|
|
round_index=round_index,
|
|
usable_count=usable_count,
|
|
digest=search_digest[:600],
|
|
)
|
|
result = await asyncio.to_thread(self._call_sync, prompt, max_tokens=150)
|
|
raw = extract_json_object(result)
|
|
if raw:
|
|
try:
|
|
return json.loads(raw).get("next_query", "").strip()
|
|
except Exception:
|
|
pass
|
|
return ""
|
|
|
|
def _call_sync(self, user_prompt: str, max_tokens: int = 200) -> str:
|
|
payload = {
|
|
"model": self.settings.llm_model,
|
|
"messages": [{"role": "user", "content": user_prompt}],
|
|
"temperature": 0.1,
|
|
"max_tokens": max_tokens,
|
|
}
|
|
try:
|
|
with httpx.Client(timeout=30) as client:
|
|
resp = client.post(
|
|
self.settings.llm_chat_endpoint,
|
|
headers={
|
|
"Authorization": f"Bearer {self.settings.llm_api_key}",
|
|
"Content-Type": "application/json",
|
|
},
|
|
json=payload,
|
|
)
|
|
resp.raise_for_status()
|
|
return resp.json()["choices"][0]["message"]["content"].strip()
|
|
except Exception:
|
|
return ""
|