feat: add source channel context history configuration and original sent timestamp
This commit is contained in:
@@ -0,0 +1,110 @@
|
||||
import asyncio
|
||||
import time
|
||||
from datetime import datetime, timezone, timedelta
|
||||
from unittest.mock import AsyncMock, patch, MagicMock
|
||||
from db.database import init_db
|
||||
from db.repository import Repository
|
||||
from db.models import TargetChannel, Post, AIProviderProfile
|
||||
from services.ai_processor import AIProcessor
|
||||
from core.llm import LLMClient
|
||||
|
||||
|
||||
async def test_source_created_at_and_context_count():
|
||||
await init_db()
|
||||
repo = Repository()
|
||||
|
||||
unique_channel_id = -10099887700 - int(time.time() % 100000)
|
||||
src_id = await repo.add_source(unique_channel_id, "Context Test Source", "ctx_test")
|
||||
assert src_id is not None
|
||||
|
||||
# Check default context_message_count is 0
|
||||
src = await repo.get_source_by_id(src_id)
|
||||
assert src.context_message_count == 0
|
||||
|
||||
# Update context_message_count to 5
|
||||
await repo.update_source_context_count(src_id, 5)
|
||||
src_updated = await repo.get_source_by_id(src_id)
|
||||
assert src_updated.context_message_count == 5
|
||||
|
||||
# Insert posts with explicit source_created_at timestamps
|
||||
t0 = datetime(2026, 8, 28, 10, 0, 0, tzinfo=timezone.utc)
|
||||
t1 = datetime(2026, 8, 28, 11, 0, 0, tzinfo=timezone.utc)
|
||||
t2 = datetime(2026, 8, 28, 12, 0, 0, tzinfo=timezone.utc)
|
||||
|
||||
p0_id = await repo.create_raw_post(
|
||||
source_channel_id=unique_channel_id,
|
||||
source_message_id=101,
|
||||
raw_text="خبر اول: مذاکرات آغاز شد.",
|
||||
source_created_at=t0
|
||||
)
|
||||
p1_id = await repo.create_raw_post(
|
||||
source_channel_id=unique_channel_id,
|
||||
source_message_id=102,
|
||||
raw_text="خبر دوم: توافقات اولیه حاصل گردید.",
|
||||
source_created_at=t1
|
||||
)
|
||||
p2_id = await repo.create_raw_post(
|
||||
source_channel_id=unique_channel_id,
|
||||
source_message_id=103,
|
||||
raw_text="خبر سوم: بیانیه مشترک امضا شد.",
|
||||
source_created_at=t2
|
||||
)
|
||||
|
||||
assert p0_id is not None
|
||||
assert p1_id is not None
|
||||
assert p2_id is not None
|
||||
|
||||
# Check source_created_at persisted
|
||||
p2_loaded = await repo.get_post_by_id(p2_id)
|
||||
assert p2_loaded.source_created_at is not None
|
||||
|
||||
# Fetch recent 2 posts excluding p2_id -> should return p0 and p1 in chronological order
|
||||
recent_posts = await repo.get_recent_source_posts(unique_channel_id, limit=2, exclude_post_id=p2_id)
|
||||
assert len(recent_posts) == 2
|
||||
assert recent_posts[0].id == p0_id
|
||||
assert recent_posts[1].id == p1_id
|
||||
|
||||
# Clean up
|
||||
await repo.delete_source(src_id)
|
||||
|
||||
|
||||
async def test_ai_processor_context_injection():
|
||||
repo_mock = AsyncMock()
|
||||
llm_mock = MagicMock()
|
||||
captured_payload = {}
|
||||
|
||||
async def fake_generate_json(prompt, system_prompt, action_name, image_path=None):
|
||||
captured_payload["prompt"] = prompt
|
||||
captured_payload["system_prompt"] = system_prompt
|
||||
return {"decision": "accept", "rewritten_text": "پست بازنویسیشده با در نظر گرفتن پیوستگی زمینه"}
|
||||
|
||||
llm_mock.generate_json = AsyncMock(side_effect=fake_generate_json)
|
||||
|
||||
processor = AIProcessor(repo=repo_mock, llm=llm_mock, double_check=False)
|
||||
target = TargetChannel(id=1, channel_id=-1001234, title="Target News", username="trg_news", personality="رسمی و خبری")
|
||||
|
||||
ctx_p1 = Post(id=10, source_channel_id=-100, source_message_id=1, raw_text="پست زمینه ۱: مرحله اول آغاز شد.", source_created_at="2026-08-28 10:00:00")
|
||||
ctx_p2 = Post(id=11, source_channel_id=-100, source_message_id=2, raw_text="پست زمینه ۲: مرحله دوم با موفقیت انجام شد.", source_created_at="2026-08-28 11:00:00")
|
||||
|
||||
result = await processor.rewrite_for_target(
|
||||
raw_text="پست جدید ۳: نتایج نهایی اعلام شد.",
|
||||
target=target,
|
||||
context_posts=[ctx_p1, ctx_p2]
|
||||
)
|
||||
|
||||
assert result.is_rejected is False
|
||||
prompt_sent = captured_payload["prompt"]
|
||||
assert "پیامها و پستهای قبلی/اخیر این کانال" in prompt_sent
|
||||
assert "پست زمینه ۱" in prompt_sent
|
||||
assert "پست زمینه ۲" in prompt_sent
|
||||
assert "پست جدید ۳: نتایج نهایی اعلام شد." in prompt_sent
|
||||
|
||||
|
||||
async def main():
|
||||
await test_source_created_at_and_context_count()
|
||||
await test_ai_processor_context_injection()
|
||||
print("All source context history and original timestamp tests passed successfully!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
Reference in New Issue
Block a user