scripts/test_reasoning_levels.py

#!/usr/bin/env python3
"""Load a locomo dataset into Honcho and test with configurable query and reasoning level."""

import argparse
import json
import os
import time
from datetime import datetime, timedelta, timezone

import httpx
from dotenv import load_dotenv

load_dotenv()

# Use environment variables with defaults matching .env.template
BASE_URL = os.getenv("HONCHO_BASE_URL")
REASONING_LEVELS = ["minimal", "low", "medium", "high", "max"]


def parse_datetime(dt_string: str) -> datetime:
    """Parse datetime string like '1:56 pm on 8 May, 2023' into datetime object."""
    parts = dt_string.split(" on ")
    time_part = parts[0]
    date_part = parts[1]

    time_obj = datetime.strptime(time_part, "%I:%M %p")
    date_obj = datetime.strptime(date_part, "%d %B, %Y")

    return datetime(
        year=date_obj.year,
        month=date_obj.month,
        day=date_obj.day,
        hour=time_obj.hour,
        minute=time_obj.minute,
        tzinfo=timezone.utc,
    )


def load_locomo(
    client: httpx.Client, filepath: str, workspace_id: str
) -> tuple[str, str]:
    """Load locomo dataset into Honcho. Returns (speaker_a, speaker_b)."""
    with open(filepath) as f:
        data = json.load(f)

    convo = data[0]["conversation"]
    speaker_a = convo["speaker_a"]
    speaker_b = convo["speaker_b"]

    print(f"Loading conversation between {speaker_a} and {speaker_b}")

    # Create workspace
    resp = client.post(f"{BASE_URL}/workspaces", json={"id": workspace_id})
    if resp.status_code >= 400:
        print(f"Failed to create workspace: {resp.status_code} {resp.text}")
        return "", ""
    print(f"Created workspace: {workspace_id}")

    # Create peers
    resp = client.post(
        f"{BASE_URL}/workspaces/{workspace_id}/peers", json={"id": speaker_a}
    )
    if resp.status_code >= 400:
        print(f"Failed to create peer {speaker_a}: {resp.status_code} {resp.text}")
        return "", ""
    resp = client.post(
        f"{BASE_URL}/workspaces/{workspace_id}/peers", json={"id": speaker_b}
    )
    if resp.status_code >= 400:
        print(f"Failed to create peer {speaker_b}: {resp.status_code} {resp.text}")
        return "", ""
    print(f"Created peers: {speaker_a}, {speaker_b}")

    session_num = 1
    while f"session_{session_num}" in convo:
        session_key = f"session_{session_num}"
        datetime_key = f"session_{session_num}_date_time"

        messages = convo[session_key]
        base_time = parse_datetime(convo[datetime_key])

        print(f"\n--- Session {session_num}: {convo[datetime_key]} ---")
        print(f"  {len(messages)} messages")

        session_id = f"locomo_session_{session_num}"

        # Create session
        resp = client.post(
            f"{BASE_URL}/workspaces/{workspace_id}/sessions",
            json={"id": session_id},
        )
        if resp.status_code >= 400:
            print(
                f"Failed to create session {session_id}: {resp.status_code} {resp.text}"
            )
            return "", ""

        # Add peers to session
        resp = client.post(
            f"{BASE_URL}/workspaces/{workspace_id}/sessions/{session_id}/peers",
            json={speaker_a: {}, speaker_b: {}},
        )
        if resp.status_code >= 400:
            print(f"Failed to add peers to session: {resp.status_code} {resp.text}")
            return "", ""
        print(f"  Created session: {session_id}")

        # Build message batch
        msg_batch = []
        for i, msg in enumerate(messages):
            msg_time = base_time + timedelta(seconds=i * 2)
            msg_batch.append(
                {
                    "peer_id": msg["speaker"],
                    "content": msg["text"],
                    "created_at": msg_time.isoformat(),
                }
            )

        # Create messages
        resp = client.post(
            f"{BASE_URL}/workspaces/{workspace_id}/sessions/{session_id}/messages",
            json={"messages": msg_batch},
        )
        if resp.status_code >= 400:
            print(f"Failed to create messages: {resp.status_code} {resp.text}")
            return "", ""
        print(f"  Loaded {len(messages)} messages")
        session_num += 1

    print(f"\nDone! Loaded {session_num - 1} sessions.")
    return speaker_a, speaker_b


def chat(
    client: httpx.Client, workspace_id: str, peer_id: str, query: str, level: str
) -> dict:
    """Call the chat endpoint with a specific reasoning level."""
    resp = client.post(
        f"{BASE_URL}/workspaces/{workspace_id}/peers/{peer_id}/chat",
        json={
            "query": query,
            "reasoning_level": level,
        },
    )
    if resp.status_code >= 400:
        return {"error": f"{resp.status_code} {resp.text}"}
    return resp.json()


def main():
    parser = argparse.ArgumentParser(
        description="Load a locomo dataset into Honcho and test with a query."
    )
    parser.add_argument(
        "filepath",
        type=str,
        help="Path to the locomo JSON file",
    )
    parser.add_argument(
        "--workspace",
        "-w",
        type=str,
        default=None,
        help="Workspace ID (default: auto-generated from timestamp)",
    )
    parser.add_argument(
        "--query",
        "-q",
        type=str,
        default="What do you know about this person?",
        help="The query to send to the chat endpoint",
    )
    parser.add_argument(
        "--peer",
        "-p",
        type=str,
        default=None,
        help="The peer ID to query (default: first speaker from dataset)",
    )
    parser.add_argument(
        "--level",
        "-l",
        type=str,
        choices=REASONING_LEVELS,
        default="medium",
        help="Reasoning level to use (default: medium)",
    )
    parser.add_argument(
        "--skip-load",
        action="store_true",
        help="Skip loading data, just run the query (requires --workspace and --peer)",
    )
    args = parser.parse_args()

    if not BASE_URL:
        print(
            "Error: HONCHO_BASE_URL is not set. Please set it in your environment or .env."
        )
        return

    # Generate workspace ID if not provided
    workspace_id = (
        args.workspace or f"locomo_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
    )

    with httpx.Client(timeout=None) as client:
        if args.skip_load:
            if not args.workspace or not args.peer:
                print(
                    "Error: --skip-load requires --workspace and --peer to be specified"
                )
                return
            speaker_a = args.peer
        else:
            # Load the dataset
            speaker_a, speaker_b = load_locomo(client, args.filepath, workspace_id)
            if not speaker_a:
                return

            print(f"\nPeers available: {speaker_a}, {speaker_b}")

        # Determine which peer to query
        peer_id = args.peer or speaker_a

        print("\n" + "=" * 60)
        print("Testing chat endpoint")
        print("=" * 60)
        print(f"Workspace: {workspace_id}")
        print(f"Peer: {peer_id}")
        print(f"Query: {args.query}")
        print(f"Level: {args.level}")
        print("=" * 60)

        start_time = time.time()
        result = chat(client, workspace_id, peer_id, args.query, args.level)
        elapsed = time.time() - start_time

        print(f"\nTime: {elapsed:.2f}s")

        if "error" in result:
            print(f"Error: {result['error']}")
        else:
            content = result.get("content", "")
            print(f"\nResponse ({len(content)} chars):")
            print("-" * 60)
            print(content)
            print("-" * 60)


if __name__ == "__main__":
    main()