121 lines
4.0 KiB
Python
121 lines
4.0 KiB
Python
import os
|
|
import asyncio
|
|
from db.database import init_db, close_db_pool
|
|
from db.repository import Repository
|
|
from db.models import Post, TargetChannel, SourceChannel
|
|
from core.llm import LLMClient
|
|
from services.ai_processor import AIProcessor
|
|
from services.collector import CollectorService
|
|
|
|
|
|
async def test_provider_vision_toggle():
|
|
await init_db()
|
|
repo = Repository()
|
|
|
|
prov_id = await repo.add_provider_profile(
|
|
name="Test Vision Model",
|
|
provider_type="openai",
|
|
model="gpt-4o",
|
|
supports_vision=False
|
|
)
|
|
assert prov_id is not None
|
|
|
|
p = await repo.get_provider_profile_by_id(prov_id)
|
|
assert p.supports_vision is False
|
|
|
|
await repo.update_provider_vision(prov_id, True)
|
|
p_updated = await repo.get_provider_profile_by_id(prov_id)
|
|
assert p_updated.supports_vision is True
|
|
|
|
await repo.update_provider_vision(prov_id, False)
|
|
p_reverted = await repo.get_provider_profile_by_id(prov_id)
|
|
assert p_reverted.supports_vision is False
|
|
|
|
await repo.delete_provider_profile(prov_id)
|
|
|
|
|
|
async def test_find_candidate_posts_by_tags():
|
|
await init_db()
|
|
repo = Repository()
|
|
|
|
import time
|
|
msg_id = int(time.time() * 1000) % 100000000
|
|
post1_id = await repo.create_raw_post(
|
|
source_channel_id=-100111222,
|
|
source_message_id=msg_id,
|
|
raw_text="قیمت بیتکوین به ۷۰ هزار دلار رسید و رکورد جدیدی ثبت کرد.",
|
|
tags=["بیت_کوین", "ارز_دیجیتال", "اقتصاد"],
|
|
subject="رکوردشکنی بیتکوین",
|
|
)
|
|
assert post1_id is not None
|
|
|
|
candidates = await repo.find_candidate_posts_by_tags(["بیت_کوین", "رمزارز"])
|
|
assert len(candidates) >= 1
|
|
found_ids = [c.id for c in candidates]
|
|
assert post1_id in found_ids
|
|
|
|
no_match = await repo.find_candidate_posts_by_tags(["ورزش", "فوتبال"])
|
|
assert post1_id not in [c.id for c in no_match]
|
|
|
|
|
|
async def test_semantic_duplicate_detection():
|
|
await init_db()
|
|
repo = Repository()
|
|
|
|
class FakeDuplicateLLM:
|
|
def __init__(self):
|
|
self.supports_vision = False
|
|
self.last_used_model = "fake"
|
|
self.last_used_provider = "fake"
|
|
|
|
async def generate_json(self, prompt, system_prompt=None, action_name=None, image_path=None):
|
|
if action_name == "extract_tags_and_subject":
|
|
return {
|
|
"subject": "رشد بیتکوین",
|
|
"tags": ["بیت_کوین", "کریپتو", "قیمت"]
|
|
}
|
|
elif action_name == "check_semantic_duplicate":
|
|
return {
|
|
"is_duplicate": True,
|
|
"duplicate_of_id": 55,
|
|
"reason": "پوشش یکسان خبر افزایش نرخ بیتکوین"
|
|
}
|
|
return {}
|
|
|
|
fake_llm = FakeDuplicateLLM()
|
|
processor = AIProcessor(repo=repo, llm=fake_llm)
|
|
|
|
tags, subject = await processor.extract_tags_and_subject("بیتکوین ۷۰ هزار دلار شد")
|
|
assert tags == ["بیت_کوین", "کریپتو", "قیمت"]
|
|
assert subject == "رشد بیتکوین"
|
|
|
|
candidates = [
|
|
Post(
|
|
id=55,
|
|
source_channel_id=-100123,
|
|
source_message_id=10,
|
|
raw_text="بیت کوین به ۷۰۰۰۰ دلار رسید",
|
|
tags=["بیت_کوین"],
|
|
subject="افزایش قیمت بیتکوین"
|
|
)
|
|
]
|
|
|
|
is_dup, dup_id, reason = await processor.check_semantic_duplicate(
|
|
"بیتکوین به مرز هفتاد هزار دلار دست یافت",
|
|
candidates
|
|
)
|
|
assert is_dup is True
|
|
assert dup_id == 55
|
|
assert "پوشش یکسان" in reason
|
|
|
|
|
|
async def main():
|
|
await test_provider_vision_toggle()
|
|
await test_find_candidate_posts_by_tags()
|
|
await test_semantic_duplicate_detection()
|
|
print("All semantic deduplication and vision configuration tests passed successfully!")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|