tests/test_agent.py

"""
Test script for the LangGraph agent pipeline.
Runs several queries with a short wait between them to verify the full flow.
Requires gemini_api in environment for real LLM calls; otherwise only tests prepare_generation (no API).
"""

import os
import time
from rag_engine import RAGEngine
from agent import build_agent_graph, run_stream


def main():
    print("Loading RAG Engine and building agent graph...")
    engine = RAGEngine()
    graph = build_agent_graph(engine)
    print("OK.\n")

    api_key = os.environ.get("gemini_api")
    if not api_key:
        print("⚠️  gemini_api not set. Testing only prepare_generation (no LLM calls).\n")
        test_queries = [
            "Tell me about the Audi RS3",
            "Compare Audi RS3 vs Hyundai Elantra N",
            "מה דעתך על BMW X5?",  # should trigger refusal
        ]
        for i, query in enumerate(test_queries, 1):
            print(f"--- Test {i}: prepare_generation ---")
            print(f"Query: {query!r}")
            refusal, sys_p, user_p, steps = engine.prepare_generation(query)
            if refusal:
                print(f"Refusal (expected for unsupported car): {refusal[:150]}...")
            else:
                print(f"Steps: {len(steps)}; system_prompt length: {len(sys_p or '')}; user_prompt length: {len(user_p or '')}")
            print()
        print("Done (prepare_generation only). Set gemini_api to run full agent.")
        return

    test_queries = [
        "Tell me about the Audi RS3",
        "Compare Audi RS3 vs Hyundai Elantra N",
        "מה היתרונות של קיה EV9?",
        "מה דעתך על BMW X5?",  # should trigger refusal (unsupported model)
    ]
    wait_seconds = 8

    for i, query in enumerate(test_queries, 1):
        print(f"--- Test {i}/{len(test_queries)} ---")
        print(f"Query: {query!r}")
        last_output = None
        step_count = 0
        try:
            for out in run_stream(engine, graph, query, api_key):
                last_output = out
                step_count += 1
            if last_output:
                preview = last_output[:400] + "..." if len(last_output) > 400 else last_output
                print(f"Steps yielded: {step_count}; final length: {len(last_output)}")
                print(f"Final preview:\n{preview}\n")
            else:
                print("No output received.\n")
        except Exception as e:
            print(f"Error: {e}\n")
        if i < len(test_queries):
            print(f"Waiting {wait_seconds}s before next query...")
            time.sleep(wait_seconds)

    print("All tests finished.")


if __name__ == "__main__":
    main()