test_agent.py

import os
from dotenv import load_dotenv
from graph import AgentState, run_pm_agent, run_synthesis_agent

# Load environment variables from .env file
load_dotenv()

print("---🔬 Interactive Agent Test Bed ---")

# 1. Create a mock state to test the PM Agent
mock_state_for_pm = AgentState(
    userInput="How do I fine-tune a Llama-3 model on a custom dataset?",
    coreObjectivePrompt="Provide a detailed, step-by-step guide for fine-tuning a Llama-3 model on a custom dataset, including code examples and best practices.",
    retrievedMemory="Memory: Fine-tuning requires a powerful GPU and careful data preparation.",
    qaFeedback=None,
    execution_path=[]
)

print("\n--- Testing PM Agent ---")
pm_output = run_pm_agent(mock_state_for_pm)
pm_plan = pm_output.get('pmPlan', {})
print(f"PM Plan Generated: {pm_plan}")


# 2. Use the output from the PM test to test the Synthesis Agent
if pm_plan:
    mock_state_for_synthesis = AgentState(
        coreObjectivePrompt=mock_state_for_pm['coreObjectivePrompt'],
        pmPlan=pm_plan,
        experimentResults=None # Mocking that no experiment was run
    )

    print("\n--- Testing Synthesis Agent ---")
    synthesis_output = run_synthesis_agent(mock_state_for_synthesis)
    print(f"Synthesized Draft (first 300 chars): {synthesis_output['draftResponse'][:300]}...")