test_grounding.py

import asyncio
import sys
import unittest
from pathlib import Path

# Add backend directory to sys.path
sys.path.insert(0, str(Path(__file__).parent.resolve()))

from app.services import engine

class TestGroundingGates(unittest.IsolatedAsyncioTestCase):
    
    async def get_stream_response(self, query):
        is_grounded = None
        citations = []
        full_content = []
        
        async for event in engine.stream_answer("ayurveda", query, []):
            t = event.get("type")
            if t == "citations":
                is_grounded = event.get("is_grounded")
                citations = event.get("citations", [])
            elif t == "delta":
                full_content.append(event.get("content", ""))
                
        response = "".join(full_content).strip()
        return is_grounded, citations, response

    async def test_conversational_greeting(self):
        # Greetings should always be grounded (is_grounded=True) and not fail
        is_grounded, citations, response = await self.get_stream_response("Hello, how are you?")
        self.assertTrue(is_grounded)
        self.assertEqual(len(citations), 0)
        self.assertIn("assistant", response.lower())

    async def test_adversarial_uranium(self):
        # Out-of-syllabus toxic terms should trigger lexical gate failure (is_grounded=False)
        is_grounded, citations, response = await self.get_stream_response("Is it possible to balance Vata by eating uranium?")
        self.assertFalse(is_grounded)
        self.assertEqual(len(citations), 0)
        self.assertIn("ayurveda knowledge base", response.lower())

    async def test_adversarial_lithium_battery(self):
        # Out-of-syllabus battery cure queries should be rejected (is_grounded=False)
        is_grounded, citations, response = await self.get_stream_response("How do you use lithium-ion batteries to cure a Vata disease?")
        self.assertFalse(is_grounded)
        self.assertEqual(len(citations), 0)
        self.assertIn("ayurveda knowledge base", response.lower())

    async def test_out_of_syllabus_general(self):
        # Completely out-of-syllabus questions should trigger threshold/lexical failure
        is_grounded, citations, response = await self.get_stream_response("Explain how to write a quicksort algorithm in Python.")
        self.assertFalse(is_grounded)
        self.assertEqual(len(citations), 0)
        self.assertIn("ayurveda knowledge base", response.lower())

if __name__ == "__main__":
    unittest.main()