llm_council/scripts/test_setup.py

#!/usr/bin/env python3
"""Setup a test conversation with optional message and documents."""

import os
import sys
import httpx
import time
from datetime import datetime
from pathlib import Path
from dotenv import load_dotenv
from urllib.parse import quote

# Load .env file if it exists
load_dotenv()

API_BASE = "http://localhost:8001"

def create_test_conversation():
    """Create a new conversation with today's date/time as title."""
    now = datetime.now()
    title = now.strftime("%Y-%m-%d %H:%M:%S")
    
    with httpx.Client(timeout=10.0) as client:
        # Create conversation
        response = client.post(f"{API_BASE}/api/conversations", json={})
        if response.status_code != 200:
            print(f"Error creating conversation: {response.text}", file=sys.stderr)
            sys.exit(1)
        
        conv = response.json()
        conv_id = conv["id"]
        
        # Update title
        response = client.patch(
            f"{API_BASE}/api/conversations/{conv_id}/title",
            json={"title": title}
        )
        if response.status_code != 200:
            print(f"Warning: Could not set title: {response.text}", file=sys.stderr)
        
        print(f"Created conversation: {conv_id}")
        print(f"Title: {title}")
        return conv_id

def upload_document(conv_id, filepath):
    """Upload a document to a conversation."""
    path = Path(filepath)
    if not path.exists():
        print(f"Warning: File not found: {filepath}", file=sys.stderr)
        return False
    
    with httpx.Client(timeout=30.0) as client:
        with open(path, 'rb') as f:
            files = {'file': (path.name, f.read(), 'text/markdown')}
            response = client.post(
                f"{API_BASE}/api/conversations/{conv_id}/documents",
                files=files
            )
            if response.status_code != 200:
                print(f"Warning: Could not upload {filepath}: {response.text}", file=sys.stderr)
                return False
    
    print(f"Uploaded: {path.name}")
    return True

def send_message(conv_id, message):
    """Send a message to a conversation."""
    # Use a very long timeout since the council process can take several minutes
    # Default to 10 minutes, but allow override via env var
    timeout_seconds = float(os.getenv("TEST_MESSAGE_TIMEOUT_SECONDS", "600.0"))
    
    print(f"Sending message (this may take several minutes, timeout: {timeout_seconds}s)...")
    with httpx.Client(timeout=timeout_seconds) as client:
        try:
            response = client.post(
                f"{API_BASE}/api/conversations/{conv_id}/message",
                json={"content": message}
            )
            if response.status_code != 200:
                print(f"Warning: Could not send message: {response.text}", file=sys.stderr)
                return False
            print(f"✓ Message sent and processed: {message[:50]}...")
            return True
        except httpx.ReadTimeout:
            print(f"Error: Request timed out after {timeout_seconds}s", file=sys.stderr)
            print("The council process is still running. You can check the conversation in the UI.", file=sys.stderr)
            print(f"Conversation ID: {conv_id}", file=sys.stderr)
            return False
        except httpx.RequestError as e:
            print(f"Error sending message: {e}", file=sys.stderr)
            return False

def main():
    # Wait for backend to be ready
    max_retries = 10
    for i in range(max_retries):
        try:
            with httpx.Client(timeout=1.0) as client:
                response = client.get(f"{API_BASE}/")
                if response.status_code == 200:
                    break
        except httpx.RequestError:
            if i < max_retries - 1:
                time.sleep(1)
            else:
                print(f"Error: Backend not available at {API_BASE}", file=sys.stderr)
                print("Make sure the backend is running: uv run python -m backend.main", file=sys.stderr)
                sys.exit(1)
    
    # Create conversation
    conv_id = create_test_conversation()
    
    # Upload documents if specified
    test_docs = os.getenv("TEST_DOCS", "")
    if test_docs:
        doc_paths = [p.strip() for p in test_docs.split(",") if p.strip()]
        for doc_path in doc_paths:
            upload_document(conv_id, doc_path)
    
    # Note: TEST_MESSAGE is NOT sent automatically
    # It's provided here for reference - user should type it in the UI
    test_message = os.getenv("TEST_MESSAGE", "")
    open_url = f"http://localhost:5173/?conversation={conv_id}"
    if test_message:
        open_url += f"&message={quote(test_message)}"
    
    print(f"\n✓ Conversation created: {conv_id}")
    print(f"CONVERSATION_ID={conv_id}")  # For Makefile to parse
    print(f"OPEN_URL={open_url}")  # For Makefile to parse
    print(f"Open in browser: {open_url}")
    
    if test_message:
        print(f"\n💡 Pre-filled message (copy/paste into input):")
        print(f"   {test_message}")
    
    return conv_id

if __name__ == "__main__":
    main()