%pip install "google-genai>=1"

from google import genai

# 创建客户端
api_key = "YOUR_API_KEY"  # 请替换为你的API密钥
client = genai.Client(api_key=api_key)

# 定义要使用的模型
model_id = "gemini-2.0-flash"  # 也可以使用 "gemini-2.0-flash-lite-preview-02-05" 或 "gemini-2.0-pro-exp-02-05"

# 上传PDF文件
invoice_pdf = client.files.upload(file="invoice.pdf", config={'display_name': 'invoice'})

# 查看文件转换为多少个令牌
file_size = client.models.count_tokens(model=model_id, contents=invoice_pdf)
print(f'File: {invoice_pdf.display_name} equals to {file_size.total_tokens} tokens')

# 输出示例: File: invoice equals to 821 tokens

def extract_structured_data(file_path: str, model: BaseModel):
    # 上传文件到文件API
    file = client.files.upload(file=file_path, config={'display_name': file_path.split('/')[-1].split('.')[0]})
    
    # 使用Gemini API生成结构化响应
    prompt = f"Extract the structured data from the following PDF file"
    response = client.models.generate_content(model=model_id, 
                                             contents=[prompt, file], 
                                             config={'response_mime_type': 'application/json', 
                                                     'response_schema': model})
    
    # 将响应转换为Pydantic模型并返回
    return response.parsed

from pydantic import BaseModel, Field

class Item(BaseModel):
    description: str = Field(description="The description of the item")
    quantity: float = Field(description="The Qty of the item")
    gross_worth: float = Field(description="The gross worth of the item")

class Invoice(BaseModel):
    """Extract the invoice number, date and all list items with description, quantity and gross worth and the total gross worth."""
    invoice_number: str = Field(description="The invoice number e.g. 1234567890")
    date: str = Field(description="The date of the invoice e.g. 2024-01-01")
    items: list[Item] = Field(description="The list of items with description, quantity and gross worth")
    total_gross_worth: float = Field(description="The total gross worth of the invoice")

# 使用该模型提取数据
result = extract_structured_data("invoice.pdf", Invoice)

# 输出结果
print(f"Extracted Invoice: {result.invoice_number} on {result.date} with total gross worth {result.total_gross_worth}")
for item in result.items:
    print(f"Item: {item.description} with quantity {item.quantity} and gross worth {item.gross_worth}")

class Form(BaseModel):
    """Extract the form number, fiscal start date, fiscal end date, and the plan liabilities beginning of the year and end of the year."""
    form_number: str = Field(description="The Form Number")
    start_date: str = Field(description="Effective Date")
    beginning_of_year: float = Field(description="The plan liabilities beginning of the year")
    end_of_year: float = Field(description="The plan liabilities end of the year")

# 提取数据
result = extract_structured_data("handwriting_form.pdf", Form)

# 输出结果
print(f'Extracted Form Number: {result.form_number} with start date {result.start_date}.\nPlan liabilities beginning of the year {result.beginning_of_year} and end of the year {result.end_of_year}')
# 输出示例: Extracted Form Number: CA530082 with start date 02/05/2022.
# Plan liabilities beginning of the year 40000.0 and end of the year 55000.0

CHUNKING_PROMPT = """OCR the following page into Markdown. Tables should be formatted as HTML. 
Do not surround your output with triple backticks.
Chunk the document into sections of roughly 250 - 1000 words. Our goal is 
to identify parts of the page with same semantic theme. These chunks will 
be embedded and used in a RAG pipeline. 
Surround the chunks with <chunk> </chunk> html tags."""

# 使用该提示词进行处理
response = client.models.generate_content(
    model=model_id,
    contents=[CHUNKING_PROMPT, pdf_file],
)

chunked_content = response.text

class Person(BaseModel):
    first_name: str = Field(description="The first name of the person")
    last_name: str = Field(description="The last name of the person")
    age: int = Field(description="The age of the person, if not provided please return 0")
    work_topics: list[Topic] = Field(description="The fields of interest of the person, if not provided please return an empty list")

# 使用Person模型生成响应
prompt = "Philipp Schmid is a Senior AI Developer Relations Engineer at Google DeepMind working on Gemini, Gemma with the mission to help every developer to build and benefit from AI in a responsible way."
response = client.models.generate_content(
    model=model_id, 
    contents=prompt, 
    config={'response_mime_type': 'application/json', 'response_schema': Person}
)

# SDK会自动将响应转换为Pydantic模型
philipp: Person = response.parsed
print(f"First name is {philipp.first_name}")

async def batch_process_pdfs(file_paths, model, batch_size=10):
    results = []
    for i in range(0, len(file_paths), batch_size):
        batch = file_paths[i:i+batch_size]
        tasks = [extract_structured_data(path, model) for path in batch]
        batch_results = await asyncio.gather(*tasks)
        results.extend(batch_results)
        print(f"Processed batch {i//batch_size + 1}/{(len(file_paths)+batch_size-1)//batch_size}")
    return results

def extract_with_retry(file_path, model, max_retries=3):
    for attempt in range(max_retries):
        try:
            return extract_structured_data(file_path, model)
        except Exception as e:
            if attempt == max_retries - 1:
                print(f"Failed to process {file_path} after {max_retries} attempts: {e}")
                return None
            print(f"Attempt {attempt+1} failed, retrying: {e}")
            time.sleep(2 ** attempt)  # 指数退避策略

TABLE_EXTRACTION_PROMPT = """Extract all tables from the PDF as HTML tables. 
Preserve the exact structure, including merged cells, headers, and formatting. 
Each table should be semantically complete and maintain the relationships between cells.
For numeric values, maintain their exact format as shown in the document."""