Spaces:
Running
Running
File size: 1,070 Bytes
a63c61f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 | 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}")
|