remittance-poc-with-verifier / anthropic_api_invoice_extractor.py
eelang's picture
Upload 8 files
7850a69 verified
import anthropic
import os
from remittance_pdf_processing_types import InvoiceNumbers, PaymentAmount
from remittance_pdf_processing_utils import remittance_logger, remove_duplicate_lists
from anthropic.types import ContentBlock, ImageBlockParam
def extract_invoice_numbers_with_anthropic_ai(base64_images: list[str], multi_hop: bool = False) -> list[InvoiceNumbers]:
"""
Extracts invoice numbers from one or more base64-encoded images using Anthropic's Claude 3.5 Sonnet model.
Args:
base64_images (list[str]): A list of base64-encoded image strings.
multi_hop (bool): Whether to use multi-hop processing.
Returns:
list[InvoiceNumbers]: A list containing lists of extracted invoice numbers.
"""
if multi_hop:
return extract_invoice_numbers_with_anthropic_ai_multi_hop(base64_images)
else:
return extract_invoice_numbers_with_anthropic_ai_single_hop(base64_images)
def extract_invoice_numbers_with_anthropic_ai_single_hop(base64_images: list[str]) -> list[InvoiceNumbers]:
client = anthropic.Anthropic(
api_key=os.environ.get("ANTHROPIC_API_KEY"),
)
content: list[ContentBlock] = [
{
"type": "image",
"source": {
"type": "base64",
"media_type": "image/png",
"data": image
}
} for image in base64_images
]
message = client.messages.create(
model="claude-3-5-sonnet-20240620",
max_tokens=1024,
temperature=0,
system="Given the remittance letter images, extract all invoice numbers. Respond with a comma-separated list only.",
messages=[
{
"role": "user",
"content": content
}
]
)
remittance_logger.debug(f'Anthropic (raw) response: {message.content}')
invoice_numbers = parse_anthropic_response(message.content[0].text)
return [invoice_numbers]
def extract_invoice_numbers_with_anthropic_ai_multi_hop(base64_images: list[str]) -> list[InvoiceNumbers]:
# First hop: Extract column headers
column_headers = extract_column_headers_from_images(base64_images)
remittance_logger.debug(f"Extracted column headers: {column_headers}")
# Second hop: Extract invoice numbers for each column (up to 3 columns)
all_invoice_numbers = []
for column_name in column_headers[:3]:
invoice_numbers = extract_invoice_numbers_for_column_from_images(base64_images, column_name)
remittance_logger.debug(f"Extracted invoice numbers for column '{column_name}': {invoice_numbers}")
if invoice_numbers: # Only add non-empty lists
all_invoice_numbers.append(invoice_numbers)
# Remove duplicate lists using the utility function
unique_invoice_numbers = remove_duplicate_lists(all_invoice_numbers)
return unique_invoice_numbers
def extract_column_headers_from_images(base64_images: list[str]) -> list[str]:
client = anthropic.Anthropic(
api_key=os.environ.get("ANTHROPIC_API_KEY"),
)
content: list[ContentBlock] = [
{
"type": "image",
"source": {
"type": "base64",
"media_type": "image/png",
"data": image
}
} for image in base64_images
]
message = client.messages.create(
model="claude-3-5-sonnet-20240620",
max_tokens=1024,
temperature=0,
system="Given the remittance letter images, extract all column header names that could contain invoice numbers. Respond with a comma-separated list only.",
messages=[
{
"role": "user",
"content": content
}
]
)
remittance_logger.debug(f'Anthropic (raw) response for column headers: {message.content}')
return parse_anthropic_response(message.content[0].text)
def extract_invoice_numbers_for_column_from_images(base64_images: list[str], column_name: str) -> InvoiceNumbers:
client = anthropic.Anthropic(
api_key=os.environ.get("ANTHROPIC_API_KEY"),
)
content: list[ContentBlock] = [
{
"type": "image",
"source": {
"type": "base64",
"media_type": "image/png",
"data": image
}
} for image in base64_images
]
message = client.messages.create(
model="claude-3-5-sonnet-20240620",
max_tokens=1024,
temperature=0,
system=f"Given the remittance letter images, extract all invoice numbers from the column '{column_name}'. Respond with a comma-separated list only.",
messages=[
{
"role": "user",
"content": content
}
]
)
remittance_logger.debug(f'Anthropic (raw) response for invoice numbers in column {column_name}: {message.content}')
return parse_anthropic_response(message.content[0].text)
def parse_anthropic_response(response: str) -> list[str]:
"""
Parses the response from Claude 3.5 Sonnet model and extracts a list of items.
Args:
response (str): The response string from Claude 3.5 Sonnet model.
Returns:
list[str]: A list of extracted items (invoice numbers or column headers).
"""
return [item.strip() for item in response.split(',') if item.strip()]
def extract_invoice_numbers_from_single_base64_image(base64_image: str, multi_hop: bool = False) -> list[InvoiceNumbers]:
"""
Extracts invoice numbers from a single base64-encoded image using Anthropic's Claude 3.5 Sonnet model.
Args:
base64_image (str): The base64-encoded image string.
multi_hop (bool): Whether to use multi-hop processing.
Returns:
list[InvoiceNumbers]: A list containing lists of extracted invoice numbers.
"""
return extract_invoice_numbers_with_anthropic_ai([base64_image], multi_hop)
def extract_invoice_numbers_from_multi_page_images(base64_images: list[str], multi_hop: bool = False) -> list[InvoiceNumbers]:
"""
Extracts invoice numbers from multiple base64-encoded images using Anthropic's Claude 3.5 Sonnet model.
Args:
base64_images (list[str]): A list of base64-encoded image strings.
multi_hop (bool): Whether to use multi-hop processing.
Returns:
list[InvoiceNumbers]: A list containing lists of extracted invoice numbers.
"""
return extract_invoice_numbers_with_anthropic_ai(base64_images, multi_hop)
def extract_payment_amounts_with_anthropic_ai(base64_images: list[str]) -> list[PaymentAmount]:
"""
Extracts payment amounts from one or more base64-encoded images using Anthropic's Claude 3.5 Sonnet model.
Args:
base64_images (list[str]): A list of base64-encoded image strings.
Returns:
list[PaymentAmount]: A list containing extracted payment amounts.
"""
client = anthropic.Anthropic(
api_key=os.environ.get("ANTHROPIC_API_KEY"),
)
# Prepare the message content
content = []
for image in base64_images:
content.append({
"type": "image",
"source": {
"type": "base64",
"media_type": "image/png",
"data": image
}
})
# Call the Anthropic API
message = client.messages.create(
model="claude-3-5-sonnet-20240620",
max_tokens=1024,
temperature=0,
system="You are a precise payment amount extractor. Given remittance letter images, extract only the total payment amount. Respond with the numerical amount only, including any decimal places and currency symbols if present. Do not include any additional text or explanations.",
messages=[
{
"role": "user",
"content": content
}
]
)
remittance_logger.debug(f'Anthropic (raw) response for payment amount: {message.content}')
# assert(isinstance(message.content, anthropic.TextBlock))
# Parse the response
payment_amount = parse_anthropic_payment_response(message.content[0].text)
return payment_amount
def parse_anthropic_payment_response(response: str) -> list[PaymentAmount]:
"""
Parses the response from Claude 3.5 Sonnet model and extracts the payment amount.
Args:
response (str): The response string from Claude 3.5 Sonnet model.
Returns:
list[PaymentAmount]: A list containing the extracted payment amount.
"""
# Strip whitespace and return as a single-item list
return [response.strip()]