“ydy0615”
2c7a02f587
- Added support for PRO mode in LLM with specific instruction handling and context awareness. - Updated prompt building functions to include prefill options for better context management. - Introduced new inline examples for PRO mode in JSON format. - Enhanced system prompts to reflect PRO mode capabilities and rules. - Modified API endpoints to accommodate new parameters and ensure backward compatibility. - Improved test cases to validate new functionality and ensure comprehensive coverage.
118 lines
3.7 KiB
Python
118 lines
3.7 KiB
Python
import os
|
|
import sys
|
|
import types
|
|
import asyncio
|
|
import threading
|
|
from fastapi.testclient import TestClient
|
|
|
|
|
|
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
BACKEND_DIR = os.path.abspath(os.path.join(CURRENT_DIR, ".."))
|
|
if BACKEND_DIR not in sys.path:
|
|
sys.path.insert(0, BACKEND_DIR)
|
|
|
|
if "tts_asr" not in sys.modules:
|
|
fake_tts_asr = types.ModuleType("tts_asr")
|
|
fake_tts_asr.register_tts_asr_routes = lambda app: None
|
|
sys.modules["tts_asr"] = fake_tts_asr
|
|
|
|
import main # type: ignore
|
|
import pro_completions # type: ignore
|
|
|
|
|
|
HEADERS = {"X-API-Key": main.API_KEY}
|
|
|
|
|
|
def _payload():
|
|
return {
|
|
"prefix": "Before",
|
|
"suffix": "After",
|
|
"languageId": "markdown",
|
|
"instruction": "expand",
|
|
"pro_thinking": "medium",
|
|
"privacy_mode": True,
|
|
}
|
|
|
|
|
|
def setup_function():
|
|
pro_completions.PRO_STATES.clear()
|
|
|
|
|
|
def teardown_function():
|
|
pro_completions.PRO_STATES.clear()
|
|
|
|
|
|
def test_pro_queue_full_returns_429(monkeypatch):
|
|
monkeypatch.setattr(pro_completions, "PRO_QUEUE_MAX_SIZE", 0)
|
|
client = TestClient(main.app)
|
|
response = client.post("/v1/pro/completions", headers=HEADERS, json=_payload())
|
|
assert response.status_code == 429
|
|
assert response.json()["error"] == "PRO queue is full"
|
|
|
|
|
|
def test_pro_status_missing_returns_404():
|
|
client = TestClient(main.app)
|
|
response = client.get("/v1/pro/completions/status/missing", headers=HEADERS)
|
|
assert response.status_code == 404
|
|
|
|
|
|
def test_pro_prompt_uses_pro_specific_instruction():
|
|
system_prompt, user_prompt = pro_completions._build_pro_prompts(
|
|
prefix="欢迎使用 LLM-IN-TEXT\n\n即时可用的 LLM 系统",
|
|
suffix="",
|
|
language_id="markdown",
|
|
instruction="",
|
|
pro_thinking="high",
|
|
)
|
|
combined = f"{system_prompt}\n{user_prompt}".lower()
|
|
assert "[pro] model for llm-in-text" in combined
|
|
assert "pro_mode: true" in combined
|
|
assert "pro_thinking_level: high" in combined
|
|
assert "long paragraphs or section-level output are allowed" in combined
|
|
assert "highest priority" in combined
|
|
assert "never copy tags to output" in combined
|
|
assert "write only the markdown that belongs at the cursor" not in combined
|
|
assert "continue the markdown naturally" in combined
|
|
|
|
|
|
def test_pro_cancel_waits_for_stream_cleanup(monkeypatch):
|
|
started = threading.Event()
|
|
cleaned = threading.Event()
|
|
|
|
async def fake_stream_events(*args, **kwargs):
|
|
started.set()
|
|
try:
|
|
yield "thinking", ""
|
|
while True:
|
|
await asyncio.sleep(0.05)
|
|
finally:
|
|
cleaned.set()
|
|
|
|
monkeypatch.setattr(pro_completions, "stream_ollama_events", fake_stream_events)
|
|
request_id = "pro-cancel-cleanup"
|
|
headers = {**HEADERS, "X-Request-Id": request_id}
|
|
response_box = {}
|
|
|
|
with TestClient(main.app) as client:
|
|
def send_stream():
|
|
with client.stream("POST", "/v1/pro/completions", headers=headers, json=_payload()) as response:
|
|
response_box["status_code"] = response.status_code
|
|
response_box["body"] = "".join(response.iter_text())
|
|
|
|
stream_thread = threading.Thread(target=send_stream, daemon=True)
|
|
stream_thread.start()
|
|
|
|
assert started.wait(timeout=2.0)
|
|
cancel_response = client.post(
|
|
"/v1/pro/completions/cancel",
|
|
headers=HEADERS,
|
|
json={"request_id": request_id, "reason": "test"},
|
|
)
|
|
|
|
assert cancel_response.status_code == 200
|
|
assert cancel_response.json() == {"cancelled": True, "status": "ok"}
|
|
assert cleaned.wait(timeout=2.0)
|
|
|
|
stream_thread.join(timeout=5.0)
|
|
assert not stream_thread.is_alive()
|