84 lines
4.1 KiB
Python
84 lines
4.1 KiB
Python
"""Opt-in live probe using synthetic input only; never prints keys or provider bodies.
|
|
|
|
Run from backend: .venv/Scripts/python scripts/probe_context_recovery.py --live --model NAME --rows 40000
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import re
|
|
import time
|
|
|
|
from langchain.agents.middleware.types import ModelRequest, ModelResponse
|
|
from langchain_core.messages import HumanMessage, message_chunk_to_message
|
|
|
|
from deerflow.agents.middlewares.context_window_guard_middleware import ContextWindowGuardMiddleware, estimate_model_request_tokens
|
|
from deerflow.agents.middlewares.llm_error_handling_middleware import is_llm_error_message
|
|
from deerflow.config.app_config import get_app_config
|
|
from deerflow.models.context_recovery import CONTEXT_RECOVERY_MESSAGE, is_context_length_error
|
|
from deerflow.models.factory import create_chat_model
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--live", action="store_true", required=True)
|
|
parser.add_argument("--model", required=True)
|
|
parser.add_argument("--rows", type=int, default=100)
|
|
parser.add_argument("--stream", action="store_true")
|
|
args = parser.parse_args()
|
|
if not 0 <= args.rows <= 100000:
|
|
parser.error("rows must be between 0 and 100000")
|
|
logging.disable(logging.CRITICAL)
|
|
started = time.monotonic()
|
|
attempts = []
|
|
report = {"model": args.model, "rows": args.rows, "stream": args.stream, "attempts": attempts}
|
|
try:
|
|
config = get_app_config().model_copy(deep=True)
|
|
selected = config.get_model_config(args.model)
|
|
selected.max_retries = 0
|
|
selected.request_timeout = 90
|
|
selected.max_tokens = 1024
|
|
model = create_chat_model(name=args.model, app_config=config, thinking_enabled=False, force_disable_thinking=True)
|
|
text = "以下是合成测试材料。只回复：验证通过。\n" + "".join(f"记录{i:06d}：甲乙丙丁，气温正常。\n" for i in range(args.rows)) + "\n只回复：验证通过。"
|
|
request = ModelRequest(model=model, messages=[HumanMessage(content=text)])
|
|
|
|
def invoke(current):
|
|
attempt = {"estimated_input": estimate_model_request_tokens(current)}
|
|
attempts.append(attempt)
|
|
try:
|
|
messages = ([current.system_message] if current.system_message else []) + current.messages
|
|
if args.stream:
|
|
chunks = list(model.stream(messages, **current.model_settings))
|
|
response = message_chunk_to_message(sum(chunks[1:], chunks[0]))
|
|
else:
|
|
response = model.invoke(messages, **current.model_settings)
|
|
attempt["outcome"] = "response"
|
|
return ModelResponse(result=[response])
|
|
except Exception as exc:
|
|
attempt["outcome"] = "context_length" if is_context_length_error(exc) else type(exc).__name__
|
|
body = getattr(exc, "body", None)
|
|
if isinstance(body, dict):
|
|
error = body.get("error", body)
|
|
if isinstance(error, dict):
|
|
attempt["provider_code"] = error.get("code")
|
|
attempt["provider_message"] = str(error.get("message", ""))[:500]
|
|
raise
|
|
|
|
response = ContextWindowGuardMiddleware().wrap_model_call(request, invoke)
|
|
message = response.result[0] if isinstance(response, ModelResponse) else response
|
|
text = message.text.strip()
|
|
text = re.sub(r"<think>.*?</think>", "", text, flags=re.S).strip()
|
|
if "<think>" in text:
|
|
text = text.split("<think>", 1)[0].strip()
|
|
report["outcome"] = "answered" if "验证通过" in text and not is_llm_error_message(message) else "clarification" if text == CONTEXT_RECOVERY_MESSAGE else "unexpected_response"
|
|
report["visible_chars"] = len(text)
|
|
except Exception as exc:
|
|
report["outcome"] = type(exc).__name__
|
|
report["elapsed_seconds"] = round(time.monotonic() - started, 2)
|
|
print(json.dumps(report, ensure_ascii=True))
|
|
return 0 if report["outcome"] == "answered" else 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|