agents/data_retrieval.py

def sample_data_retrieval_agent(state):
    """Retrieves sample data (3 rows) from each relevant table."""
    relevant_tables = state.get("relevant_tables", [])
    client = state["client"]
    
    if not relevant_tables:
        return {"sample_data": {}}
    
    if client is None:
        return {"sample_data": {"error": "Failed to connect to BigQuery."}}
    
    sample_data = {}
    
    for table in relevant_tables:
        try:
            # Clean table name if needed
            clean_table = table.replace("`", "").strip()
            
            # Query to get 3 sample rows
            query = f"""
            SELECT * FROM `{clean_table}` 
            LIMIT 3
            """
            
            query_job = client.query(query)
            results = query_job.result()
            
            # Convert results to a list of dictionaries
            rows = []
            for row in results:
                rows.append(dict(row.items()))
            
            sample_data[clean_table] = rows
            
        except Exception as e:
            sample_data[clean_table] = {"error": str(e)}
    
    return {"sample_data": sample_data}