76 lines
3.1 KiB
Python
76 lines
3.1 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 a file to paths (e.g., aggregate.py itself)
|
|
if "aggregate.py" not in proj['project']['files']['paths']:
|
|
proj['project']['files']['paths'].append("aggregate.py")
|
|
|
|
# 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(2)
|
|
|
|
# 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")
|
|
print(f"[Sim] Token budget pct: {pct}")
|
|
assert pct > 0, "Expected token_budget_pct > 0 after generation"
|
|
|
|
# 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)
|