test_research_assistant.py
1 import os 2 import pytest 3 import sys 4 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../..'))) 5 6 from examples.test_utils.get_trace_data import ( 7 run_command, 8 extract_information, 9 load_trace_data 10 ) 11 12 from examples.test_utils.get_components import ( 13 get_component_structure_and_sequence 14 ) 15 16 @pytest.mark.parametrize("model, provider, async_llm, syntax", [ 17 ("gpt-4o-mini", "openai", False, "chat"), 18 ("gemini-1.5-flash", "google_genai", False, "chat"), 19 # ("gemini-1.5-flash", "google_vertexai", False, "chat"), 20 # ("gpt-3.5-turbo", "azure", False, "chat"), 21 # ("gemini-1.5-flash", "anthropic", False, "chat"), 22 ]) 23 def test_research_assistant(model: str, provider: str, async_llm: bool, syntax: str): 24 # Build the command to run research_assistant.py with the provided arguments 25 command = f'python research_assistant.py --model {model} --provider {provider} --async_llm {async_llm} --syntax {syntax}' 26 cwd = os.path.dirname(os.path.abspath(__file__)) # Use the current directory 27 output = run_command(command, cwd=cwd) 28 29 # Extract trace file location from logs 30 locations = extract_information(output) 31 32 # Load and validate the trace data 33 data = load_trace_data(locations) 34 35 # Get component structure and sequence 36 component_sequence = get_component_structure_and_sequence(data) 37 38 # Print component sequence 39 print("Component sequence:", component_sequence) 40 41 # Validate component sequence 42 assert len(component_sequence) >= 2, f"Expected at least 2 components, got {len(component_sequence)}"