/ tests / examples / langgraph / personal_research_assistant / test_research_assistant.py
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)}"