210 lines
7.4 KiB
Python
210 lines
7.4 KiB
Python
"""The rewrite prompt embeds a literal JSON example, which str.format() must not choke on."""
|
|
import sys
|
|
|
|
sys.path.insert(0, "/app")
|
|
|
|
from services.ai_processor import build_rewrite_system_prompt, build_verify_system_prompt, AIProcessor
|
|
from core.llm import LLMClient
|
|
from db.models import TargetChannel
|
|
import asyncio
|
|
|
|
|
|
def test_prompt_builds_with_literal_json_braces():
|
|
prompt = build_rewrite_system_prompt(
|
|
channel_title="کانال تست",
|
|
personality="لحن دوستانه",
|
|
custom_footer="@test_channel",
|
|
language="fa",
|
|
has_media=False,
|
|
)
|
|
assert "کانال تست" in prompt
|
|
assert "لحن دوستانه" in prompt
|
|
assert "@test_channel" in prompt
|
|
# The JSON schema example must survive verbatim.
|
|
assert '"rewritten_text"' in prompt
|
|
assert "{" in prompt and "}" in prompt
|
|
assert "4000 characters" in prompt
|
|
|
|
|
|
def test_prompt_languages():
|
|
for lang, expected in [
|
|
("en", "English"),
|
|
("es", "Spanish"),
|
|
("ar", "Arabic"),
|
|
("tr", "Turkish"),
|
|
("fa", "Persian"),
|
|
]:
|
|
prompt = build_rewrite_system_prompt("Test", "Tone", "@tag", language=lang)
|
|
assert expected in prompt, f"Expected {expected} in prompt for lang {lang}"
|
|
|
|
|
|
def test_prompt_media_length_limits():
|
|
media_prompt = build_rewrite_system_prompt("Test", "Tone", "@tag", has_media=True)
|
|
assert "1000 characters" in media_prompt
|
|
assert "1024-character caption limit" in media_prompt
|
|
|
|
text_prompt = build_rewrite_system_prompt("Test", "Tone", "@tag", has_media=False)
|
|
assert "4000 characters" in text_prompt
|
|
assert "4096-character message limit" in text_prompt
|
|
|
|
|
|
def test_verify_prompt():
|
|
v_prompt = build_verify_system_prompt("English Channel", "Humorous Tone", "@eng_footer", language="en", has_media=True)
|
|
assert "English Channel" in v_prompt
|
|
assert "English" in v_prompt
|
|
assert "Humorous Tone" in v_prompt
|
|
assert "@eng_footer" in v_prompt
|
|
assert "1000 characters" in v_prompt
|
|
assert '"final_text"' in v_prompt
|
|
|
|
|
|
|
|
def test_llm_client_agy_provider():
|
|
client = LLMClient(provider="agy", base_url="http://127.0.0.1:8000/v1")
|
|
assert client.provider == "agy"
|
|
assert client.model == "antigravity"
|
|
assert client.base_url == "http://127.0.0.1:8000/v1"
|
|
|
|
|
|
def test_llm_client_reasoning_effort():
|
|
client = LLMClient(reasoning_effort="high")
|
|
assert client.reasoning_effort == "high"
|
|
|
|
|
|
|
|
class FakeLLM:
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
async def generate_json(self, prompt: str, system_prompt: str = None, action_name: str = "general", *args, **kwargs):
|
|
self.calls.append((prompt, system_prompt, action_name))
|
|
if action_name == "rewrite_target_post":
|
|
return {"rewritten_text": "Draft with Spanish and English mixed"}
|
|
elif action_name == "verify_target_post":
|
|
return {"final_text": "Clean English only verified post"}
|
|
return {}
|
|
|
|
|
|
async def async_test_ai_processor_double_check():
|
|
fake_llm = FakeLLM()
|
|
processor = AIProcessor(repo=None, llm=fake_llm, double_check=True)
|
|
target = TargetChannel(id=1, channel_id=123, title="Target", username="tgt", language="en")
|
|
|
|
result = await processor.rewrite_for_target("Original Raw Post", target, has_media=False)
|
|
assert result.rewritten_text == "Clean English only verified post"
|
|
assert str(result) == "Clean English only verified post"
|
|
assert result.is_rejected is False
|
|
assert len(fake_llm.calls) == 2
|
|
assert fake_llm.calls[0][2] == "rewrite_target_post"
|
|
assert fake_llm.calls[1][2] == "verify_target_post"
|
|
|
|
|
|
|
|
class FakeRepo:
|
|
def __init__(self):
|
|
self.settings = {"ai_provider": "gemini", "ai_model": "gemini-2.5-flash", "ai_double_check": "true"}
|
|
self.logs = []
|
|
self.active_provider = None
|
|
|
|
async def get_setting(self, key, default=None):
|
|
return self.settings.get(key, default)
|
|
|
|
async def set_setting(self, key, value, description=None):
|
|
self.settings[key] = value
|
|
|
|
async def record_ai_log(self, action_name, provider, model, prompt, system_prompt=None, response_text=None, duration_sec=0.0, status="success", error_message=None):
|
|
self.logs.append({
|
|
"action_name": action_name,
|
|
"provider": provider,
|
|
"model": model,
|
|
"prompt": prompt,
|
|
"status": status,
|
|
})
|
|
return len(self.logs)
|
|
|
|
async def get_active_provider_profile(self):
|
|
return self.active_provider
|
|
|
|
|
|
async def async_test_dynamic_config_and_logging():
|
|
fake_repo = FakeRepo()
|
|
client = LLMClient(provider="openai", model="orcarouter/auto", repo=fake_repo)
|
|
await client.sync_config_from_repo()
|
|
assert client.provider == "gemini"
|
|
assert client.model == "gemini-2.5-flash"
|
|
|
|
processor = AIProcessor(repo=fake_repo)
|
|
assert await processor._is_double_check_enabled() is True
|
|
|
|
|
|
async def async_test_provider_profile_sync():
|
|
from db.models import AIProviderProfile
|
|
fake_repo = FakeRepo()
|
|
fake_repo.active_provider = AIProviderProfile(
|
|
id=1,
|
|
name="Local AGY",
|
|
provider_type="agy",
|
|
model="antigravity-turbo",
|
|
base_url="http://localhost:8000/v1",
|
|
api_key="",
|
|
is_active=True
|
|
)
|
|
client = LLMClient(provider="openai", model="orcarouter/auto", repo=fake_repo)
|
|
await client.sync_config_from_repo()
|
|
assert client.provider == "agy"
|
|
assert client.model == "antigravity-turbo"
|
|
assert client.base_url == "http://localhost:8000/v1"
|
|
|
|
|
|
async def async_test_ai_processor_rejection():
|
|
class RejectLLM:
|
|
async def generate_json(self, prompt: str, system_prompt: str = None, action_name: str = "general", *args, **kwargs):
|
|
return {
|
|
"decision": "reject",
|
|
"rejection_reason": "تبلیغات نامرتبط و غیرمجاز",
|
|
"rewritten_text": ""
|
|
}
|
|
|
|
processor = AIProcessor(repo=None, llm=RejectLLM())
|
|
target = TargetChannel(id=1, channel_id=123, title="News Channel", username="news", custom_prompt="فقط خبرهای فناوری")
|
|
res = await processor.rewrite_for_target("تبلیغ خرید پکیج آموزشی", target)
|
|
assert res.is_rejected is True
|
|
assert res.rejection_reason == "تبلیغات نامرتبط و غیرمجاز"
|
|
assert str(res) == ""
|
|
|
|
|
|
def test_custom_target_prompt_in_system_prompt():
|
|
prompt = build_rewrite_system_prompt(
|
|
channel_title="Tech News",
|
|
personality="رسمی",
|
|
custom_footer="@tech_news",
|
|
custom_prompt="همیشه تیتر را با ایموجی 🔴 شروع کن و خلاصهای در ۳ خط بنویس",
|
|
language="fa"
|
|
)
|
|
assert "TARGET CHANNEL SPECIFIC COMMANDS & FORMATTING RULES:" in prompt
|
|
assert "همیشه تیتر را با ایموجی 🔴 شروع کن" in prompt
|
|
assert "decision" in prompt
|
|
assert "rejection_reason" in prompt
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_prompt_builds_with_literal_json_braces()
|
|
test_prompt_languages()
|
|
test_prompt_media_length_limits()
|
|
test_verify_prompt()
|
|
test_custom_target_prompt_in_system_prompt()
|
|
test_llm_client_agy_provider()
|
|
test_llm_client_reasoning_effort()
|
|
asyncio.run(async_test_ai_processor_double_check())
|
|
asyncio.run(async_test_ai_processor_rejection())
|
|
asyncio.run(async_test_dynamic_config_and_logging())
|
|
asyncio.run(async_test_provider_profile_sync())
|
|
print("All prompt, custom instructions, rejection, and multi-provider profile unit tests passed!")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|