gistfile1.txt

from agentlightning import Algorithm, LightningStore, Dataset 
from agentlightning.types import Rollout, Span, ResourcesUpdate 
 
class MyCustomAlgorithm(Algorithm): 
 """Custom algorithm that learns from agent traces.""" 
 
 async def fit( 
 self, 
 train_dataset: Dataset, 
 val_dataset: Dataset, 
 store: LightningStore 
 ) -> None: 
 # 1. Enqueue initial rollouts 
 for task in train_dataset: 
 await store.enqueue_rollout( 
 input=task, 
 resource_keys=["main_llm", "main_prompt"] 
 ) 
 
 # 2. Wait for rollouts to complete 
 rollout_ids = [r.rollout_id for r in await store.query_rollouts(…)] 
 await store.wait_for_rollouts(rollout_ids) 
 
 # 3. Query spans and learn 
 for rollout_id in rollout_ids: 
 spans = await store.query_spans(rollout_id=rollout_id) 
 # Transform spans to training data 
 training_data = self.adapter.adapt(spans) 
 # Update model/prompts 
 updated_resources = self.learn(training_data) 
 await store.update_resources(updated_resources)