"""
حفاظهای بازبینیِ امنیتیِ پیش از تحویل (دستیار). مدل و Qdrant همه ساختگیاند.
"""
import json
from types import SimpleNamespace
from unittest.mock import patch
from django.core.cache import cache
from django.test import TestCase, override_settings
from rest_framework.test import APIClient
from accounts.authentication import issue_access_token
from accounts.models import User
from ai_tutor.models import Message
from rag.services.clients import reset_clients_cache
RAG_ON = dict(OPENAI_API_KEY='t', QDRANT_API_KEY='t', RAG_FORCE_MOCK=False, AI_DAILY_QUOTA_FREE=3)
ANSWER = '۱) دو بهعلاوهی دو میشود چهار.\n✅ جواب: ۴\n'
OVERFLOW_DIAGRAM = '{"kind":"numberline","min":1' + '0' * 400 + '*2,"max":5}'
DEEP_DIAGRAM = '' + '[' * 3000 + ']' * 3000 + ''
class FakeOpenAI:
def __init__(self, reply):
self.reply = reply
self.calls = 0
self.chat = SimpleNamespace(completions=SimpleNamespace(create=self._create))
self.embeddings = SimpleNamespace(
create=lambda **k: SimpleNamespace(data=[SimpleNamespace(embedding=[0.01] * 1536)]))
def _create(self, **kw):
self.calls += 1
if kw.get('stream'):
parts = [self.reply[i:i + 25] for i in range(0, len(self.reply), 25)]
events = [SimpleNamespace(choices=[SimpleNamespace(delta=SimpleNamespace(content=p))], usage=None)
for p in parts]
events.append(SimpleNamespace(choices=[], usage=SimpleNamespace(prompt_tokens=900, completion_tokens=300)))
return iter(events)
return SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(content=self.reply))],
usage=SimpleNamespace(prompt_tokens=900, completion_tokens=300))
def sse(resp):
out = []
for block in b''.join(resp.streaming_content).decode().split('\n\n'):
ev = data = None
for line in block.strip().splitlines():
if line.startswith('event: '):
ev = line[7:]
elif line.startswith('data: '):
data = json.loads(line[6:])
if ev:
out.append((ev, data))
return out
@override_settings(**RAG_ON)
class _Base(TestCase):
def setUp(self):
cache.clear()
reset_clients_cache()
self.user = User.objects.create(phone='09120009901')
self.c = APIClient()
self.c.credentials(HTTP_AUTHORIZATION=f'Bearer {issue_access_token(self.user.id)}')
def ask(self, q, **extra):
return self.c.post('/api/ai/ask', {'userQuestion': q, **extra}, format='json')
def remaining(self):
return self.c.get('/api/ai/quota').data['remainingQuota']
class BrokenDiagramDoesNotMakeAnswersFreeTests(_Base):
"""پارسرِ شکل OverflowError/RecursionError را بالا میداد و جوابِ استریمشده رایگان و بیثبت میماند."""
def test_stream_and_json_answers_are_charged(self):
for diagram in (OVERFLOW_DIAGRAM, DEEP_DIAGRAM):
with patch('rag.services.clients.get_openai_client', return_value=FakeOpenAI(ANSWER + diagram)):
events = sse(self.ask(f'دو بهعلاوهی دو؟ {len(diagram)}', stream=True))
self.assertEqual(events[-1][0], 'done')
self.assertIsNone(events[-1][1]['diagram'])
with patch('rag.services.clients.get_openai_client', return_value=FakeOpenAI(ANSWER + OVERFLOW_DIAGRAM)):
self.assertEqual(self.ask('دو بهعلاوهی دو؟').status_code, 200)
self.assertEqual(self.remaining(), 0)
self.assertEqual(Message.objects.filter(role='assistant').count(), 3)
@override_settings(RAG_ANSWER_CACHE_TTL=300)
class LateOffTopicMarkerTests(_Base):
def test_rest_is_not_streamed_stored_or_cached(self):
from rag.services.guardrails import OFF_TOPIC_MARKER, REFUSAL
reply = ('۱) این یک جملهی ریاضیِ بلند برای عبور از چهل نویسه است.\n'
f'{OFF_TOPIC_MARKER} متنِ غیرِ ریاضی: یک انشای کامل دربارهی فوتبال.\n')
fake = FakeOpenAI(reply)
q = 'سوال دربارهی چیزی'
with patch('rag.services.clients.get_openai_client', return_value=fake):
streamed = ''.join(d['delta'] for e, d in sse(self.ask(q, stream=True)) if e == 'chunk')
other = User.objects.create(phone='09120009902')
c2 = APIClient()
c2.credentials(HTTP_AUTHORIZATION=f'Bearer {issue_access_token(other.id)}')
second = c2.post('/api/ai/ask', {'userQuestion': q}, format='json')
self.assertNotIn('فوتبال', streamed)
self.assertIn(REFUSAL, streamed)
self.assertEqual(Message.objects.filter(role='assistant').first().text, REFUSAL)
self.assertNotIn('فوتبال', second.data['answer'])
self.assertEqual(fake.calls, 2) # جوابِ ردشده در کش نرفت
class DisconnectMidStreamTests(_Base):
def test_partial_answer_is_saved(self):
with patch('rag.services.clients.get_openai_client', return_value=FakeOpenAI(ANSWER * 10)):
resp = self.ask('سوال', stream=True)
it = iter(resp.streaming_content)
next(it)
next(it)
resp._iterator.close() # همان GeneratorExit که بستنِ صفحه میسازد
saved = Message.objects.filter(role='assistant')
self.assertEqual(saved.count(), 1)
self.assertTrue(saved.first().text.endswith('…'))
self.assertEqual(self.remaining(), 2)
class UnknownCollectionIsNotSearchedTests(_Base):
def test_client_chosen_collection_is_ignored(self):
seen = []
def fake_search(collection, vec, **kw):
seen.append(collection)
return [{'score': 0.9, 'payload': {'title': 'doc', 'text': 'PAYLOAD-OF-UNRELATED-COLLECTION'}}]
with patch('rag.services.clients.get_openai_client', return_value=FakeOpenAI('پاسخ')), \
patch('rag.services.retrieval.search', side_effect=fake_search):
r = self.ask('هر چیزی', subject='any', grade='other collection')
self.assertEqual(r.status_code, 200)
self.assertEqual(seen, [])
self.assertEqual(r.data['usedChunks'], [])