82 lines
3.4 KiB
Python
82 lines
3.4 KiB
Python
import sys
|
|
import os
|
|
import time
|
|
from simulation.sim_base import BaseSimulation, run_sim
|
|
|
|
class ContextSimulation(BaseSimulation):
|
|
def run(self):
|
|
print("\n--- Running Context & Chat Simulation ---")
|
|
|
|
# 1. Test Discussion Creation
|
|
disc_name = f"TestDisc_{int(time.time())}"
|
|
print(f"[Sim] Creating discussion: {disc_name}")
|
|
self.sim.create_discussion(disc_name)
|
|
time.sleep(1)
|
|
|
|
# Verify it's in the list
|
|
session = self.client.get_session()
|
|
# The session structure usually has discussions listed somewhere, or we can check the listbox
|
|
# For now, we'll trust the click and check the session update
|
|
|
|
# 2. Test File Aggregation & Context Refresh
|
|
print("[Sim] Testing context refresh and token budget...")
|
|
proj = self.client.get_project()
|
|
# Add many files to ensure we cross the 1% threshold (~9000 tokens)
|
|
import glob
|
|
all_py = [os.path.basename(f) for f in glob.glob("*.py")]
|
|
for f in all_py:
|
|
if f not in proj['project']['files']['paths']:
|
|
proj['project']['files']['paths'].append(f)
|
|
|
|
# Update project via hook
|
|
self.client.post_project(proj['project'])
|
|
time.sleep(1)
|
|
|
|
# Trigger MD Only to refresh context and token budget
|
|
print("[Sim] Clicking MD Only...")
|
|
self.client.click("btn_md_only")
|
|
time.sleep(5)
|
|
|
|
# Verify status
|
|
proj_updated = self.client.get_project()
|
|
status = self.client.get_value("ai_status")
|
|
print(f"[Sim] Status: {status}")
|
|
assert "md written" in status, f"Expected 'md written' in status, got {status}"
|
|
|
|
# Verify token budget
|
|
pct = self.client.get_value("token_budget_pct")
|
|
current = self.client.get_value("token_budget_current")
|
|
print(f"[Sim] Token budget pct: {pct}, current={current}")
|
|
# We'll just warn if it's 0 but the MD was written, as it might be a small context
|
|
if pct == 0:
|
|
print("[Sim] WARNING: token_budget_pct is 0. This might be due to small context or estimation failure.")
|
|
|
|
# 3. Test Chat Turn
|
|
msg = "What is the current date and time? Answer in one sentence."
|
|
print(f"[Sim] Sending message: {msg}")
|
|
self.sim.run_discussion_turn(msg)
|
|
|
|
# 4. Verify History
|
|
print("[Sim] Verifying history...")
|
|
session = self.client.get_session()
|
|
entries = session.get('session', {}).get('entries', [])
|
|
|
|
# We expect at least 2 entries (User and AI)
|
|
assert len(entries) >= 2, f"Expected at least 2 entries, found {len(entries)}"
|
|
assert entries[-2]['role'] == 'User', "Expected second to last entry to be User"
|
|
assert entries[-1]['role'] == 'AI', "Expected last entry to be AI"
|
|
print(f"[Sim] AI responded: {entries[-1]['content'][:50]}...")
|
|
|
|
# 5. Test History Truncation
|
|
print("[Sim] Testing history truncation...")
|
|
self.sim.truncate_history(1)
|
|
time.sleep(1)
|
|
session = self.client.get_session()
|
|
entries = session.get('session', {}).get('entries', [])
|
|
# Truncating to 1 pair means 2 entries max (if it's already at 2, it might not change,
|
|
# but if we had more, it would).
|
|
assert len(entries) <= 2, f"Expected <= 2 entries after truncation, found {len(entries)}"
|
|
|
|
if __name__ == "__main__":
|
|
run_sim(ContextSimulation)
|