Spaces:
Sleeping
Sleeping
File size: 3,976 Bytes
ba5110e |
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 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 |
import asyncio
import httpx
import sys
import os
# Add project root to path
sys.path.append(os.getcwd())
from backend.utils.memory import memory_tracker, WARNING_TOKENS, BLOCK_TOKENS, KIMI_K2_CONTEXT_LENGTH
async def get_latest_session_id():
"""Fetch the most recent conversation ID from the database."""
try:
import sqlite3
conn = sqlite3.connect("algebra_chat.db")
cursor = conn.cursor()
cursor.execute("SELECT id FROM conversations ORDER BY created_at DESC LIMIT 1")
result = cursor.fetchone()
conn.close()
return result[0] if result else None
except Exception as e:
print(f"Error fetching latest session: {e}")
return None
async def test_memory_limits():
"""Test memory warning and blocking behavior."""
# Try to get latest session if not specified
session_id = await get_latest_session_id()
if not session_id:
session_id = "test_memory_session_v1"
print(f"! Không tìm thấy session nào trong DB, sử dụng ID mặc định: {session_id}")
else:
print(f"✨ Đã tìm thấy session mới nhất: {session_id}")
print(f"\n--- Testing Memory Limits for Session: {session_id} ---")
print(f"Max Tokens: {KIMI_K2_CONTEXT_LENGTH}")
print(f"Warning Threshold: {WARNING_TOKENS} (80%)")
print(f"Block Threshold: {BLOCK_TOKENS} (95%)")
# 1. Create a new session (implicitly via chat or explicit reset)
print("\n1. Resetting session memory...")
memory_tracker.reset_usage(session_id)
current = memory_tracker.get_usage(session_id)
print(f"Current Usage: {current}")
# 2. Test Normal State
print("\n2. Testing Normal State...")
print("Simulating 1000 tokens usage...")
memory_tracker.set_usage(session_id, 1000)
status = memory_tracker.check_status(session_id)
print(f"Status: {status.status}, Percentage: {status.percentage:.2f}%")
if status.status != "ok":
print("❌ FAILED: Should be 'ok'")
else:
print("✅ PASSED: Status is 'ok'")
# 3. Test Warning State
print("\n3. Testing Warning State (81%)...")
# Set usage to just above warning threshold
warning_val = int(KIMI_K2_CONTEXT_LENGTH * 0.81)
memory_tracker.set_usage(session_id, warning_val)
status = memory_tracker.check_status(session_id)
print(f"Current Usage: {warning_val}")
print(f"Status: {status.status}, Percentage: {status.percentage:.2f}%")
print(f"Message: {status.message}")
if status.status != "warning":
print("❌ FAILED: Should be 'warning'")
else:
print("✅ PASSED: Status is 'warning'")
# 4. Test Blocked State
print("\n4. Testing Blocked State (96%)...")
# Set usage to above block threshold
block_val = int(KIMI_K2_CONTEXT_LENGTH * 0.96)
memory_tracker.set_usage(session_id, block_val)
status = memory_tracker.check_status(session_id)
print(f"Current Usage: {block_val}")
print(f"Status: {status.status}, Percentage: {status.percentage:.2f}%")
print(f"Message: {status.message}")
if status.status != "blocked":
print("❌ FAILED: Should be 'blocked'")
else:
print("✅ PASSED: Status is 'blocked'")
# 5. Verify API Response (Logic simulation)
# We can't easily call the running API from here without successful auth/db setup
# unless we run this script in the same environment.
# But since we share the memory_tracker instance if running locally with same cache dir,
# we can verify the logic directly.
print("\n--- Test Complete ---")
print("To verify in UI:")
print(f"1. Start the app")
print(f"2. Send a message to session '{session_id}' (or any session)")
print(f"3. Use this script to set usage for that session ID high")
print(f"4. Refresh or send another message to see the effect")
if __name__ == "__main__":
asyncio.run(test_memory_limits())
|