State Management in Agent Loops

class AgentState:
  def __init__(self, user_question):
    self.original_question = user_question
    self.sources_found = []
    self.facts_learned = {}
    self.response = ""
    self.iteration_count = 0
    self.total_cost = 0.0
  
  def add_source(self, url, content):
    self.sources_found.append(url)
    self.facts_learned[url] = content
  
  def add_to_response(self, text):
    self.response += text + " "
  
  def get_summary(self):
    return f"""
    Progress:
    - Original question: {self.original_question}
    - Sources found: {len(self.sources_found)}
    - Facts learned: {len(self.facts_learned)}
    - Iterations: {self.iteration_count}
    - Cost: ${self.total_cost:.2f}
    """

# Use the state
state = AgentState("What is renewable energy?")
state.add_source("https://example.com", "Article about renewable energy...")
print(state.get_summary())

def run_agent_with_state(user_question):
  state = AgentState(user_question)
  conversation = []
  max_iterations = 10
  max_cost = 1.0
  
  for iteration in range(max_iterations):
    state.iteration_count = iteration + 1
    
    # Ask model
    response = call_model(conversation)
    
    # Check if done
    if "<tool_use>" not in response:
      state.add_to_response(response)
      return state
    
    # Call tool
    tool_call = parse_tool_call(response)
    tool_name = tool_call["name"]
    result = call_tool(tool_name, tool_call["args"])
    cost = calculate_cost(tool_name, result)
    
    # Update state
    state.total_cost += cost
    if tool_name == "web_search":
      for item in result["results"]:
        state.sources_found.append(item["url"])
    
    # Check limits
    if state.total_cost > max_cost:
      return {"error": f"Cost limit exceeded: ${state.total_cost:.2f}"}
    
    # Add to conversation
    conversation.append({"role": "assistant", "content": response})
    conversation.append({"role": "user", "content": str(result)})
    
    print(state.get_summary())
  
  return state

class ResearchAgentState:
  def __init__(self, topic):
    self.topic = topic
    self.status = "searching"  # searching, reading, synthesizing, done
    
    # Search phase
    self.search_queries = []
    self.search_results = []
    
    # Read phase
    self.pages_fetched = {}
    self.key_facts = []
    
    # Synthesis phase
    self.draft = ""
    self.final_answer = ""
    
    # Metadata
    self.start_time = time.time()
    self.elapsed_seconds = 0
    self.tokens_used = 0
    self.cost = 0.0
  
  def move_to_next_phase(self):
    phases = ["searching", "reading", "synthesizing", "done"]
    current_index = phases.index(self.status)
    if current_index < len(phases) - 1:
      self.status = phases[current_index + 1]
  
  def update_metadata(self, tokens, cost):
    self.tokens_used += tokens
    self.cost += cost
    self.elapsed_seconds = time.time() - self.start_time

def detect_looping(state, previous_state):
  # If no new sources were found, might be looping
  if len(state.sources_found) == len(previous_state.sources_found):
    loop_count += 1
  else:
    loop_count = 0
  
  if loop_count >= 3:
    return True  # Looping detected
  
  return False

# In the agent loop
previous_state = None
for iteration in range(max_iterations):
  # ... run agent step ...
  
  if detect_looping(state, previous_state):
    print("Agent is looping, breaking out")
    break
  
  previous_state = state.copy()

import json

def save_state(state, filename):
  data = {
    "original_question": state.original_question,
    "sources_found": state.sources_found,
    "facts_learned": state.facts_learned,
    "iteration_count": state.iteration_count,
    "total_cost": state.total_cost,
  }
  with open(filename, "w") as f:
    json.dump(data, f)

def load_state(filename):
  with open(filename, "r") as f:
    data = json.load(f)
  state = AgentState(data["original_question"])
  state.sources_found = data["sources_found"]
  state.facts_learned = data["facts_learned"]
  state.iteration_count = data["iteration_count"]
  state.total_cost = data["total_cost"]
  return state

# Use checkpoints
if checkpoint_file_exists:
  state = load_state(checkpoint_file)
else:
  state = AgentState(user_question)

# ... run agent ...

save_state(state, "agent_state.json")