agentlightning

import os from typing import TypedDict from agentlightning import rollout, Trainer, PromptTemplate, NamedResources, Rollout # Ensure OPENAI_API_KEY is set in your environment or replace 'YOUR_OPENAI_KEY' os.environ.setdefault("OPENAI_API_KEY", os.environ.get('OPENAI_API_KEY', 'YOUR_OPENAI_KEY')) class RoomSelectionTask(TypedDict): attendee_count: int time: str has_whiteboard: bool expected_choice: str def room_selection_grader(final_choice: str, expected_choice: str) -> float: """Grades the agent's room selection.""" return 1.0 if final_choice == expected_choice else 0.0 @rollout def room_selector_agent(task: RoomSelectionTask, prompt_template: PromptTemplate) -> float: # Simulate agent logic using the prompt_template and task # In a real scenario, this would involve LLM calls and tool usage. prompt_text = prompt_template.format(task=task) # Placeholder for LLM interaction and tool calls # For demonstration, we'll simulate a choice. if task['attendee_count'] > 5 and task['has_whiteboard']: final_choice = 'Large Conference Room with Whiteboard' else: final_choice = 'Small Meeting Room' reward = room_selection_grader(final_choice, task['expected_choice']) return reward # Define a simple prompt template (this would be optimized by Agent Lightning) initial_prompt = PromptTemplate( template="""You are a room selection agent. Given the following task: Attendees: {task[attendee_count]}, Time: {task[time]}, Whiteboard needed: {task[has_whiteboard]}. Select the best room.""" ) # Example usage with a dummy trainer (full training requires more setup) # For a complete training loop, you would typically define a dataset and an algorithm. # This snippet focuses on demonstrating the @rollout decorator. # Create a dummy task for a single rollout demonstration dummy_task = RoomSelectionTask( attendee_count=7, time="10 AM", has_whiteboard=True, expected_choice="Large Conference Room with Whiteboard" ) # Manually run the agent with the initial prompt for demonstration # In a real setup, Trainer would orchestrate this. print(f"Running agent with task: {dummy_task}") resources = NamedResources(prompt_template=initial_prompt) reward = room_selector_agent(dummy_task, resources) print(f"Agent received reward: {reward}") # To initialize a trainer (requires more setup including an algorithm and dataset): # trainer = Trainer(n_runners=1, algorithm=your_algorithm_instance) # trainer.fit(agent=room_selector_agent, tasks=[dummy_task], resources={'prompt_template': initial_prompt})