tests/test_rag_pipeline.py

import sys
import os
from unittest.mock import MagicMock

# Add src to path
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), 'src')))

try:
    import tiktoken
    from src.api.routes.rag import limit_context_tokens
    print("Tiktoken and limit_context_tokens imported successfully.")
except ImportError as e:
    print(f"Import error: {e}")

def test_context_limiter():
    sources = [
        {"content": "This is a long piece of text " * 100, "metadata": {}},
        {"content": "Another short text", "metadata": {}},
        {"content": "Third piece of text", "metadata": {}}
    ]
    
    # Test with 100 tokens
    context, filtered = limit_context_tokens(sources, max_tokens=100)
    enc = tiktoken.get_encoding("cl100k_base")
    tokens = len(enc.encode(context))
    print(f"Context token count: {tokens}")
    assert tokens <= 110 # slight buffer for delimiters
    print("Context limiter test passed!")

if __name__ == "__main__":
    try:
        test_context_limiter()
    except Exception as e:
        print(f"Test failed: {e}")