feat: add topic tagging, semantic duplicate detection, and per-provider vision toggle
This commit is contained in:
@@ -0,0 +1,120 @@
|
||||
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())
|
||||
@@ -71,6 +71,7 @@ async def test_multimodal_vision_image_payload():
|
||||
name="OpenAI Vision",
|
||||
provider_type="openai",
|
||||
model="gpt-4o",
|
||||
supports_vision=True,
|
||||
is_active=True
|
||||
)
|
||||
repo_mock = AsyncMock()
|
||||
|
||||
Reference in New Issue
Block a user