mirror of
https://github.com/run-llama/llama_cloud_services.git
synced 2026-07-19 16:43:32 -04:00
Compare commits
6 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 853aecc625 | |||
| 918190d7de | |||
| 346b1f1f86 | |||
| 145bf3476c | |||
| 6118111691 | |||
| 018dae3166 |
@@ -34,7 +34,7 @@ repos:
|
||||
rev: v1.0.1
|
||||
hooks:
|
||||
- id: mypy
|
||||
exclude: ^py/tests|^py/unit_tests
|
||||
exclude: ^py/tests|^py/unit_tests|^examples
|
||||
additional_dependencies:
|
||||
[
|
||||
"types-requests",
|
||||
|
||||
@@ -0,0 +1,73 @@
|
||||
This project uses LlamaSheets to extract data from spreadsheets for analysis.
|
||||
|
||||
## Current Project Structure
|
||||
|
||||
- `data/` - Contains extracted parquet files from LlamaSheets
|
||||
- `{name}_region_{N}.parquet` - Table data files
|
||||
- `{name}_metadata_{N}.parquet` - Cell metadata files
|
||||
- `{name}_job_metadata.json` - Extraction job information
|
||||
- `scripts/` - Analysis and helper scripts
|
||||
- `reports/` - Your generated reports and outputs
|
||||
|
||||
## Working with LlamaSheets Data
|
||||
|
||||
### Understanding the Files
|
||||
|
||||
When a spreadsheet is extracted, you'll find:
|
||||
|
||||
1. **Table parquet files** (`region_*.parquet`): The actual table data
|
||||
- Columns correspond to spreadsheet columns
|
||||
- Data types are preserved (dates, numbers, strings, booleans)
|
||||
|
||||
2. **Metadata parquet files** (`metadata_*.parquet`): Rich cell-level metadata
|
||||
- Formatting: `font_bold`, `font_italic`, `font_size`, `background_color_rgb`
|
||||
- Position: `row_number`, `column_number`, `coordinate` (e.g., "A1")
|
||||
- Type detection: `data_type`, `is_date_like`, `is_percentage`, `is_currency`
|
||||
- Layout: `is_in_first_row`, `is_merged_cell`, `horizontal_alignment`
|
||||
- Content: `cell_value`, `raw_cell_value`
|
||||
|
||||
3. **Job metadata JSON** (`job_metadata.json`): Overall extraction results
|
||||
- `regions[]`: List of extracted regions with IDs, locations, and titles/descriptions
|
||||
- `worksheet_metadata[]`: Generated titles and descriptions
|
||||
- `status`: Success/failure status
|
||||
|
||||
### Key Principles
|
||||
|
||||
1. **Use metadata to understand structure**: Bold cells often indicate headers, colors indicate groupings
|
||||
2. **Validate before analysis**: Check data types, look for missing values
|
||||
3. **Preserve formatting context**: The metadata tells you what the spreadsheet author emphasized
|
||||
4. **Save intermediate results**: Store cleaned data as new parquet files
|
||||
|
||||
### Common Patterns
|
||||
|
||||
**Loading data:**
|
||||
```python
|
||||
import pandas as pd
|
||||
|
||||
df = pd.read_parquet("data/region_1_Sheet1.parquet")
|
||||
meta_df = pd.read_parquet("data/metadata_1_Sheet1.parquet")
|
||||
```
|
||||
|
||||
**Finding headers:**
|
||||
```python
|
||||
headers = meta_df[meta_df["font_bold"] == True]["cell_value"].tolist()
|
||||
```
|
||||
|
||||
**Finding date columns:**
|
||||
```python
|
||||
date_cols = meta_df[meta_df["is_date_like"] == True]["column_number"].unique()
|
||||
```
|
||||
|
||||
## Tools Available
|
||||
|
||||
- **Python 3.11+**: For data analysis
|
||||
- **pandas**: DataFrame manipulation
|
||||
- **pyarrow**: Parquet file reading
|
||||
- **matplotlib**: Visualization (optional)
|
||||
|
||||
## Guidelines
|
||||
|
||||
- Always read the job_metadata.json first to understand what was extracted
|
||||
- Check both table data and metadata before making assumptions
|
||||
- Write reusable functions for common operations
|
||||
- Document any data quality issues discovered
|
||||
@@ -0,0 +1,278 @@
|
||||
"""
|
||||
Generate sample spreadsheets for LlamaSheets + Claude workflows.
|
||||
|
||||
This script creates example Excel files that demonstrate different use cases:
|
||||
1. Simple data table (for Workflow 1)
|
||||
2. Regional sales data (for Workflow 2)
|
||||
3. Complex budget with formatting (for Workflow 3)
|
||||
4. Weekly sales report (for Workflow 4)
|
||||
|
||||
Usage:
|
||||
python generate_sample_data.py
|
||||
"""
|
||||
|
||||
import random
|
||||
from datetime import datetime, timedelta
|
||||
from pathlib import Path
|
||||
|
||||
import pandas as pd
|
||||
from openpyxl import Workbook
|
||||
from openpyxl.styles import Font, PatternFill, Alignment
|
||||
|
||||
|
||||
def generate_workflow_1_data(output_dir: Path) -> None:
|
||||
"""Generate simple financial report for Workflow 1."""
|
||||
print("📊 Generating Workflow 1: financial_report_q1.xlsx")
|
||||
|
||||
# Create sample quarterly data
|
||||
months = ["January", "February", "March"]
|
||||
categories = ["Revenue", "Cost of Goods Sold", "Operating Expenses", "Net Income"]
|
||||
|
||||
data = []
|
||||
for category in categories:
|
||||
row: dict[str, str | int] = {"Category": category}
|
||||
for month in months:
|
||||
if category == "Revenue":
|
||||
value = random.randint(80000, 120000)
|
||||
elif category == "Cost of Goods Sold":
|
||||
value = random.randint(30000, 50000)
|
||||
elif category == "Operating Expenses":
|
||||
value = random.randint(20000, 35000)
|
||||
else: # Net Income
|
||||
value = int(
|
||||
int(row.get("January", 0))
|
||||
+ int(row.get("February", 0))
|
||||
+ int(row.get("March", 0))
|
||||
)
|
||||
value = random.randint(15000, 40000)
|
||||
row[month] = value
|
||||
data.append(row)
|
||||
|
||||
df = pd.DataFrame(data)
|
||||
|
||||
# Write to Excel
|
||||
output_file = output_dir / "financial_report_q1.xlsx"
|
||||
with pd.ExcelWriter(output_file, engine="openpyxl") as writer:
|
||||
df.to_excel(writer, sheet_name="Q1 Summary", index=False)
|
||||
|
||||
# Format it nicely
|
||||
worksheet = writer.sheets["Q1 Summary"]
|
||||
for cell in worksheet[1]: # Header row
|
||||
cell.font = Font(bold=True)
|
||||
cell.fill = PatternFill(
|
||||
start_color="4F81BD", end_color="4F81BD", fill_type="solid"
|
||||
)
|
||||
cell.font = Font(color="FFFFFF", bold=True)
|
||||
|
||||
print(f" ✅ Created {output_file}")
|
||||
|
||||
|
||||
def generate_workflow_2_data(output_dir: Path) -> None:
|
||||
"""Generate regional sales data for Workflow 2."""
|
||||
print("\n📊 Generating Workflow 2: Regional sales data")
|
||||
|
||||
regions = ["northeast", "southeast", "west"]
|
||||
products = ["Widget A", "Widget B", "Widget C", "Gadget X", "Gadget Y"]
|
||||
|
||||
for region in regions:
|
||||
data = []
|
||||
start_date = datetime(2024, 1, 1)
|
||||
|
||||
# Generate 90 days of sales data
|
||||
for day in range(90):
|
||||
date = start_date + timedelta(days=day)
|
||||
# Random number of sales per day (3-8)
|
||||
for _ in range(random.randint(3, 8)):
|
||||
product = random.choice(products)
|
||||
units_sold = random.randint(1, 20)
|
||||
price_per_unit = random.randint(50, 200)
|
||||
revenue = units_sold * price_per_unit
|
||||
|
||||
data.append(
|
||||
{
|
||||
"Date": date.strftime("%Y-%m-%d"),
|
||||
"Product": product,
|
||||
"Units_Sold": units_sold,
|
||||
"Revenue": revenue,
|
||||
}
|
||||
)
|
||||
|
||||
df = pd.DataFrame(data)
|
||||
|
||||
# Write to Excel
|
||||
output_file = output_dir / f"sales_{region}.xlsx"
|
||||
df.to_excel(output_file, sheet_name="Sales", index=False)
|
||||
print(f" ✅ Created {output_file} ({len(df)} rows)")
|
||||
|
||||
|
||||
def generate_workflow_3_data(output_dir: Path) -> None:
|
||||
"""Generate complex budget spreadsheet with formatting for Workflow 3."""
|
||||
print("\n📊 Generating Workflow 3: company_budget_2024.xlsx")
|
||||
|
||||
wb = Workbook()
|
||||
ws = wb.active
|
||||
ws.title = "Budget"
|
||||
|
||||
# Define departments with colors
|
||||
departments = {
|
||||
"Engineering": "C6E0B4",
|
||||
"Marketing": "FFD966",
|
||||
"Sales": "F4B084",
|
||||
"Operations": "B4C7E7",
|
||||
}
|
||||
|
||||
# Define categories
|
||||
categories = {
|
||||
"Personnel": ["Salaries", "Benefits", "Training"],
|
||||
"Infrastructure": ["Office Rent", "Equipment", "Software Licenses"],
|
||||
"Operations": ["Travel", "Supplies", "Miscellaneous"],
|
||||
}
|
||||
|
||||
# Styles
|
||||
header_font = Font(bold=True, size=12)
|
||||
category_font = Font(bold=True, size=11)
|
||||
|
||||
row = 1
|
||||
|
||||
# Title
|
||||
ws.merge_cells(f"A{row}:E{row}")
|
||||
ws[f"A{row}"] = "2024 Annual Budget"
|
||||
ws[f"A{row}"].font = Font(bold=True, size=14)
|
||||
ws[f"A{row}"].alignment = Alignment(horizontal="center")
|
||||
row += 2
|
||||
|
||||
# Headers
|
||||
ws[f"A{row}"] = "Category"
|
||||
ws[f"B{row}"] = "Item"
|
||||
for i, dept in enumerate(departments.keys()):
|
||||
ws.cell(row, 3 + i, dept)
|
||||
ws.cell(row, 3 + i).font = header_font
|
||||
|
||||
for cell in ws[row]:
|
||||
cell.font = header_font
|
||||
row += 1
|
||||
|
||||
# Data
|
||||
for category, items in categories.items():
|
||||
# Category header (bold)
|
||||
ws[f"A{row}"] = category
|
||||
ws[f"A{row}"].font = category_font
|
||||
row += 1
|
||||
|
||||
# Items with department budgets
|
||||
for item in items:
|
||||
ws[f"A{row}"] = ""
|
||||
ws[f"B{row}"] = item
|
||||
|
||||
# Add budget amounts for each department (with color)
|
||||
for i, (dept, color) in enumerate(departments.items()):
|
||||
amount = random.randint(5000, 50000)
|
||||
cell = ws.cell(row, 3 + i, amount)
|
||||
cell.fill = PatternFill(
|
||||
start_color=color, end_color=color, fill_type="solid"
|
||||
)
|
||||
cell.number_format = "$#,##0"
|
||||
|
||||
row += 1
|
||||
|
||||
row += 1 # Blank row between categories
|
||||
|
||||
# Adjust column widths
|
||||
ws.column_dimensions["A"].width = 20
|
||||
ws.column_dimensions["B"].width = 25
|
||||
for i in range(len(departments)):
|
||||
ws.column_dimensions[chr(67 + i)].width = 15 # C, D, E, F
|
||||
|
||||
output_file = output_dir / "company_budget_2024.xlsx"
|
||||
wb.save(output_file)
|
||||
print(f" ✅ Created {output_file}")
|
||||
print(" • Bold categories, colored departments, merged title cell")
|
||||
|
||||
|
||||
def generate_workflow_4_data(output_dir: Path) -> None:
|
||||
"""Generate weekly sales report for Workflow 4."""
|
||||
print("\n📊 Generating Workflow 4: sales_weekly.xlsx")
|
||||
|
||||
products = [
|
||||
"Product A",
|
||||
"Product B",
|
||||
"Product C",
|
||||
"Product D",
|
||||
"Product E",
|
||||
"Product F",
|
||||
"Product G",
|
||||
"Product H",
|
||||
]
|
||||
|
||||
# Generate one week of data
|
||||
data = []
|
||||
start_date = datetime(2024, 11, 4) # Monday
|
||||
|
||||
for day in range(7):
|
||||
date = start_date + timedelta(days=day)
|
||||
# Each product has 3-10 transactions per day
|
||||
for product in products:
|
||||
for _ in range(random.randint(3, 10)):
|
||||
units = random.randint(1, 15)
|
||||
price = random.randint(20, 150)
|
||||
revenue = units * price
|
||||
|
||||
data.append(
|
||||
{
|
||||
"Date": date.strftime("%Y-%m-%d"),
|
||||
"Product": product,
|
||||
"Units": units,
|
||||
"Revenue": revenue,
|
||||
}
|
||||
)
|
||||
|
||||
df = pd.DataFrame(data)
|
||||
|
||||
# Write to Excel with some formatting
|
||||
output_file = output_dir / "sales_weekly.xlsx"
|
||||
with pd.ExcelWriter(output_file, engine="openpyxl") as writer:
|
||||
df.to_excel(writer, sheet_name="Weekly Sales", index=False)
|
||||
|
||||
# Format header
|
||||
worksheet = writer.sheets["Weekly Sales"]
|
||||
for cell in worksheet[1]:
|
||||
cell.font = Font(bold=True)
|
||||
|
||||
print(f" ✅ Created {output_file} ({len(df)} rows)")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Generate all sample data files."""
|
||||
print("=" * 60)
|
||||
print("Generating Sample Data for LlamaSheets + Coding Agent Workflows")
|
||||
print("=" * 60)
|
||||
|
||||
# Create output directory
|
||||
output_dir = Path("input_data")
|
||||
output_dir.mkdir(exist_ok=True)
|
||||
|
||||
# Generate data for each workflow
|
||||
generate_workflow_1_data(output_dir)
|
||||
generate_workflow_2_data(output_dir)
|
||||
generate_workflow_3_data(output_dir)
|
||||
generate_workflow_4_data(output_dir)
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("✅ All sample data generated!")
|
||||
print("=" * 60)
|
||||
print(f"\nFiles created in {output_dir.absolute()}:")
|
||||
print("\nWorkflow 1 (Understanding a New Spreadsheet):")
|
||||
print(" • financial_report_q1.xlsx")
|
||||
print("\nWorkflow 2 (Generating Analysis Scripts):")
|
||||
print(" • sales_northeast.xlsx")
|
||||
print(" • sales_southeast.xlsx")
|
||||
print(" • sales_west.xlsx")
|
||||
print("\nWorkflow 3 (Using Cell Metadata):")
|
||||
print(" • company_budget_2024.xlsx")
|
||||
print("\nWorkflow 4 (Complete Automation):")
|
||||
print(" • sales_weekly.xlsx")
|
||||
print("\nYou can now use these files with the workflows in the documentation!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,5 @@
|
||||
llama-cloud-services # LlamaSheets SDK
|
||||
pandas>=2.0.0
|
||||
pyarrow>=12.0.0
|
||||
openpyxl>=3.0.0 # For Excel file support
|
||||
matplotlib>=3.7.0 # For visualizations (optional)
|
||||
@@ -0,0 +1,100 @@
|
||||
"""Helper script to extract spreadsheets using LlamaSheets."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import dotenv
|
||||
from pathlib import Path
|
||||
|
||||
from llama_cloud_services.beta.sheets import LlamaSheets
|
||||
from llama_cloud_services.beta.sheets.types import (
|
||||
SpreadsheetParsingConfig,
|
||||
SpreadsheetResultType,
|
||||
)
|
||||
|
||||
dotenv.load_dotenv()
|
||||
|
||||
|
||||
async def extract_spreadsheet(
|
||||
file_path: str, output_dir: str = "data", generate_metadata: bool = True
|
||||
) -> dict:
|
||||
"""Extract a spreadsheet using LlamaSheets."""
|
||||
|
||||
client = LlamaSheets(
|
||||
base_url="https://api.cloud.llamaindex.ai",
|
||||
api_key=os.getenv("LLAMA_CLOUD_API_KEY"),
|
||||
)
|
||||
|
||||
print(f"Extracting {file_path}...")
|
||||
|
||||
# Extract regions
|
||||
config = SpreadsheetParsingConfig(
|
||||
sheet_names=None, # Extract all sheets
|
||||
generate_additional_metadata=generate_metadata,
|
||||
)
|
||||
|
||||
job_result = await client.aextract_regions(file_path, config=config)
|
||||
|
||||
print(f"Extracted {len(job_result.regions)} region(s)")
|
||||
|
||||
# Create output directory
|
||||
output_path = Path(output_dir)
|
||||
output_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Get base name for files
|
||||
base_name = Path(file_path).stem
|
||||
|
||||
# Save job metadata
|
||||
job_metadata_path = output_path / f"{base_name}_job_metadata.json"
|
||||
with open(job_metadata_path, "w") as f:
|
||||
json.dump(job_result.model_dump(mode="json"), f, indent=2)
|
||||
print(f"Saved job metadata to {job_metadata_path}")
|
||||
|
||||
# Download each region
|
||||
for idx, region in enumerate(job_result.regions, 1):
|
||||
sheet_name = region.sheet_name.replace(" ", "_")
|
||||
|
||||
# Download region data
|
||||
region_bytes = await client.adownload_region_result(
|
||||
job_id=job_result.id,
|
||||
region_id=region.region_id,
|
||||
result_type=region.region_type,
|
||||
)
|
||||
|
||||
region_path = output_path / f"{base_name}_region_{idx}_{sheet_name}.parquet"
|
||||
with open(region_path, "wb") as f:
|
||||
f.write(region_bytes)
|
||||
print(f" Table {idx}: {region_path}")
|
||||
|
||||
# Download metadata
|
||||
metadata_bytes = await client.adownload_region_result(
|
||||
job_id=job_result.id,
|
||||
region_id=region.region_id,
|
||||
result_type=SpreadsheetResultType.CELL_METADATA,
|
||||
)
|
||||
|
||||
metadata_path = output_path / f"{base_name}_metadata_{idx}_{sheet_name}.parquet"
|
||||
with open(metadata_path, "wb") as f:
|
||||
f.write(metadata_bytes)
|
||||
print(f" Metadata {idx}: {metadata_path}")
|
||||
|
||||
print(f"\nAll files saved to {output_path}/")
|
||||
|
||||
return job_result.model_dump(mode="json")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
if len(sys.argv) < 2:
|
||||
print("Usage: python scripts/extract.py <spreadsheet_file>")
|
||||
sys.exit(1)
|
||||
|
||||
file_path = sys.argv[1]
|
||||
|
||||
if not Path(file_path).exists():
|
||||
print(f"❌ File not found: {file_path}")
|
||||
sys.exit(1)
|
||||
|
||||
result = asyncio.run(extract_spreadsheet(file_path))
|
||||
print(f"\n✅ Extraction complete! Job ID: {result['id']}")
|
||||
@@ -0,0 +1,278 @@
|
||||
"""
|
||||
Generate sample spreadsheets for LlamaSheets + LlamaIndex Agent workflows.
|
||||
|
||||
This script creates example Excel files that demonstrate different use cases:
|
||||
1. Simple data table (for Workflow 1)
|
||||
2. Regional sales data (for Workflow 2)
|
||||
3. Complex budget with formatting (for Workflow 3)
|
||||
4. Weekly sales report (for Workflow 4)
|
||||
|
||||
Usage:
|
||||
python generate_sample_data.py
|
||||
"""
|
||||
|
||||
import random
|
||||
from datetime import datetime, timedelta
|
||||
from pathlib import Path
|
||||
|
||||
import pandas as pd
|
||||
from openpyxl import Workbook
|
||||
from openpyxl.styles import Font, PatternFill, Alignment
|
||||
|
||||
|
||||
def generate_workflow_1_data(output_dir: Path) -> None:
|
||||
"""Generate simple financial report for Workflow 1."""
|
||||
print("📊 Generating Workflow 1: financial_report_q1.xlsx")
|
||||
|
||||
# Create sample quarterly data
|
||||
months = ["January", "February", "March"]
|
||||
categories = ["Revenue", "Cost of Goods Sold", "Operating Expenses", "Net Income"]
|
||||
|
||||
data = []
|
||||
for category in categories:
|
||||
row: dict[str, str | int] = {"Category": category}
|
||||
for month in months:
|
||||
if category == "Revenue":
|
||||
value = random.randint(80000, 120000)
|
||||
elif category == "Cost of Goods Sold":
|
||||
value = random.randint(30000, 50000)
|
||||
elif category == "Operating Expenses":
|
||||
value = random.randint(20000, 35000)
|
||||
else: # Net Income
|
||||
value = int(
|
||||
int(row.get("January", 0))
|
||||
+ int(row.get("February", 0))
|
||||
+ int(row.get("March", 0))
|
||||
)
|
||||
value = random.randint(15000, 40000)
|
||||
row[month] = value
|
||||
data.append(row)
|
||||
|
||||
df = pd.DataFrame(data)
|
||||
|
||||
# Write to Excel
|
||||
output_file = output_dir / "financial_report_q1.xlsx"
|
||||
with pd.ExcelWriter(output_file, engine="openpyxl") as writer:
|
||||
df.to_excel(writer, sheet_name="Q1 Summary", index=False)
|
||||
|
||||
# Format it nicely
|
||||
worksheet = writer.sheets["Q1 Summary"]
|
||||
for cell in worksheet[1]: # Header row
|
||||
cell.font = Font(bold=True)
|
||||
cell.fill = PatternFill(
|
||||
start_color="4F81BD", end_color="4F81BD", fill_type="solid"
|
||||
)
|
||||
cell.font = Font(color="FFFFFF", bold=True)
|
||||
|
||||
print(f" ✅ Created {output_file}")
|
||||
|
||||
|
||||
def generate_workflow_2_data(output_dir: Path) -> None:
|
||||
"""Generate regional sales data for Workflow 2."""
|
||||
print("\n📊 Generating Workflow 2: Regional sales data")
|
||||
|
||||
regions = ["northeast", "southeast", "west"]
|
||||
products = ["Widget A", "Widget B", "Widget C", "Gadget X", "Gadget Y"]
|
||||
|
||||
for region in regions:
|
||||
data = []
|
||||
start_date = datetime(2024, 1, 1)
|
||||
|
||||
# Generate 90 days of sales data
|
||||
for day in range(90):
|
||||
date = start_date + timedelta(days=day)
|
||||
# Random number of sales per day (3-8)
|
||||
for _ in range(random.randint(3, 8)):
|
||||
product = random.choice(products)
|
||||
units_sold = random.randint(1, 20)
|
||||
price_per_unit = random.randint(50, 200)
|
||||
revenue = units_sold * price_per_unit
|
||||
|
||||
data.append(
|
||||
{
|
||||
"Date": date.strftime("%Y-%m-%d"),
|
||||
"Product": product,
|
||||
"Units_Sold": units_sold,
|
||||
"Revenue": revenue,
|
||||
}
|
||||
)
|
||||
|
||||
df = pd.DataFrame(data)
|
||||
|
||||
# Write to Excel
|
||||
output_file = output_dir / f"sales_{region}.xlsx"
|
||||
df.to_excel(output_file, sheet_name="Sales", index=False)
|
||||
print(f" ✅ Created {output_file} ({len(df)} rows)")
|
||||
|
||||
|
||||
def generate_workflow_3_data(output_dir: Path) -> None:
|
||||
"""Generate complex budget spreadsheet with formatting for Workflow 3."""
|
||||
print("\n📊 Generating Workflow 3: company_budget_2024.xlsx")
|
||||
|
||||
wb = Workbook()
|
||||
ws = wb.active
|
||||
ws.title = "Budget"
|
||||
|
||||
# Define departments with colors
|
||||
departments = {
|
||||
"Engineering": "C6E0B4",
|
||||
"Marketing": "FFD966",
|
||||
"Sales": "F4B084",
|
||||
"Operations": "B4C7E7",
|
||||
}
|
||||
|
||||
# Define categories
|
||||
categories = {
|
||||
"Personnel": ["Salaries", "Benefits", "Training"],
|
||||
"Infrastructure": ["Office Rent", "Equipment", "Software Licenses"],
|
||||
"Operations": ["Travel", "Supplies", "Miscellaneous"],
|
||||
}
|
||||
|
||||
# Styles
|
||||
header_font = Font(bold=True, size=12)
|
||||
category_font = Font(bold=True, size=11)
|
||||
|
||||
row = 1
|
||||
|
||||
# Title
|
||||
ws.merge_cells(f"A{row}:E{row}")
|
||||
ws[f"A{row}"] = "2024 Annual Budget"
|
||||
ws[f"A{row}"].font = Font(bold=True, size=14)
|
||||
ws[f"A{row}"].alignment = Alignment(horizontal="center")
|
||||
row += 2
|
||||
|
||||
# Headers
|
||||
ws[f"A{row}"] = "Category"
|
||||
ws[f"B{row}"] = "Item"
|
||||
for i, dept in enumerate(departments.keys()):
|
||||
ws.cell(row, 3 + i, dept)
|
||||
ws.cell(row, 3 + i).font = header_font
|
||||
|
||||
for cell in ws[row]:
|
||||
cell.font = header_font
|
||||
row += 1
|
||||
|
||||
# Data
|
||||
for category, items in categories.items():
|
||||
# Category header (bold)
|
||||
ws[f"A{row}"] = category
|
||||
ws[f"A{row}"].font = category_font
|
||||
row += 1
|
||||
|
||||
# Items with department budgets
|
||||
for item in items:
|
||||
ws[f"A{row}"] = ""
|
||||
ws[f"B{row}"] = item
|
||||
|
||||
# Add budget amounts for each department (with color)
|
||||
for i, (dept, color) in enumerate(departments.items()):
|
||||
amount = random.randint(5000, 50000)
|
||||
cell = ws.cell(row, 3 + i, amount)
|
||||
cell.fill = PatternFill(
|
||||
start_color=color, end_color=color, fill_type="solid"
|
||||
)
|
||||
cell.number_format = "$#,##0"
|
||||
|
||||
row += 1
|
||||
|
||||
row += 1 # Blank row between categories
|
||||
|
||||
# Adjust column widths
|
||||
ws.column_dimensions["A"].width = 20
|
||||
ws.column_dimensions["B"].width = 25
|
||||
for i in range(len(departments)):
|
||||
ws.column_dimensions[chr(67 + i)].width = 15 # C, D, E, F
|
||||
|
||||
output_file = output_dir / "company_budget_2024.xlsx"
|
||||
wb.save(output_file)
|
||||
print(f" ✅ Created {output_file}")
|
||||
print(" • Bold categories, colored departments, merged title cell")
|
||||
|
||||
|
||||
def generate_workflow_4_data(output_dir: Path) -> None:
|
||||
"""Generate weekly sales report for Workflow 4."""
|
||||
print("\n📊 Generating Workflow 4: sales_weekly.xlsx")
|
||||
|
||||
products = [
|
||||
"Product A",
|
||||
"Product B",
|
||||
"Product C",
|
||||
"Product D",
|
||||
"Product E",
|
||||
"Product F",
|
||||
"Product G",
|
||||
"Product H",
|
||||
]
|
||||
|
||||
# Generate one week of data
|
||||
data = []
|
||||
start_date = datetime(2024, 11, 4) # Monday
|
||||
|
||||
for day in range(7):
|
||||
date = start_date + timedelta(days=day)
|
||||
# Each product has 3-10 transactions per day
|
||||
for product in products:
|
||||
for _ in range(random.randint(3, 10)):
|
||||
units = random.randint(1, 15)
|
||||
price = random.randint(20, 150)
|
||||
revenue = units * price
|
||||
|
||||
data.append(
|
||||
{
|
||||
"Date": date.strftime("%Y-%m-%d"),
|
||||
"Product": product,
|
||||
"Units": units,
|
||||
"Revenue": revenue,
|
||||
}
|
||||
)
|
||||
|
||||
df = pd.DataFrame(data)
|
||||
|
||||
# Write to Excel with some formatting
|
||||
output_file = output_dir / "sales_weekly.xlsx"
|
||||
with pd.ExcelWriter(output_file, engine="openpyxl") as writer:
|
||||
df.to_excel(writer, sheet_name="Weekly Sales", index=False)
|
||||
|
||||
# Format header
|
||||
worksheet = writer.sheets["Weekly Sales"]
|
||||
for cell in worksheet[1]:
|
||||
cell.font = Font(bold=True)
|
||||
|
||||
print(f" ✅ Created {output_file} ({len(df)} rows)")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Generate all sample data files."""
|
||||
print("=" * 60)
|
||||
print("Generating Sample Data for LlamaSheets + Coding Agent Workflows")
|
||||
print("=" * 60)
|
||||
|
||||
# Create output directory
|
||||
output_dir = Path("input_data")
|
||||
output_dir.mkdir(exist_ok=True)
|
||||
|
||||
# Generate data for each workflow
|
||||
generate_workflow_1_data(output_dir)
|
||||
generate_workflow_2_data(output_dir)
|
||||
generate_workflow_3_data(output_dir)
|
||||
generate_workflow_4_data(output_dir)
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("✅ All sample data generated!")
|
||||
print("=" * 60)
|
||||
print(f"\nFiles created in {output_dir.absolute()}:")
|
||||
print("\nWorkflow 1 (Understanding a New Spreadsheet):")
|
||||
print(" • financial_report_q1.xlsx")
|
||||
print("\nWorkflow 2 (Generating Analysis Scripts):")
|
||||
print(" • sales_northeast.xlsx")
|
||||
print(" • sales_southeast.xlsx")
|
||||
print(" • sales_west.xlsx")
|
||||
print("\nWorkflow 3 (Using Cell Metadata):")
|
||||
print(" • company_budget_2024.xlsx")
|
||||
print("\nWorkflow 4 (Complete Automation):")
|
||||
print(" • sales_weekly.xlsx")
|
||||
print("\nYou can now use these files with the workflows in the documentation!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,308 @@
|
||||
"""
|
||||
LlamaSheets Agent with LlamaIndex
|
||||
|
||||
This example shows how to build an agent that can work with spreadsheet data
|
||||
extracted by LlamaSheets using Python code execution.
|
||||
|
||||
The agent has minimal tools but maximum flexibility - it can execute arbitrary
|
||||
pandas code against the extracted data, similar to a coding agent.
|
||||
|
||||
NOTE: Code execution should be handled safely in a sandboxed environment for security.
|
||||
"""
|
||||
|
||||
import io
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import dotenv
|
||||
import pandas as pd
|
||||
from llama_index.core.agent import FunctionAgent, ToolCall, ToolCallResult, AgentStream
|
||||
from llama_index.llms.openai import OpenAI
|
||||
from workflows import Context
|
||||
|
||||
dotenv.load_dotenv()
|
||||
|
||||
# Global context for loaded dataframes
|
||||
_dataframe_context: Dict[str, Any] = {}
|
||||
|
||||
|
||||
# Helper function for initial agent context
|
||||
def list_extracted_data(data_dir: str = "data") -> str:
|
||||
"""
|
||||
List all regions and metadata files extracted by LlamaSheets.
|
||||
|
||||
This helps discover what data is available to work with.
|
||||
|
||||
Args:
|
||||
data_dir: Directory containing extracted parquet files (default: "data")
|
||||
|
||||
Returns:
|
||||
JSON string with information about available files
|
||||
"""
|
||||
data_path = Path(data_dir)
|
||||
|
||||
if not data_path.exists():
|
||||
return json.dumps({"error": f"Data directory '{data_dir}' not found"})
|
||||
|
||||
# Find all parquet and metadata files
|
||||
region_files = list(data_path.glob("*_region_*.parquet"))
|
||||
job_metadata_files = list(data_path.glob("*_job_metadata.json"))
|
||||
|
||||
regions = []
|
||||
for region_file in region_files:
|
||||
# Quick peek at dimensions
|
||||
df = pd.read_parquet(region_file)
|
||||
|
||||
# Find corresponding metadata file
|
||||
base_name = region_file.stem.replace("_region_", "_metadata_")
|
||||
metadata_path = region_file.parent / f"{base_name}.parquet"
|
||||
|
||||
regions.append(
|
||||
{
|
||||
"region_file": str(region_file),
|
||||
"metadata_file": str(metadata_path) if metadata_path.exists() else None,
|
||||
"shape": {"rows": len(df), "columns": len(df.columns)},
|
||||
"columns": list(df.columns),
|
||||
}
|
||||
)
|
||||
|
||||
result = {
|
||||
"data_directory": str(data_path.absolute()),
|
||||
"num_regions": len(regions),
|
||||
"regions": regions,
|
||||
"job_metadata_files": [str(f) for f in job_metadata_files],
|
||||
}
|
||||
|
||||
return json.dumps(result, indent=2)
|
||||
|
||||
|
||||
# Agent tool for code execution against dataframes
|
||||
def execute_dataframe_code(
|
||||
code: str, load_files: Optional[Dict[str, str]] = None
|
||||
) -> str:
|
||||
"""
|
||||
Execute Python pandas code against LlamaSheets extracted data.
|
||||
|
||||
This tool allows flexible data analysis by executing arbitrary pandas code.
|
||||
You can load parquet files, manipulate dataframes, and return results.
|
||||
|
||||
The code executes in a context where:
|
||||
- pandas is available as 'pd'
|
||||
- json is available for formatting output
|
||||
- Previously loaded dataframes are accessible by their variable names
|
||||
|
||||
Args:
|
||||
code: Python code to execute. Any print() statements or stdout/stderr
|
||||
will be captured and returned. Optionally set a 'result' variable
|
||||
for structured output.
|
||||
load_files: Optional dict mapping variable names to file paths to load
|
||||
Example: {"df": "data/sales_region_1.parquet",
|
||||
"meta": "data/sales_metadata_1.parquet"}
|
||||
|
||||
Returns:
|
||||
String containing:
|
||||
- Any stdout/stderr output from the code execution
|
||||
- The 'result' variable if it was set (formatted appropriately)
|
||||
- Error message if execution failed
|
||||
|
||||
Example usage:
|
||||
code = '''
|
||||
# Load and inspect data
|
||||
df = pd.read_parquet("data/sales_region_1.parquet")
|
||||
print(f"Loaded {len(df)} rows")
|
||||
|
||||
result = {
|
||||
"shape": df.shape,
|
||||
"columns": list(df.columns),
|
||||
"sample": df.head(3).to_dict(orient="records")
|
||||
}
|
||||
'''
|
||||
"""
|
||||
global _dataframe_context
|
||||
|
||||
# Capture stdout and stderr
|
||||
stdout_capture = io.StringIO()
|
||||
stderr_capture = io.StringIO()
|
||||
old_stdout = sys.stdout
|
||||
old_stderr = sys.stderr
|
||||
|
||||
try:
|
||||
# Redirect stdout/stderr
|
||||
sys.stdout = stdout_capture
|
||||
sys.stderr = stderr_capture
|
||||
|
||||
# Create execution context with pandas, json, and previously loaded dfs
|
||||
exec_context = {
|
||||
"pd": pd,
|
||||
"json": json,
|
||||
"Path": Path,
|
||||
**_dataframe_context, # Include previously loaded dataframes
|
||||
}
|
||||
|
||||
# Load any requested files into context
|
||||
if load_files:
|
||||
for var_name, file_path in load_files.items():
|
||||
if file_path.endswith(".parquet"):
|
||||
exec_context[var_name] = pd.read_parquet(file_path)
|
||||
# Also save to global context for future calls
|
||||
_dataframe_context[var_name] = exec_context[var_name]
|
||||
elif file_path.endswith(".json"):
|
||||
with open(file_path, "r") as f:
|
||||
exec_context[var_name] = json.load(f)
|
||||
_dataframe_context[var_name] = exec_context[var_name]
|
||||
|
||||
# Execute the code
|
||||
exec(code, exec_context)
|
||||
|
||||
# Restore stdout/stderr
|
||||
sys.stdout = old_stdout
|
||||
sys.stderr = old_stderr
|
||||
|
||||
# Collect output
|
||||
stdout_output = stdout_capture.getvalue()
|
||||
stderr_output = stderr_capture.getvalue()
|
||||
|
||||
output_parts = []
|
||||
|
||||
# Add stdout if any
|
||||
if stdout_output:
|
||||
output_parts.append(f"<stdout>{stdout_output}</stdout>")
|
||||
|
||||
# Add stderr if any
|
||||
if stderr_output:
|
||||
output_parts.append(f"<stderr>{stderr_output}</stderr>")
|
||||
|
||||
# Try to get a result (if code set a 'result' variable)
|
||||
if "result" in exec_context:
|
||||
result = exec_context["result"]
|
||||
result_str = None
|
||||
|
||||
if isinstance(result, pd.DataFrame):
|
||||
# Convert DataFrame to readable format
|
||||
result_str = result.to_string()
|
||||
elif isinstance(result, (dict, list)):
|
||||
result_str = json.dumps(result, indent=2, default=str)
|
||||
else:
|
||||
result_str = str(result)
|
||||
|
||||
if result_str:
|
||||
output_parts.append(f"<result_var>{result_str}</result_var>")
|
||||
|
||||
# Return combined output or success message
|
||||
if output_parts:
|
||||
return "\n\n".join(output_parts)
|
||||
else:
|
||||
return "Code executed successfully (no output or result)"
|
||||
|
||||
except Exception as e:
|
||||
# Restore stdout/stderr in case of error
|
||||
sys.stdout = old_stdout
|
||||
sys.stderr = old_stderr
|
||||
|
||||
# Get any partial output
|
||||
stdout_output = stdout_capture.getvalue()
|
||||
stderr_output = stderr_capture.getvalue()
|
||||
|
||||
error_parts = []
|
||||
if stdout_output:
|
||||
error_parts.append(f"=== STDOUT (before error) ===\n{stdout_output}")
|
||||
if stderr_output:
|
||||
error_parts.append(f"=== STDERR (before error) ===\n{stderr_output}")
|
||||
|
||||
error_parts.append(f"=== ERROR ===\n{str(e)}")
|
||||
error_parts.append(f"\n=== CODE ===\n{code}")
|
||||
|
||||
return "\n\n".join(error_parts)
|
||||
|
||||
|
||||
def create_llamasheets_agent(
|
||||
llm_model: str = "gpt-4.1", api_key: Optional[str] = None
|
||||
) -> FunctionAgent:
|
||||
# Initialize LLM
|
||||
llm = OpenAI(model=llm_model, api_key=api_key)
|
||||
|
||||
# Create tools - just 4 simple but powerful tools
|
||||
tools = [execute_dataframe_code]
|
||||
|
||||
# System prompt to guide the agent
|
||||
available_regions = list_extracted_data()
|
||||
system_prompt = f"""You are an AI assistant that helps analyze spreadsheet data extracted by LlamaSheets.
|
||||
|
||||
LlamaSheets extracts messy spreadsheets into clean parquet files with two types of outputs:
|
||||
1. Region files (*_region_*.parquet) - The actual data with columns and rows
|
||||
2. Metadata files (*_metadata_*.parquet) - Rich cell-level metadata including:
|
||||
- Formatting: font_bold, font_italic, font_size, background_color_rgb
|
||||
- Position: row_number, column_number, coordinate
|
||||
- Type detection: data_type, is_date_like, is_percentage, is_currency
|
||||
- Layout: is_in_first_row, is_merged_cell, horizontal_alignment
|
||||
|
||||
Your approach:
|
||||
1. Use list_extracted_data() to discover available files
|
||||
2. Use execute_dataframe_code() to load and analyze data with pandas
|
||||
3. Use metadata to understand structure (bold = headers, colors = groups)
|
||||
4. Use save_dataframe() to export results
|
||||
|
||||
Key tips:
|
||||
- Bold cells in metadata often indicate headers
|
||||
- Background colors often indicate groupings or departments
|
||||
- Load both region and metadata files for complete analysis
|
||||
- Write clear pandas code - you have full pandas functionality available
|
||||
- Store results in variables for reuse across multiple code executions
|
||||
|
||||
Existing Processed Regions:
|
||||
{available_regions}
|
||||
"""
|
||||
|
||||
# Configure agent
|
||||
return FunctionAgent(tools=tools, llm=llm, system_prompt=system_prompt)
|
||||
|
||||
|
||||
async def main():
|
||||
"""Example of using the LlamaSheets agent."""
|
||||
|
||||
# Create the agent
|
||||
agent = create_llamasheets_agent()
|
||||
ctx = Context(agent)
|
||||
|
||||
# Example queries the agent can handle:
|
||||
queries = [
|
||||
# Discovery
|
||||
"What spreadsheet data is available?",
|
||||
# Simple analysis
|
||||
"Load the sales data and show me the first few rows with column info",
|
||||
# Using metadata
|
||||
"Find all bold cells in the metadata - these are likely headers",
|
||||
]
|
||||
|
||||
# Example: Run a query
|
||||
for query in queries:
|
||||
print(f"\n=== Query: {query} ===")
|
||||
handler = agent.run(query, ctx=ctx)
|
||||
async for ev in handler.stream_events():
|
||||
if isinstance(ev, ToolCall):
|
||||
tool_kwargs_str = (
|
||||
str(ev.tool_kwargs)[:500] + " ..."
|
||||
if len(str(ev.tool_kwargs)) > 500
|
||||
else str(ev.tool_kwargs)
|
||||
)
|
||||
print(f"\n[Tool Call] {ev.tool_name} with args:\n{tool_kwargs_str}\n\n")
|
||||
elif isinstance(ev, ToolCallResult):
|
||||
result_str = (
|
||||
str(ev.tool_output)[:500] + " ..."
|
||||
if len(str(ev.tool_output)) > 500
|
||||
else str(ev.tool_output)
|
||||
)
|
||||
print(f"\n[Tool Result] {ev.tool_name}:\n{result_str}\n\n")
|
||||
elif isinstance(ev, AgentStream):
|
||||
print(ev.delta, end="", flush=True)
|
||||
|
||||
_ = await handler
|
||||
print("=== End Query ===\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import asyncio
|
||||
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,7 @@
|
||||
llama-cloud-services # LlamaSheets SDK
|
||||
llama-index-core
|
||||
llama-index-llms-openai
|
||||
pandas>=2.0.0
|
||||
pyarrow>=12.0.0
|
||||
openpyxl>=3.0.0 # For Excel file support
|
||||
matplotlib>=3.7.0 # For visualizations (optional)
|
||||
@@ -0,0 +1,100 @@
|
||||
"""Helper script to extract spreadsheets using LlamaSheets."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import dotenv
|
||||
from pathlib import Path
|
||||
|
||||
from llama_cloud_services.beta.sheets import LlamaSheets
|
||||
from llama_cloud_services.beta.sheets.types import (
|
||||
SpreadsheetParsingConfig,
|
||||
SpreadsheetResultType,
|
||||
)
|
||||
|
||||
dotenv.load_dotenv()
|
||||
|
||||
|
||||
async def extract_spreadsheet(
|
||||
file_path: str, output_dir: str = "data", generate_metadata: bool = True
|
||||
) -> dict:
|
||||
"""Extract a spreadsheet using LlamaSheets."""
|
||||
|
||||
client = LlamaSheets(
|
||||
base_url="https://api.cloud.llamaindex.ai",
|
||||
api_key=os.getenv("LLAMA_CLOUD_API_KEY"),
|
||||
)
|
||||
|
||||
print(f"Extracting {file_path}...")
|
||||
|
||||
# Extract regions
|
||||
config = SpreadsheetParsingConfig(
|
||||
sheet_names=None, # Extract all sheets
|
||||
generate_additional_metadata=generate_metadata,
|
||||
)
|
||||
|
||||
job_result = await client.aextract_regions(file_path, config=config)
|
||||
|
||||
print(f"Extracted {len(job_result.regions)} region(s)")
|
||||
|
||||
# Create output directory
|
||||
output_path = Path(output_dir)
|
||||
output_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Get base name for files
|
||||
base_name = Path(file_path).stem
|
||||
|
||||
# Save job metadata
|
||||
job_metadata_path = output_path / f"{base_name}_job_metadata.json"
|
||||
with open(job_metadata_path, "w") as f:
|
||||
json.dump(job_result.model_dump(mode="json"), f, indent=2)
|
||||
print(f"Saved job metadata to {job_metadata_path}")
|
||||
|
||||
# Download each region
|
||||
for idx, region in enumerate(job_result.regions, 1):
|
||||
sheet_name = region.sheet_name.replace(" ", "_")
|
||||
|
||||
# Download region data
|
||||
region_bytes = await client.adownload_region_result(
|
||||
job_id=job_result.id,
|
||||
region_id=region.region_id,
|
||||
result_type=region.region_type,
|
||||
)
|
||||
|
||||
region_path = output_path / f"{base_name}_region_{idx}_{sheet_name}.parquet"
|
||||
with open(region_path, "wb") as f:
|
||||
f.write(region_bytes)
|
||||
print(f" Table {idx}: {region_path}")
|
||||
|
||||
# Download metadata
|
||||
metadata_bytes = await client.adownload_region_result(
|
||||
job_id=job_result.id,
|
||||
region_id=region.region_id,
|
||||
result_type=SpreadsheetResultType.CELL_METADATA,
|
||||
)
|
||||
|
||||
metadata_path = output_path / f"{base_name}_metadata_{idx}_{sheet_name}.parquet"
|
||||
with open(metadata_path, "wb") as f:
|
||||
f.write(metadata_bytes)
|
||||
print(f" Metadata {idx}: {metadata_path}")
|
||||
|
||||
print(f"\nAll files saved to {output_path}/")
|
||||
|
||||
return job_result.model_dump(mode="json")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
if len(sys.argv) < 2:
|
||||
print("Usage: python scripts/extract.py <spreadsheet_file>")
|
||||
sys.exit(1)
|
||||
|
||||
file_path = sys.argv[1]
|
||||
|
||||
if not Path(file_path).exists():
|
||||
print(f"❌ File not found: {file_path}")
|
||||
sys.exit(1)
|
||||
|
||||
result = asyncio.run(extract_spreadsheet(file_path))
|
||||
print(f"\n✅ Extraction complete! Job ID: {result['id']}")
|
||||
@@ -0,0 +1,43 @@
|
||||
"""LlamaCloud Spreadsheet API SDK
|
||||
|
||||
This module provides a Python SDK for the LlamaCloud Spreadsheet API.
|
||||
"""
|
||||
|
||||
from llama_cloud_services.beta.sheets.client import (
|
||||
LlamaSheets,
|
||||
SpreadsheetAPIError,
|
||||
SpreadsheetJobError,
|
||||
SpreadsheetTimeoutError,
|
||||
)
|
||||
from llama_cloud_services.beta.sheets.types import (
|
||||
ExtractedRegionSummary,
|
||||
FileUploadResponse,
|
||||
JobStatus,
|
||||
PresignedUrlResponse,
|
||||
SpreadsheetJob,
|
||||
SpreadsheetJobResult,
|
||||
SpreadsheetParseResult,
|
||||
SpreadsheetParsingConfig,
|
||||
SpreadsheetResultType,
|
||||
WorksheetMetadata,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
# Client
|
||||
"LlamaSheets",
|
||||
# Exceptions
|
||||
"SpreadsheetAPIError",
|
||||
"SpreadsheetJobError",
|
||||
"SpreadsheetTimeoutError",
|
||||
# Types
|
||||
"ExtractedRegionSummary",
|
||||
"FileUploadResponse",
|
||||
"JobStatus",
|
||||
"PresignedUrlResponse",
|
||||
"SpreadsheetJob",
|
||||
"SpreadsheetJobResult",
|
||||
"SpreadsheetParseResult",
|
||||
"SpreadsheetParsingConfig",
|
||||
"SpreadsheetResultType",
|
||||
"WorksheetMetadata",
|
||||
]
|
||||
@@ -0,0 +1,518 @@
|
||||
import asyncio
|
||||
import io
|
||||
import os
|
||||
import time
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import httpx
|
||||
from llama_cloud.client import AsyncLlamaCloud
|
||||
from tenacity import (
|
||||
AsyncRetrying,
|
||||
retry_if_exception,
|
||||
stop_after_attempt,
|
||||
wait_exponential,
|
||||
)
|
||||
|
||||
from llama_cloud_services.beta.sheets.types import (
|
||||
FileUploadResponse,
|
||||
JobStatus,
|
||||
PresignedUrlResponse,
|
||||
SpreadsheetJob,
|
||||
SpreadsheetJobResult,
|
||||
SpreadsheetParsingConfig,
|
||||
SpreadsheetResultType,
|
||||
)
|
||||
from llama_cloud_services.constants import BASE_URL
|
||||
from llama_cloud_services.files.client import FileClient
|
||||
from llama_cloud_services.utils import (
|
||||
augment_async_errors,
|
||||
FileInput,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import pandas as pd
|
||||
|
||||
|
||||
def _should_retry_exception(exception: BaseException) -> bool:
|
||||
"""Determine if an exception should be retried."""
|
||||
if isinstance(exception, httpx.HTTPStatusError):
|
||||
return exception.response.status_code in (429, 500, 502, 503, 504)
|
||||
return False
|
||||
|
||||
|
||||
class SpreadsheetAPIError(Exception):
|
||||
"""Base exception for spreadsheet API errors"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class SpreadsheetJobError(SpreadsheetAPIError):
|
||||
"""Exception raised when a spreadsheet job fails"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class SpreadsheetTimeoutError(SpreadsheetAPIError):
|
||||
"""Exception raised when a job times out"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class LlamaSheets:
|
||||
"""Client for the LlamaCloud Spreadsheet API"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str | None = None,
|
||||
base_url: str | None = None,
|
||||
max_timeout: int = 300,
|
||||
poll_interval: int = 5,
|
||||
max_retries: int = 3,
|
||||
async_httpx_client: httpx.AsyncClient | None = None,
|
||||
) -> None:
|
||||
"""Initialize the LlamaSheets client.
|
||||
|
||||
Args:
|
||||
api_key: API key for authentication. If not provided, will use LLAMA_CLOUD_API_KEY env var
|
||||
base_url: Base URL for the API
|
||||
max_timeout: Maximum time to wait for job completion in seconds
|
||||
poll_interval: Interval between status checks in seconds
|
||||
max_retries: Maximum number of retries for failed requests
|
||||
async_httpx_client: Optional custom async httpx client
|
||||
"""
|
||||
self.api_key = api_key or os.environ.get("LLAMA_CLOUD_API_KEY")
|
||||
if not self.api_key:
|
||||
raise ValueError(
|
||||
"An API key must be provided either as an argument or via the LLAMA_CLOUD_API_KEY environment variable."
|
||||
)
|
||||
|
||||
base_url = base_url or os.environ.get("LLAMA_CLOUD_BASE_URL", BASE_URL)
|
||||
self.base_url = str(base_url).rstrip("/")
|
||||
|
||||
self.max_timeout = max_timeout
|
||||
self.poll_interval = poll_interval
|
||||
self.max_retries = max_retries
|
||||
|
||||
self._async_client: httpx.AsyncClient | None = async_httpx_client
|
||||
self._files_client = FileClient(
|
||||
AsyncLlamaCloud(
|
||||
token=self.api_key,
|
||||
base_url=self.base_url,
|
||||
httpx_client=async_httpx_client,
|
||||
)
|
||||
)
|
||||
|
||||
def _get_async_client(self) -> httpx.AsyncClient:
|
||||
"""Get or create the async httpx client"""
|
||||
if self._async_client is None:
|
||||
self._async_client = httpx.AsyncClient(
|
||||
timeout=httpx.Timeout(60.0),
|
||||
follow_redirects=True,
|
||||
)
|
||||
return self._async_client
|
||||
|
||||
def _get_headers(self) -> dict[str, str]:
|
||||
"""Get common headers for API requests"""
|
||||
return {
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
# Sync methods
|
||||
|
||||
def upload_file(
|
||||
self, file_obj: FileInput, file_name: str | None = None
|
||||
) -> FileUploadResponse:
|
||||
"""Upload a file to the Files API.
|
||||
|
||||
Args:
|
||||
file_obj: File to upload (path, bytes, or file-like object)
|
||||
file_name: Optional name for the uploaded filename
|
||||
|
||||
Returns:
|
||||
FileUploadResponse with the uploaded file ID
|
||||
"""
|
||||
with augment_async_errors():
|
||||
return asyncio.run(self.aupload_file(file_obj))
|
||||
|
||||
def create_job(
|
||||
self,
|
||||
file_id: str,
|
||||
config: dict | SpreadsheetParsingConfig | None = None,
|
||||
) -> SpreadsheetJob:
|
||||
"""Create a new spreadsheet parsing job.
|
||||
|
||||
Args:
|
||||
file_id: ID of the uploaded file
|
||||
config: Parsing configuration
|
||||
|
||||
Returns:
|
||||
SpreadsheetJob with job details
|
||||
"""
|
||||
with augment_async_errors():
|
||||
return asyncio.run(self.acreate_job(file_id, config))
|
||||
|
||||
def get_job(
|
||||
self, job_id: str, include_results_metadata: bool = True
|
||||
) -> SpreadsheetJobResult:
|
||||
"""Get the status of a spreadsheet parsing job.
|
||||
|
||||
Args:
|
||||
job_id: ID of the job
|
||||
include_results_metadata: Whether to include results metadata in the response
|
||||
|
||||
Returns:
|
||||
SpreadsheetJobResult with job status and optionally results
|
||||
"""
|
||||
with augment_async_errors():
|
||||
return asyncio.run(self.aget_job(job_id, include_results_metadata))
|
||||
|
||||
def wait_for_completion(self, job_id: str) -> SpreadsheetJobResult:
|
||||
"""Wait for a job to complete by polling.
|
||||
|
||||
Args:
|
||||
job_id: ID of the job to wait for
|
||||
|
||||
Returns:
|
||||
SpreadsheetJobResult when job is complete
|
||||
|
||||
Raises:
|
||||
SpreadsheetTimeoutError: If job doesn't complete within max_timeout
|
||||
SpreadsheetJobError: If job fails
|
||||
"""
|
||||
with augment_async_errors():
|
||||
return asyncio.run(self.await_for_completion(job_id))
|
||||
|
||||
def download_region_result(
|
||||
self,
|
||||
job_id: str,
|
||||
region_id: str,
|
||||
result_type: SpreadsheetResultType = SpreadsheetResultType.TABLE,
|
||||
) -> bytes:
|
||||
"""Download a region result (either region data or cell metadata).
|
||||
|
||||
Args:
|
||||
job_id: ID of the job
|
||||
region_id: ID of the region
|
||||
result_type: Type of result to download (region or cell_metadata)
|
||||
|
||||
Returns:
|
||||
Raw bytes of the parquet file
|
||||
"""
|
||||
with augment_async_errors():
|
||||
return asyncio.run(
|
||||
self.adownload_region_result(job_id, region_id, result_type)
|
||||
)
|
||||
|
||||
def download_region_as_dataframe(
|
||||
self,
|
||||
job_id: str,
|
||||
region_id: str,
|
||||
result_type: SpreadsheetResultType = SpreadsheetResultType.TABLE,
|
||||
) -> "pd.DataFrame":
|
||||
"""Download a region result as a pandas DataFrame.
|
||||
|
||||
Args:
|
||||
job_id: ID of the job
|
||||
region_id: ID of the region
|
||||
result_type: Type of result to download (region or cell_metadata)
|
||||
|
||||
Returns:
|
||||
pandas DataFrame
|
||||
"""
|
||||
with augment_async_errors():
|
||||
return asyncio.run(
|
||||
self.adownload_region_as_dataframe(job_id, region_id, result_type)
|
||||
)
|
||||
|
||||
def extract_regions(
|
||||
self,
|
||||
file_obj: FileInput,
|
||||
config: dict | SpreadsheetParsingConfig | None = None,
|
||||
) -> SpreadsheetJobResult:
|
||||
"""High-level method to parse a spreadsheet file.
|
||||
|
||||
This method handles the entire workflow:
|
||||
1. Upload the file
|
||||
2. Create a parsing job
|
||||
3. Wait for completion
|
||||
4. Return results
|
||||
|
||||
Args:
|
||||
file_obj: File to parse (path, bytes, or file-like object)
|
||||
config: Parsing configuration
|
||||
|
||||
Returns:
|
||||
SpreadsheetJobResult with parsing results
|
||||
"""
|
||||
with augment_async_errors():
|
||||
return asyncio.run(self.aextract_regions(file_obj, config))
|
||||
|
||||
# Async methods
|
||||
|
||||
async def aupload_file(
|
||||
self, file_obj: FileInput, file_name: str | None = None
|
||||
) -> FileUploadResponse:
|
||||
"""Upload a file to the Files API.
|
||||
|
||||
Args:
|
||||
file_obj: File to upload (path, bytes, or file-like object)
|
||||
file_name: Optional name for the uploaded filename
|
||||
|
||||
Returns:
|
||||
FileUploadResponse with the uploaded file ID
|
||||
"""
|
||||
try:
|
||||
async for attempt in AsyncRetrying(
|
||||
stop=stop_after_attempt(self.max_retries),
|
||||
wait=wait_exponential(multiplier=1, min=1, max=32),
|
||||
retry=retry_if_exception(_should_retry_exception),
|
||||
reraise=True,
|
||||
):
|
||||
with attempt:
|
||||
return await self._files_client.upload_content(
|
||||
file_obj, external_file_id=file_name
|
||||
)
|
||||
except Exception as e:
|
||||
raise SpreadsheetAPIError(f"Failed to upload file: {e}") from e
|
||||
raise RuntimeError("Tenacity did not execute")
|
||||
|
||||
async def acreate_job(
|
||||
self,
|
||||
file_id: str,
|
||||
config: dict | SpreadsheetParsingConfig | None = None,
|
||||
) -> SpreadsheetJob:
|
||||
"""Create a new spreadsheet parsing job.
|
||||
|
||||
Args:
|
||||
file_id: ID of the uploaded file
|
||||
config: Parsing configuration
|
||||
|
||||
Returns:
|
||||
SpreadsheetJob with job details
|
||||
"""
|
||||
if config is None:
|
||||
config = SpreadsheetParsingConfig()
|
||||
elif isinstance(config, dict):
|
||||
config = SpreadsheetParsingConfig.model_validate(config)
|
||||
|
||||
if not isinstance(config, SpreadsheetParsingConfig):
|
||||
raise ValueError(
|
||||
"config must be a dict or SpreadsheetParsingConfig instance"
|
||||
)
|
||||
|
||||
payload = {
|
||||
"file_id": file_id,
|
||||
"config": config.model_dump(mode="json", exclude_none=True),
|
||||
}
|
||||
|
||||
try:
|
||||
async for attempt in AsyncRetrying(
|
||||
stop=stop_after_attempt(self.max_retries),
|
||||
wait=wait_exponential(multiplier=1, min=1, max=32),
|
||||
retry=retry_if_exception(_should_retry_exception),
|
||||
reraise=True,
|
||||
):
|
||||
with attempt:
|
||||
client = self._get_async_client()
|
||||
response = await client.post(
|
||||
f"{self.base_url}/api/v1/beta/sheets/jobs",
|
||||
headers=self._get_headers(),
|
||||
json=payload,
|
||||
)
|
||||
response.raise_for_status()
|
||||
return SpreadsheetJob.model_validate(response.json())
|
||||
except Exception as e:
|
||||
raise SpreadsheetAPIError(f"Failed to create job: {e}") from e
|
||||
raise RuntimeError("Tenacity did not execute")
|
||||
|
||||
async def aget_job(
|
||||
self, job_id: str, include_results_metadata: bool = True
|
||||
) -> SpreadsheetJobResult:
|
||||
"""Get the status of a spreadsheet parsing job.
|
||||
|
||||
Args:
|
||||
job_id: ID of the job
|
||||
include_results_metadata: Whether to include results in the response
|
||||
|
||||
Returns:
|
||||
SpreadsheetJobResult with job status and optionally results
|
||||
"""
|
||||
try:
|
||||
async for attempt in AsyncRetrying(
|
||||
stop=stop_after_attempt(self.max_retries),
|
||||
wait=wait_exponential(multiplier=1, min=1, max=32),
|
||||
retry=retry_if_exception(_should_retry_exception),
|
||||
reraise=True,
|
||||
):
|
||||
with attempt:
|
||||
client = self._get_async_client()
|
||||
response = await client.get(
|
||||
f"{self.base_url}/api/v1/beta/sheets/jobs/{job_id}",
|
||||
headers=self._get_headers(),
|
||||
params={"include_results": include_results_metadata},
|
||||
)
|
||||
response.raise_for_status()
|
||||
return SpreadsheetJobResult.model_validate(response.json())
|
||||
except Exception as e:
|
||||
raise SpreadsheetAPIError(f"Failed to get job status: {e}") from e
|
||||
raise RuntimeError("Tenacity did not execute")
|
||||
|
||||
async def await_for_completion(self, job_id: str) -> SpreadsheetJobResult:
|
||||
"""Wait for a job to complete by polling.
|
||||
|
||||
Args:
|
||||
job_id: ID of the job to wait for
|
||||
|
||||
Returns:
|
||||
SpreadsheetJobResult when job is complete
|
||||
|
||||
Raises:
|
||||
SpreadsheetTimeoutError: If job doesn't complete within max_timeout
|
||||
SpreadsheetJobError: If job fails
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
while (time.time() - start_time) < self.max_timeout:
|
||||
job_result = await self.aget_job(job_id, include_results_metadata=True)
|
||||
|
||||
if job_result.status in (
|
||||
JobStatus.SUCCESS,
|
||||
JobStatus.PARTIAL_SUCCESS,
|
||||
JobStatus.ERROR,
|
||||
JobStatus.FAILURE,
|
||||
):
|
||||
if job_result.status in (JobStatus.SUCCESS, JobStatus.PARTIAL_SUCCESS):
|
||||
return job_result
|
||||
else:
|
||||
error_msg = f"Job failed with status: {job_result.status}"
|
||||
if job_result.errors:
|
||||
error_msg += f"\nErrors: {', '.join(job_result.errors)}"
|
||||
raise SpreadsheetJobError(error_msg)
|
||||
|
||||
await asyncio.sleep(self.poll_interval)
|
||||
|
||||
raise SpreadsheetTimeoutError(
|
||||
f"Job did not complete within {self.max_timeout} seconds"
|
||||
)
|
||||
|
||||
async def adownload_region_result(
|
||||
self,
|
||||
job_id: str,
|
||||
region_id: str,
|
||||
result_type: SpreadsheetResultType = SpreadsheetResultType.TABLE,
|
||||
) -> bytes:
|
||||
"""Download a region result (either region data or cell metadata).
|
||||
|
||||
Args:
|
||||
job_id: ID of the job
|
||||
region_id: ID of the region
|
||||
result_type: Type of result to download (region or cell_metadata)
|
||||
|
||||
Returns:
|
||||
Raw bytes of the parquet file
|
||||
"""
|
||||
# Get presigned URL
|
||||
presigned_response = None
|
||||
result_type_str = str(result_type)
|
||||
try:
|
||||
async for attempt in AsyncRetrying(
|
||||
stop=stop_after_attempt(self.max_retries),
|
||||
wait=wait_exponential(multiplier=1, min=1, max=32),
|
||||
retry=retry_if_exception(_should_retry_exception),
|
||||
reraise=True,
|
||||
):
|
||||
with attempt:
|
||||
client = self._get_async_client()
|
||||
response = await client.get(
|
||||
f"{self.base_url}/api/v1/beta/sheets/jobs/{job_id}/regions/{region_id}/result/{result_type_str}",
|
||||
headers=self._get_headers(),
|
||||
)
|
||||
response.raise_for_status()
|
||||
presigned_response = PresignedUrlResponse.model_validate(
|
||||
response.json()
|
||||
)
|
||||
except Exception as e:
|
||||
raise SpreadsheetAPIError(f"Failed to get presigned URL: {e}") from e
|
||||
|
||||
# Download using presigned URL
|
||||
if presigned_response is None:
|
||||
raise SpreadsheetAPIError("Failed to obtain presigned URL.")
|
||||
|
||||
try:
|
||||
async for attempt in AsyncRetrying(
|
||||
stop=stop_after_attempt(self.max_retries),
|
||||
wait=wait_exponential(multiplier=1, min=1, max=32),
|
||||
retry=retry_if_exception(_should_retry_exception),
|
||||
reraise=True,
|
||||
):
|
||||
with attempt:
|
||||
download_response = await client.get(presigned_response.url)
|
||||
download_response.raise_for_status()
|
||||
return download_response.content
|
||||
except Exception as e:
|
||||
raise SpreadsheetAPIError(f"Failed to download result: {e}") from e
|
||||
raise RuntimeError("Tenacity did not execute")
|
||||
|
||||
async def adownload_region_as_dataframe(
|
||||
self,
|
||||
job_id: str,
|
||||
region_id: str,
|
||||
result_type: SpreadsheetResultType = SpreadsheetResultType.TABLE,
|
||||
) -> "pd.DataFrame":
|
||||
"""Download a region result as a pandas DataFrame.
|
||||
|
||||
Args:
|
||||
job_id: ID of the job
|
||||
region_id: ID of the region
|
||||
result_type: Type of result to download (region or cell_metadata)
|
||||
|
||||
Returns:
|
||||
pandas DataFrame
|
||||
"""
|
||||
import pandas as pd
|
||||
|
||||
parquet_bytes = await self.adownload_region_result(
|
||||
job_id, region_id, result_type
|
||||
)
|
||||
return pd.read_parquet(io.BytesIO(parquet_bytes))
|
||||
|
||||
async def aextract_regions(
|
||||
self,
|
||||
file_obj: FileInput,
|
||||
config: dict | SpreadsheetParsingConfig | None = None,
|
||||
) -> SpreadsheetJobResult:
|
||||
"""High-level method to parse a spreadsheet file.
|
||||
|
||||
This method handles the entire workflow:
|
||||
1. Upload the file
|
||||
2. Create a parsing job
|
||||
3. Wait for completion
|
||||
4. Return results
|
||||
|
||||
Args:
|
||||
file_obj: File to parse (path, bytes, or file-like object)
|
||||
config: Parsing configuration
|
||||
|
||||
Returns:
|
||||
SpreadsheetJobResult with parsing results
|
||||
"""
|
||||
# Upload file
|
||||
file_response = await self.aupload_file(file_obj)
|
||||
|
||||
# Create job
|
||||
job = await self.acreate_job(file_response.id, config)
|
||||
|
||||
# Wait for completion
|
||||
return await self.await_for_completion(job.id)
|
||||
|
||||
async def aclose(self) -> None:
|
||||
"""Close all HTTP clients (async)"""
|
||||
if self._async_client:
|
||||
await self._async_client.aclose()
|
||||
|
||||
async def __aenter__(self) -> "LlamaSheets":
|
||||
return self
|
||||
|
||||
async def __aexit__(self, _exc_type, _exc_val, _exc_tb) -> None: # type: ignore
|
||||
await self.aclose()
|
||||
@@ -0,0 +1,156 @@
|
||||
from datetime import datetime
|
||||
from enum import Enum
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator
|
||||
|
||||
|
||||
class SpreadsheetResultType(str, Enum):
|
||||
TABLE = "table"
|
||||
EXTRA = "extra"
|
||||
CELL_METADATA = "cell_metadata"
|
||||
|
||||
def __str__(self) -> str:
|
||||
return self.value
|
||||
|
||||
|
||||
class ExtractedRegionSummary(BaseModel):
|
||||
"""A summary of a single extracted region from a spreadsheet"""
|
||||
|
||||
region_id: str = Field(
|
||||
...,
|
||||
description="Unique identifier for this region within the file",
|
||||
)
|
||||
sheet_name: str = Field(..., description="Worksheet name where region was found")
|
||||
location: str = Field(..., description="Location of the region in the spreadsheet")
|
||||
title: str | None = Field(None, description="Generated title for the region")
|
||||
description: str | None = Field(
|
||||
None, description="Generated description of the region"
|
||||
)
|
||||
region_type: SpreadsheetResultType = Field(
|
||||
..., description="Type of the extracted region"
|
||||
)
|
||||
|
||||
|
||||
class WorksheetMetadata(BaseModel):
|
||||
"""Metadata about a worksheet in a spreadsheet"""
|
||||
|
||||
sheet_name: str = Field(..., description="Name of the worksheet")
|
||||
title: str | None = Field(None, description="Generated title for the worksheet")
|
||||
description: str | None = Field(
|
||||
None, description="Generated description of the worksheet"
|
||||
)
|
||||
|
||||
|
||||
class SpreadsheetParseResult(BaseModel):
|
||||
"""Result of parsing a single spreadsheet file"""
|
||||
|
||||
success: bool = Field(..., description="Whether parsing was successful")
|
||||
file_name: str = Field(..., description="Original filename")
|
||||
|
||||
regions: list[ExtractedRegionSummary] = Field(
|
||||
default_factory=list, description="All successfully extracted regions"
|
||||
)
|
||||
worksheet_metadata: list[WorksheetMetadata] = Field(
|
||||
default_factory=list, description="Metadata for each processed worksheet"
|
||||
)
|
||||
|
||||
# Error information
|
||||
errors: list[str] = Field(
|
||||
default_factory=list, description="Any errors encountered during parsing"
|
||||
)
|
||||
|
||||
|
||||
class SpreadsheetParsingConfig(BaseModel):
|
||||
"""Configuration for spreadsheet parsing and region extraction"""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
sheet_names: list[str] | None = Field(
|
||||
default=None,
|
||||
description="The names of the sheets to extract regions from. If empty, the default sheet is extracted.",
|
||||
)
|
||||
include_hidden_cells: bool = Field(
|
||||
default=True,
|
||||
description="Whether to include hidden cells when extracting regions from the spreadsheet.",
|
||||
)
|
||||
extraction_range: str | None = Field(
|
||||
default=None,
|
||||
description="A1 notation of the range to extract a single region from. If None, the entire sheet is used.",
|
||||
)
|
||||
generate_additional_metadata: bool = Field(
|
||||
default=True,
|
||||
description="Whether to generate additional metadata (title, description) for each extracted region.",
|
||||
)
|
||||
use_experimental_processing: bool = Field(
|
||||
default=False,
|
||||
description="Enables experimental processing. Accuracy may be impacted.",
|
||||
)
|
||||
|
||||
|
||||
class SpreadsheetJob(BaseModel):
|
||||
"""A spreadsheet parsing job"""
|
||||
|
||||
id: str = Field(..., description="The ID of the job")
|
||||
user_id: str = Field(..., description="The ID of the user")
|
||||
project_id: str = Field(..., description="The ID of the project")
|
||||
file: dict = Field(..., description="The file object being parsed")
|
||||
config: SpreadsheetParsingConfig = Field(
|
||||
..., description="Configuration for the parsing job"
|
||||
)
|
||||
status: str = Field(..., description="The status of the parsing job")
|
||||
created_at: str = Field(..., description="When the job was created")
|
||||
updated_at: str = Field(..., description="When the job was last updated")
|
||||
|
||||
@field_validator("created_at", "updated_at", mode="before")
|
||||
def validate_dates(cls, v: str) -> str:
|
||||
"""Validate that the dates are in the correct format"""
|
||||
if isinstance(v, datetime):
|
||||
return v.isoformat()
|
||||
else:
|
||||
return v
|
||||
|
||||
|
||||
class SpreadsheetJobResult(SpreadsheetJob):
|
||||
"""A spreadsheet parsing job result."""
|
||||
|
||||
# Results are included when the job is complete
|
||||
success: bool | None = Field(
|
||||
None, description="Whether the job completed successfully"
|
||||
)
|
||||
regions: list[ExtractedRegionSummary] = Field(
|
||||
default_factory=list,
|
||||
description="All extracted regions (populated when job is complete)",
|
||||
)
|
||||
worksheet_metadata: list[WorksheetMetadata] = Field(
|
||||
default_factory=list,
|
||||
description="Metadata for each processed worksheet (populated when job is complete)",
|
||||
)
|
||||
errors: list[str] = Field(
|
||||
default_factory=list, description="Any errors encountered"
|
||||
)
|
||||
|
||||
|
||||
class JobStatus(str, Enum):
|
||||
"""Status of a spreadsheet parsing job"""
|
||||
|
||||
PENDING = "PENDING"
|
||||
IN_PROGRESS = "IN_PROGRESS"
|
||||
SUCCESS = "SUCCESS"
|
||||
PARTIAL_SUCCESS = "PARTIAL_SUCCESS"
|
||||
ERROR = "ERROR"
|
||||
FAILURE = "FAILURE"
|
||||
|
||||
|
||||
class PresignedUrlResponse(BaseModel):
|
||||
"""Response containing a presigned URL for downloading results"""
|
||||
|
||||
url: str = Field(..., description="The presigned URL for downloading")
|
||||
|
||||
|
||||
class FileUploadResponse(BaseModel):
|
||||
"""Response from uploading a file"""
|
||||
|
||||
id: str = Field(..., description="The ID of the uploaded file")
|
||||
name: str = Field(..., description="The name of the file")
|
||||
project_id: str = Field(..., description="The project ID")
|
||||
user_id: str = Field(..., description="The user ID")
|
||||
@@ -1,2 +1,3 @@
|
||||
BASE_URL = "https://api.cloud.llamaindex.ai"
|
||||
EU_BASE_URL = "https://api.cloud.eu.llamaindex.ai"
|
||||
POLLING_TIMEOUT_SECONDS = 300.0
|
||||
|
||||
+4
-1
@@ -14,7 +14,10 @@ dev = [
|
||||
"ipython>=8.12.3,<9",
|
||||
"jupyter>=1.1.1,<2",
|
||||
"mypy>=1.14.1,<2",
|
||||
"pydantic-settings>=2.10.1"
|
||||
"pydantic-settings>=2.10.1",
|
||||
"pandas",
|
||||
"openpyxl",
|
||||
"pyarrow"
|
||||
]
|
||||
|
||||
[project]
|
||||
|
||||
@@ -0,0 +1,163 @@
|
||||
import os
|
||||
import tempfile
|
||||
import pytest
|
||||
import pandas as pd
|
||||
|
||||
from llama_cloud_services.beta.sheets import LlamaSheets
|
||||
from llama_cloud_services.beta.sheets.types import SpreadsheetParsingConfig
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sheets_client():
|
||||
"""Create a LlamaSheets client for testing."""
|
||||
api_key = os.getenv("LLAMA_CLOUD_API_KEY")
|
||||
base_url = os.getenv("LLAMA_CLOUD_BASE_URL", "https://api.cloud.llamaindex.ai")
|
||||
|
||||
client = LlamaSheets(
|
||||
api_key=api_key,
|
||||
base_url=base_url,
|
||||
max_timeout=300,
|
||||
poll_interval=2,
|
||||
)
|
||||
return client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_excel_file():
|
||||
"""Create a temporary Excel file with sample data."""
|
||||
# Create a simple dataframe with various data types
|
||||
data = {
|
||||
"Name": ["Alice", "Bob", "Charlie", "David", "Eve"],
|
||||
"Age": [25, 30, 35, 40, 45],
|
||||
"City": ["New York", "Los Angeles", "Chicago", "Houston", "Phoenix"],
|
||||
"Salary": [50000.50, 75000.75, 100000.00, 125000.25, 150000.50],
|
||||
}
|
||||
df = pd.DataFrame(data)
|
||||
|
||||
# Create a temporary file
|
||||
with tempfile.NamedTemporaryFile(suffix=".xlsx", delete=False) as tmp:
|
||||
tmp_path = tmp.name
|
||||
df.to_excel(tmp_path, index=False, sheet_name="TestSheet")
|
||||
|
||||
yield tmp_path
|
||||
|
||||
# Cleanup
|
||||
try:
|
||||
os.unlink(tmp_path)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
os.environ.get("LLAMA_CLOUD_API_KEY", "") == "",
|
||||
reason="LLAMA_CLOUD_API_KEY not set",
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_spreadsheet_extraction_e2e(
|
||||
sheets_client: LlamaSheets, sample_excel_file: str
|
||||
):
|
||||
"""End-to-end test for spreadsheet extraction.
|
||||
|
||||
This test:
|
||||
1. Creates a temporary Excel file with sample data
|
||||
2. Uploads and extracts tables from the file
|
||||
3. Downloads the extracted table as a DataFrame
|
||||
4. Verifies the extracted data matches the original data
|
||||
"""
|
||||
# Extract tables from the spreadsheet
|
||||
result = await sheets_client.aextract_tables(sample_excel_file)
|
||||
|
||||
# Verify job completed successfully
|
||||
assert result.status in ("SUCCESS", "PARTIAL_SUCCESS")
|
||||
assert result.success is True
|
||||
|
||||
# Verify we extracted at least one table
|
||||
assert len(result.tables) > 0, "Expected at least one table to be extracted"
|
||||
|
||||
# Get the first table
|
||||
first_table = result.tables[0]
|
||||
assert first_table.sheet_name == "TestSheet"
|
||||
|
||||
# Download the table as a DataFrame
|
||||
extracted_df = await sheets_client.adownload_table_as_dataframe(
|
||||
job_id=result.id,
|
||||
table_id=first_table.table_id,
|
||||
)
|
||||
|
||||
# Load the original dataframe for comparison
|
||||
original_df = pd.read_excel(sample_excel_file)
|
||||
|
||||
# Verify the extracted DataFrame has the expected shape
|
||||
breakpoint()
|
||||
assert extracted_df.shape[0] == original_df.shape[0], (
|
||||
f"Row count mismatch: extracted {extracted_df.shape[0]}, "
|
||||
f"original {original_df.shape[0]}"
|
||||
)
|
||||
assert extracted_df.shape[1] == original_df.shape[1], (
|
||||
f"Column count mismatch: extracted {extracted_df.shape[1]}, "
|
||||
f"original {original_df.shape[1]}"
|
||||
)
|
||||
|
||||
# Verify column names match
|
||||
assert list(extracted_df.columns) == list(original_df.columns), (
|
||||
f"Column names mismatch: extracted {list(extracted_df.columns)}, "
|
||||
f"original {list(original_df.columns)}"
|
||||
)
|
||||
|
||||
# Verify data types are preserved (at least numerically)
|
||||
for col in original_df.columns:
|
||||
if original_df[col].dtype in ["int64", "float64"]:
|
||||
assert extracted_df[col].dtype in ["int64", "float64"], (
|
||||
f"Column {col} type mismatch: extracted {extracted_df[col].dtype}, "
|
||||
f"original {original_df[col].dtype}"
|
||||
)
|
||||
|
||||
# Verify the data values match (allowing for minor type conversions)
|
||||
for col in original_df.columns:
|
||||
original_values = original_df[col].tolist()
|
||||
extracted_values = extracted_df[col].tolist()
|
||||
|
||||
# Convert both to strings for comparison to handle type differences
|
||||
original_str = [str(v) for v in original_values]
|
||||
extracted_str = [str(v) for v in extracted_values]
|
||||
|
||||
assert original_str == extracted_str, (
|
||||
f"Column {col} values mismatch:\n"
|
||||
f"Original: {original_str}\n"
|
||||
f"Extracted: {extracted_str}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
os.environ.get("LLAMA_CLOUD_API_KEY", "") == "",
|
||||
reason="LLAMA_CLOUD_API_KEY not set",
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_spreadsheet_extraction_with_config(
|
||||
sheets_client: LlamaSheets, sample_excel_file: str
|
||||
):
|
||||
"""Test spreadsheet extraction with custom configuration."""
|
||||
# Create a config with specific settings
|
||||
config = SpreadsheetParsingConfig(
|
||||
sheet_names=["TestSheet"],
|
||||
include_hidden_cells=True,
|
||||
generate_additional_metadata=True,
|
||||
)
|
||||
|
||||
# Extract tables with the config
|
||||
result = await sheets_client.aextract_tables(sample_excel_file, config=config)
|
||||
|
||||
# Verify job completed successfully
|
||||
assert result.status in ("SUCCESS", "PARTIAL_SUCCESS")
|
||||
assert result.success is True
|
||||
|
||||
# Verify that additional metadata was generated
|
||||
assert len(result.worksheet_metadata) > 0
|
||||
assert result.worksheet_metadata[0].title is not None
|
||||
assert result.worksheet_metadata[0].description is not None
|
||||
|
||||
# Verify we extracted at least one table
|
||||
assert len(result.tables) > 0
|
||||
|
||||
# Verify the sheet name matches
|
||||
assert result.tables[0].sheet_name == "TestSheet"
|
||||
Reference in New Issue
Block a user