Session Management

[
    {"role": "user", "content": "Hello"},
    {"role": "assistant", "content": "Hi! How can I help?"},
    {"role": "user", "content": "What did I just say?"},
]

import asyncio
import json
from pathlib import Path

from dotenv import load_dotenv
from dedalus_labs import AsyncDedalus, DedalusRunner

load_dotenv()

SESSIONS_FILE = Path(__file__).parent / "sessions.json"

MODELS = [
    "openai/gpt-5.1",
    "anthropic/claude-opus-4-5-20251101",
    "google/gemini-3-pro-preview",
]


def load_sessions() -> dict:
    if SESSIONS_FILE.exists():
        return json.loads(SESSIONS_FILE.read_text())
    return {}


def save_sessions(sessions: dict):
    SESSIONS_FILE.write_text(json.dumps(sessions, indent=2))


def get_session(session_id: str) -> list[dict]:
    sessions = load_sessions()
    return sessions.get(session_id, [])


def save_session(session_id: str, messages: list[dict]):
    sessions = load_sessions()
    sessions[session_id] = messages
    save_sessions(sessions)


async def chat(
    session_id: str,
    user_input: str,
    model: str,
    mcp_servers: list[str] | None = None,
) -> str:
    client = AsyncDedalus()
    runner = DedalusRunner(client)

    history = get_session(session_id)

    # Append user message to history (runner ignores `input` when `messages` is passed)
    history.append({"role": "user", "content": user_input})

    kwargs = {
        "messages": history,
        "model": model,
    }
    if mcp_servers:
        kwargs["mcp_servers"] = mcp_servers

    result = await runner.run(**kwargs)
    save_session(session_id, result.to_input_list())

    return result.final_output


async def demo():
    print("=" * 60)
    print("  Dedalus Session Management Demo")
    print("=" * 60)
    print("\nCommands:")
    print("  /new <name>   - Start new session")
    print("  /list         - List sessions")
    print("  /load <name>  - Load session")
    print("  /clear        - Clear current session")
    print("  /model        - List available models")
    print("  /model <num>  - Switch model")
    print("  /mcp <url>    - Add MCP server")
    print("  /mcp clear    - Clear MCP servers")
    print("  /mcp          - List active MCP servers")
    print("  /status       - Show current config")
    print("  /quit         - Exit")
    print()

    current_session = "default"
    current_model = MODELS[0]
    mcp_servers: list[str] = []

    print(f"Session: {current_session}")
    print(f"Model: {current_model}")
    print()

    while True:
        try:
            user_input = input(f"[{current_session}] You: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nGoodbye!")
            break

        if not user_input:
            continue

        if user_input.startswith("/"):
            parts = user_input.split(maxsplit=1)
            cmd = parts[0].lower()
            arg = parts[1].strip() if len(parts) > 1 else ""

            if cmd == "/quit":
                print("Goodbye!")
                break

            elif cmd == "/new":
                current_session = arg or "default"
                save_session(current_session, [])
                print(f"Started new session: {current_session}")

            elif cmd == "/list":
                sessions = load_sessions()
                print(f"Sessions: {list(sessions.keys()) or ['(none)']}")

            elif cmd == "/load":
                current_session = arg or "default"
                history = get_session(current_session)
                print(f"Loaded session: {current_session} ({len(history)} messages)")

            elif cmd == "/clear":
                save_session(current_session, [])
                print(f"Cleared session: {current_session}")

            elif cmd == "/model":
                if not arg:
                    print("Available models:")
                    for i, m in enumerate(MODELS, 1):
                        marker = "*" if m == current_model else " "
                        print(f"  {marker} {i}. {m}")
                elif arg.isdigit() and 1 <= int(arg) <= len(MODELS):
                    current_model = MODELS[int(arg) - 1]
                    print(f"Switched to: {current_model}")
                else:
                    current_model = arg
                    print(f"Switched to: {current_model}")

            elif cmd == "/mcp":
                if not arg:
                    print(f"Active MCP servers: {mcp_servers or 'None'}")
                elif arg == "clear":
                    mcp_servers = []
                    print("Cleared MCP servers")
                else:
                    mcp_servers.append(arg)
                    print(f"Added MCP server: {arg}")

            elif cmd == "/status":
                print(f"Session: {current_session}")
                print(f"Model: {current_model}")
                print(f"MCP Servers: {mcp_servers or 'None'}")
                history = get_session(current_session)
                print(f"Messages: {len(history)}")

            else:
                print(f"Unknown command: {cmd}")

            continue

        print("Assistant: ", end="", flush=True)
        response = await chat(
            current_session,
            user_input,
            model=current_model,
            mcp_servers=mcp_servers if mcp_servers else None,
        )
        print(response)
        print()


if __name__ == "__main__":
    asyncio.run(demo())