scripts/demo_queries.py

from __future__ import annotations

# File Objective  : Run all four illustrative demo queries against the Vantage agent.
# Scope           : Dev script — single-thread session demonstrating pipeline requirements.
# What it does    : Runs 4 queries, shows final response first, then tools called and timing.
# What it does not: Persist memory beyond the process or expose HTTP.

import time
from langchain_core.messages import AIMessage
from vantage_core.agent.graph import build_agent
from vantage_core.log import get_logger, timed

log = get_logger(__name__)

_THREAD = {"configurable": {"thread_id": "client-sensing-demo"}, "recursion_limit": 50}

_QUERIES = [
    (
        "Illustrative Query 1",
        "What topics related to AI in healthcare have surged recently?",
    ),
    (
        "Illustrative Query 2",
        "What are key issues APAC manufacturers faced this quarter?",
    ),
    (
        "Illustrative Query 3",
        "Which topics are under-covered versus market interest?",
    ),
    (
        "Illustrative Query 4",
        "Summarize recent signals around supply-chain resilience.",
    ),
]


def _tools_called(messages: list, start_idx: int) -> list[str]:
    """Extract unique tool names from intermediate AIMessages in the current turn."""
    return list(dict.fromkeys(
        tc["name"]
        for m in messages[start_idx:]
        if isinstance(m, AIMessage)
        for tc in (m.tool_calls or [])
    ))


def main() -> None:
    print("⏳  Loading models + connecting adapters (≈20s)…", flush=True)
    with timed("build_agent (models + adapters)", log):
        agent = build_agent()
    print("✓  Agent ready.\n", flush=True)
    results: list[tuple[str, list[str], float]] = []
    
    # Keep track of initial messages in thread (0 to start)
    messages_count = 0

    for label, query in _QUERIES:
        print(f"\n{'─'*70}")
        print(f"  {label}")
        print(f"  Q: {query}")
        print(f"{'─'*70}\n")

        log.info("query start  label=%r", label)
        t0  = time.time()
        with timed(f"agent.invoke  {label}", log):
            out = agent.invoke({"messages": [("user", query)]}, _THREAD)
        elapsed = time.time() - t0

        all_msgs = out["messages"]
        tools = _tools_called(all_msgs, messages_count)
        answer = all_msgs[-1].content

        print(answer)
        print(f"\n{'─'*70}")
        print(f"  ⚙  Tools called : {tools}")
        print(f"  ⏱  Elapsed      : {elapsed:.1f}s")
        print(f"{'─'*70}")
        results.append((label, tools, elapsed))
        
        messages_count = len(all_msgs)

    # ── Summary table ──────────────────────────────────────────────────────────
    print(f"\n\n{'═'*70}")
    print("  DEMO SUMMARY")
    print(f"{'═'*70}")
    for label, tools, elapsed in results:
        status = "✓" if tools else "✗ (no tool called)"
        print(f"  {status}  {label:<40}  {tools}  ({elapsed:.1f}s)")
    print()


if __name__ == "__main__":
    main()