Initial commit for fresh VoiceInvoiceAI repo
This commit is contained in:
43
backend/config.py
Normal file
43
backend/config.py
Normal file
@@ -0,0 +1,43 @@
|
||||
# D:\GNPL-Private\VoiceInvoice\backend\config.py
|
||||
|
||||
import os
|
||||
|
||||
# --- General Application Configuration ---
|
||||
# Setting the environment to development by default if not specified
|
||||
ENVIRONMENT = os.getenv("ENVIRONMENT", "development")
|
||||
|
||||
# --- Model Configuration ---
|
||||
QWEN2_AUDIO_MODEL_NAME = "Qwen/Qwen2-Audio-7B-Instruct" # Name of the pre-trained Qwen2-Audio model
|
||||
|
||||
# --- MinIO S3 Compatible Storage Configuration ---
|
||||
MINIO_ENDPOINT = os.getenv("MINIO_ENDPOINT", "localhost:9000") # MinIO server endpoint
|
||||
MINIO_ACCESS_KEY = os.getenv("MINIO_ACCESS_KEY", "minioadmin") # MinIO access key
|
||||
MINIO_SECRET_KEY = os.getenv("MINIO_SECRET_KEY", "minioadmin") # MinIO secret key
|
||||
MINIO_SECURE = os.getenv("MINIO_SECURE", "False").lower() == "true" # Use HTTPS if true
|
||||
|
||||
# Bucket names for audio inputs and generated PDFs
|
||||
MINIO_AUDIO_BUCKET = os.getenv("MINIO_AUDIO_BUCKET", "audio-inputs")
|
||||
MINIO_PDF_BUCKET = os.getenv("MINIO_PDF_BUCKET", "generated-invoices")
|
||||
|
||||
# --- Simulated Local Databases for Autofill (for demonstration/development) ---
|
||||
# In a production environment, these would typically be replaced with actual databases (SQL, NoSQL).
|
||||
user_db = {
|
||||
"john doe": {"name": "John Doe", "address": "123 Elm St, Springfield, IL", "email": "john.doe@example.com", "default_tax_rate": 0.07},
|
||||
"acme corp": {"name": "ACME Corporation", "address": "456 Oak Ave, Metropolis, NY", "email": "info@acmecorp.com", "default_tax_rate": 0.09},
|
||||
}
|
||||
|
||||
item_db = {
|
||||
"laptop": {"description": "Laptop Computer", "unit_price": 1200.00},
|
||||
"keyboard": {"description": "Mechanical Keyboard", "unit_price": 75.00},
|
||||
"mouse": {"description": "Wireless Mouse", "unit_price": 25.00},
|
||||
"software license": {"description": "Software License (Annual)", "unit_price": 300.00},
|
||||
"consulting services": {"description": "Consulting Services (Hourly)", "unit_price": 150.00},
|
||||
"web development": {"description": "Web Development Services", "unit_price": 100.00},
|
||||
"graphic design": {"description": "Graphic Design Services", "unit_price": 80.00},
|
||||
}
|
||||
|
||||
# No longer using local file system paths for audio/PDFs directly in services,
|
||||
# these are replaced by MinIO buckets.
|
||||
# LOCAL_AUDIO_FOLDER and LOCAL_PDF_FOLDER are removed from here.
|
||||
|
||||
print(f"Configuration Loaded: Environment={ENVIRONMENT}, MinIO Endpoint={MINIO_ENDPOINT}, Audio Bucket={MINIO_AUDIO_BUCKET}, PDF Bucket={MINIO_PDF_BUCKET}")
|
||||
115
backend/core/utils.py
Normal file
115
backend/core/utils.py
Normal file
@@ -0,0 +1,115 @@
|
||||
# backend/core/utils.py
|
||||
|
||||
import torch
|
||||
import gc
|
||||
import re
|
||||
import json
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Dict, Any, Optional
|
||||
from backend.config import user_db, item_db # These are now mock DBs, kept for autofill logic
|
||||
from backend.models import InvoiceData, InvoiceItem, ValidationError # ValidationError for type hinting/catching
|
||||
|
||||
def clear_gpu_memory():
|
||||
"""Clears CUDA memory cache and runs garbage collection."""
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
gc.collect()
|
||||
|
||||
def autofill_invoice_data(invoice_data: InvoiceData) -> InvoiceData:
|
||||
"""
|
||||
Autofills invoice data based on simulated local databases (user_db, item_db).
|
||||
Fills in missing client address, item prices, and calculates totals/dates.
|
||||
|
||||
Args:
|
||||
invoice_data (InvoiceData): The invoice data Pydantic model to autofill.
|
||||
|
||||
Returns:
|
||||
InvoiceData: The updated invoice data with autofilled information.
|
||||
"""
|
||||
# Create a mutable copy to avoid modifying the original object in place
|
||||
# if this function is called on an object that might be reused or
|
||||
# that came directly from a FastAPI request body.
|
||||
invoice_data_copy = invoice_data.model_copy(deep=True)
|
||||
|
||||
if invoice_data_copy.client_name:
|
||||
normalized_client_name = invoice_data_copy.client_name.lower().strip()
|
||||
if normalized_client_name in user_db:
|
||||
client_info = user_db[normalized_client_name]
|
||||
if not invoice_data_copy.client_address:
|
||||
invoice_data_copy.client_address = client_info.get("address")
|
||||
# Only update tax_rate if it's not explicitly set in the incoming data
|
||||
if invoice_data_copy.tax_rate is None or invoice_data_copy.tax_rate == 0.08: # Default value
|
||||
invoice_data_copy.tax_rate = client_info.get("default_tax_rate", 0.08)
|
||||
|
||||
for item in invoice_data_copy.items:
|
||||
# If unit_price is missing but description exists, try to autofill
|
||||
if not item.unit_price and item.description:
|
||||
normalized_description = item.description.lower().strip()
|
||||
for db_item, db_info in item_db.items():
|
||||
if db_item in normalized_description: # Simple substring match
|
||||
item.unit_price = db_info.get("unit_price")
|
||||
# Break after finding the first match
|
||||
break
|
||||
# Ensure item total is calculated or re-calculated after autofill of unit_price
|
||||
if item.quantity is not None and item.unit_price is not None:
|
||||
item.total = round(item.quantity * item.unit_price, 2)
|
||||
elif item.total is None: # If no quantity/unit_price but total needed
|
||||
item.total = 0.0 # Default to 0 if not enough info
|
||||
|
||||
# Recalculate totals after autofilling item prices and ensure they are not None
|
||||
invoice_data_copy.subtotal = round(sum(item.total for item in invoice_data_copy.items if item.total is not None), 2)
|
||||
effective_tax_rate = invoice_data_copy.tax_rate if invoice_data_copy.tax_rate is not None else 0.0
|
||||
invoice_data_copy.tax_amount = round(invoice_data_copy.subtotal * effective_tax_rate, 2)
|
||||
invoice_data_copy.grand_total = round(invoice_data_copy.subtotal + invoice_data_copy.tax_amount, 2)
|
||||
|
||||
# Handle dates
|
||||
if not invoice_data_copy.invoice_date:
|
||||
invoice_data_copy.invoice_date = datetime.now().strftime("%Y-%m-%d")
|
||||
if not invoice_data_copy.due_date:
|
||||
try:
|
||||
invoice_date_dt = datetime.strptime(invoice_data_copy.invoice_date, "%Y-%m-%d")
|
||||
invoice_data_copy.due_date = (invoice_date_dt + timedelta(days=30)).strftime("%Y-%m-%d")
|
||||
except (ValueError, TypeError):
|
||||
# Fallback if invoice_date is invalid or not set
|
||||
invoice_data_copy.due_date = (datetime.now() + timedelta(days=30)).strftime("%Y-%m-%d")
|
||||
|
||||
|
||||
return invoice_data_copy
|
||||
|
||||
def check_model_devices(qwen2_audio_model: Optional[Any], device: str) -> str:
|
||||
"""
|
||||
Checks the device placement and memory usage of the Qwen2-Audio model.
|
||||
"""
|
||||
status_output = []
|
||||
status_output.append("--- Checking Model Device Placement ---")
|
||||
|
||||
if qwen2_audio_model:
|
||||
try:
|
||||
if hasattr(qwen2_audio_model, 'hf_device_map') and qwen2_audio_model.hf_device_map:
|
||||
status_output.append(f"Qwen2-Audio model device map: {qwen2_audio_model.hf_device_map}")
|
||||
else:
|
||||
qwen_device = next(qwen2_audio_model.parameters()).device
|
||||
status_output.append(f"Qwen2-Audio model (first parameter) is on: {qwen_device}")
|
||||
|
||||
status_output.append(f"Qwen2-Audio model data type: {next(qwen2_audio_model.parameters()).dtype}")
|
||||
|
||||
if device == "cuda":
|
||||
status_output.append(f"Current GPU Memory Allocated: {torch.cuda.memory_allocated() / (1024**3):.2f} GB")
|
||||
status_output.append(f"Current GPU Memory Cached: {torch.cuda.memory_reserved() / (1024**3):.2f} GB")
|
||||
status_output.append(f"Max GPU Memory Allocated: {torch.cuda.max_memory_allocated() / (1024**3):.2f} GB")
|
||||
status_output.append(f"Max GPU Memory Cached: {torch.cuda.max_memory_reserved() / (1024**3):.2f} GB")
|
||||
except Exception as e:
|
||||
status_output.append(f"Error checking model device: {e}")
|
||||
else:
|
||||
status_output.append("Qwen2-Audio model is not loaded.")
|
||||
|
||||
status_output.append(f"System detected device: {device}")
|
||||
if device == "cuda":
|
||||
status_output.append(f"Number of CUDA devices: {torch.cuda.device_count()}")
|
||||
if torch.cuda.device_count() > 0:
|
||||
status_output.append(f"Current CUDA device: {torch.cuda.current_device()} ({torch.cuda.get_device_name(torch.cuda.current_device())})")
|
||||
status_output.append(f"Total GPU Memory on current device: {torch.cuda.get_device_properties(torch.cuda.current_device()).total_memory / (1024**3):.2f} GB")
|
||||
else:
|
||||
status_output.append("No CUDA devices detected. Running on CPU.")
|
||||
|
||||
return "\n".join(status_output)
|
||||
238
backend/main.py
Normal file
238
backend/main.py
Normal file
@@ -0,0 +1,238 @@
|
||||
# backend/main.py
|
||||
|
||||
from fastapi import FastAPI, UploadFile, File, HTTPException, status
|
||||
from fastapi.responses import FileResponse, JSONResponse
|
||||
from typing import Optional
|
||||
import os
|
||||
import shutil
|
||||
import uuid # For generating unique IDs for uploaded files
|
||||
from io import BytesIO
|
||||
|
||||
from backend.models import InvoiceData
|
||||
from backend.config import MINIO_AUDIO_BUCKET, MINIO_PDF_BUCKET
|
||||
# Import the llm_service module itself, not its internal global variables directly
|
||||
import backend.services.llm_service as llm_service
|
||||
from backend.services.pdf_service import generate_invoice_pdf
|
||||
from backend.core.utils import check_model_devices
|
||||
from backend.services.storage_service import minio_storage_service # Import the MinIO service instance
|
||||
|
||||
app = FastAPI(
|
||||
title="Voice-Powered Invoice Generator Backend",
|
||||
description="API for processing audio requests to generate structured invoice data and PDF invoices, with MinIO storage.",
|
||||
version="0.1.0"
|
||||
)
|
||||
|
||||
# Startup event: Initialize MinIO client and ensure buckets exist, load LLM model
|
||||
@app.on_event("startup")
|
||||
async def startup_event():
|
||||
"""
|
||||
Initializes MinIO client (which in turn ensures buckets exist)
|
||||
and loads the Qwen2-Audio LLM model on application startup.
|
||||
"""
|
||||
print("Application startup: Initializing services...")
|
||||
try:
|
||||
# MinIO storage service is instantiated globally in storage_service.py
|
||||
# Its __init__ method will ensure buckets exist.
|
||||
_ = minio_storage_service
|
||||
print("MinIO storage service initialized and buckets ensured.")
|
||||
# Call the loading function within the llm_service module
|
||||
llm_service.load_qwen2_audio_model()
|
||||
except Exception as e:
|
||||
print(f"CRITICAL ERROR during startup: {e}")
|
||||
# Depending on criticality, you might want to exit or log more severely
|
||||
# For now, just print and allow app to start, but subsequent calls will fail.
|
||||
|
||||
@app.get("/")
|
||||
async def root():
|
||||
"""Root endpoint providing a welcome message."""
|
||||
return {"message": "Welcome to the Voice-Powered Invoice Generator Backend. Visit /docs for API documentation."}
|
||||
|
||||
@app.post("/load_model/", summary="Load the Qwen2-Audio LLM Model")
|
||||
async def load_model_endpoint():
|
||||
"""
|
||||
Explicitly loads the Qwen2-Audio model and processor into memory.
|
||||
This can be called if the model failed to load at startup or for re-initialization.
|
||||
"""
|
||||
# Call the loading function within the llm_service module
|
||||
status_message = llm_service.load_qwen2_audio_model()
|
||||
return {"status": status_message}
|
||||
|
||||
@app.get("/model_status/", summary="Check Qwen2-Audio Model Device and Memory Status")
|
||||
async def get_model_status():
|
||||
"""
|
||||
Provides detailed information about the Qwen2-Audio model's device placement
|
||||
and GPU memory usage if running on CUDA.
|
||||
"""
|
||||
# Access the global variables directly from the llm_service module
|
||||
status_message = check_model_devices(llm_service.qwen2_audio_model, llm_service.device)
|
||||
return {"status": status_message}
|
||||
|
||||
@app.post("/generate_invoice_from_audio/", response_model=dict, summary="Generate Invoice from Audio Input")
|
||||
async def generate_invoice_from_audio_endpoint(
|
||||
audio_file: UploadFile = File(..., description="The audio file (e.g., MP3, WAV) containing the invoice request."),
|
||||
transcript_text: Optional[str] = File(None, description="Optional pre-provided transcript of the audio.")
|
||||
):
|
||||
"""
|
||||
Receives an audio file, uploads it to MinIO, uses the Qwen2-Audio model to extract invoice data,
|
||||
autofills missing details, generates a PDF invoice, uploads the PDF to MinIO,
|
||||
and returns the structured data along with the PDF's MinIO path.
|
||||
|
||||
Returns:
|
||||
JSONResponse: Contains the extracted InvoiceData and the MinIO object name for the PDF.
|
||||
"""
|
||||
# Check the global variables directly from the llm_service module
|
||||
if llm_service.qwen2_audio_model is None or llm_service.qwen2_audio_processor is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
|
||||
detail="Qwen2-Audio model is not loaded. Please call /load_model/ first."
|
||||
)
|
||||
|
||||
# Generate a unique object name for the audio file in MinIO
|
||||
file_extension = os.path.splitext(audio_file.filename)[1]
|
||||
audio_object_name = f"audio-{uuid.uuid4()}{file_extension}"
|
||||
|
||||
try:
|
||||
# 1. Read audio content into BytesIO
|
||||
audio_content = await audio_file.read()
|
||||
audio_bytesio = BytesIO(audio_content)
|
||||
|
||||
# 2. Upload audio to MinIO
|
||||
minio_storage_service.upload_file(
|
||||
bucket_name=MINIO_AUDIO_BUCKET,
|
||||
object_name=audio_object_name,
|
||||
data=audio_bytesio,
|
||||
length=len(audio_content),
|
||||
content_type=audio_file.content_type
|
||||
)
|
||||
print(f"Audio uploaded to MinIO: {MINIO_AUDIO_BUCKET}/{audio_object_name}")
|
||||
|
||||
# 3. Process audio with LLM to extract invoice data (LLM service will download from MinIO internally)
|
||||
# Call the function from the llm_service module
|
||||
invoice_data = await llm_service.extract_and_validate_invoice_data(audio_object_name, transcript_text)
|
||||
|
||||
# 4. Generate PDF and upload to MinIO
|
||||
pdf_object_name = generate_invoice_pdf(invoice_data)
|
||||
print(f"PDF generated and uploaded to MinIO: {MINIO_PDF_BUCKET}/{pdf_object_name}")
|
||||
|
||||
return JSONResponse(content={
|
||||
"invoice_data": invoice_data.model_dump(),
|
||||
"audio_object_name": audio_object_name, # Return the MinIO path for the uploaded audio
|
||||
"pdf_object_name": pdf_object_name, # Return the MinIO path for the generated PDF
|
||||
"message": "Invoice generated successfully."
|
||||
})
|
||||
except FileNotFoundError as e:
|
||||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e))
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e))
|
||||
except RuntimeError as e:
|
||||
raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Service error: {e}")
|
||||
except Exception as e:
|
||||
# Catch-all for unexpected errors
|
||||
print(f"Unexpected error in generate_invoice_from_audio_endpoint: {e}")
|
||||
raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"An unexpected error occurred: {e}")
|
||||
|
||||
@app.post("/generate_invoice_from_data/", response_model=dict, summary="Generate Invoice from Structured Data")
|
||||
async def generate_invoice_from_data_endpoint(invoice_data: InvoiceData):
|
||||
"""
|
||||
Generates a PDF invoice directly from provided structured invoice data and uploads it to MinIO.
|
||||
This endpoint can be used if the frontend already has the parsed invoice details.
|
||||
|
||||
Args:
|
||||
invoice_data (InvoiceData): The structured invoice data.
|
||||
|
||||
Returns:
|
||||
JSONResponse: Contains the processed InvoiceData and the MinIO object name for the PDF.
|
||||
"""
|
||||
try:
|
||||
# Create a mutable copy to apply autofill without affecting the original request payload
|
||||
invoice_data_processed = invoice_data.model_copy(deep=True)
|
||||
from backend.core.utils import autofill_invoice_data
|
||||
invoice_data_processed = autofill_invoice_data(invoice_data_processed)
|
||||
|
||||
# Generate PDF and upload to MinIO
|
||||
pdf_object_name = generate_invoice_pdf(invoice_data_processed)
|
||||
print(f"PDF generated and uploaded to MinIO: {MINIO_PDF_BUCKET}/{pdf_object_name}")
|
||||
|
||||
return JSONResponse(content={
|
||||
"invoice_data": invoice_data_processed.model_dump(),
|
||||
"pdf_object_name": pdf_object_name,
|
||||
"message": "Invoice PDF generated from provided data and stored in MinIO."
|
||||
})
|
||||
except Exception as e:
|
||||
print(f"Error in generate_invoice_from_data_endpoint: {e}")
|
||||
raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error generating invoice from data: {e}")
|
||||
|
||||
@app.get("/download_invoice/{object_name:path}", summary="Download Generated Invoice PDF from MinIO")
|
||||
async def download_invoice(object_name: str):
|
||||
"""
|
||||
Downloads a previously generated PDF invoice from MinIO by its object name (path within the bucket).
|
||||
|
||||
Args:
|
||||
object_name (str): The full object name (e.g., 'invoice_client_INV-2025-XXX.pdf') of the PDF in the MinIO bucket.
|
||||
|
||||
Returns:
|
||||
FileResponse: The PDF file.
|
||||
"""
|
||||
try:
|
||||
# Check if the object exists in MinIO
|
||||
if not minio_storage_service.object_exists(MINIO_PDF_BUCKET, object_name):
|
||||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=f"Invoice PDF '{object_name}' not found in MinIO bucket '{MINIO_PDF_BUCKET}'.")
|
||||
|
||||
# Download the file from MinIO
|
||||
pdf_data_bytesio = minio_storage_service.download_file(MINIO_PDF_BUCKET, object_name)
|
||||
|
||||
# Determine the filename for the client download
|
||||
filename_for_download = os.path.basename(object_name)
|
||||
|
||||
return FileResponse(
|
||||
content=pdf_data_bytesio.getvalue(), # Get the content bytes from BytesIO
|
||||
media_type="application/pdf",
|
||||
filename=filename_for_download,
|
||||
headers={"Content-Disposition": f"attachment; filename={filename_for_download}"}
|
||||
)
|
||||
except FileNotFoundError as e:
|
||||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e))
|
||||
except Exception as e:
|
||||
print(f"Error downloading invoice from MinIO: {e}")
|
||||
raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"An error occurred while downloading the PDF: {e}")
|
||||
|
||||
@app.get("/get_audio/{object_name:path}", summary="Get Audio Input from MinIO")
|
||||
async def get_audio(object_name: str):
|
||||
"""
|
||||
Retrieves a previously uploaded audio input file from MinIO by its object name.
|
||||
|
||||
Args:
|
||||
object_name (str): The full object name (e.g., 'audio-UUID.wav') of the audio in the MinIO bucket.
|
||||
|
||||
Returns:
|
||||
FileResponse: The audio file.
|
||||
"""
|
||||
try:
|
||||
# Check if the object exists in MinIO
|
||||
if not minio_storage_service.object_exists(MINIO_AUDIO_BUCKET, object_name):
|
||||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=f"Audio file '{object_name}' not found in MinIO bucket '{MINIO_AUDIO_BUCKET}'.")
|
||||
|
||||
# Download the file from MinIO
|
||||
audio_data_bytesio = minio_storage_service.download_file(MINIO_AUDIO_BUCKET, object_name)
|
||||
|
||||
# Determine the filename for the client download
|
||||
filename_for_download = os.path.basename(object_name)
|
||||
# Attempt to infer content type; fallback to octet-stream
|
||||
content_type = "application/octet-stream"
|
||||
if filename_for_download.lower().endswith(".wav"):
|
||||
content_type = "audio/wav"
|
||||
elif filename_for_download.lower().endswith(".mp3"):
|
||||
content_type = "audio/mpeg" # or audio/mp3
|
||||
|
||||
return FileResponse(
|
||||
content=audio_data_bytesio.getvalue(),
|
||||
media_type=content_type,
|
||||
filename=filename_for_download,
|
||||
headers={"Content-Disposition": f"inline; filename={filename_for_download}"} # Use inline to play in browser
|
||||
)
|
||||
except FileNotFoundError as e:
|
||||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e))
|
||||
except Exception as e:
|
||||
print(f"Error getting audio from MinIO: {e}")
|
||||
raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"An error occurred while retrieving the audio: {e}")
|
||||
|
||||
55
backend/models.py
Normal file
55
backend/models.py
Normal file
@@ -0,0 +1,55 @@
|
||||
# backend/models.py
|
||||
|
||||
from pydantic import BaseModel, Field, ValidationError # Ensure ValidationError is imported for potential use
|
||||
from typing import List, Optional, Any
|
||||
import datetime
|
||||
|
||||
class InvoiceItem(BaseModel):
|
||||
"""
|
||||
Represents a single item in an invoice.
|
||||
"""
|
||||
description: str = Field(..., description="Description of the item or service.")
|
||||
quantity: float = Field(..., gt=0, description="Quantity of the item, must be greater than 0.")
|
||||
unit_price: float = Field(..., gt=0, description="Unit price of the item, must be greater than 0.")
|
||||
total: Optional[float] = Field(None, description="Calculated total for the item (quantity * unit_price).")
|
||||
|
||||
def model_post_init(self, __context: Any) -> None:
|
||||
"""
|
||||
Pydantic V2 post-initialization hook. Calculates total if not provided.
|
||||
"""
|
||||
if self.total is None:
|
||||
self.total = round(self.quantity * self.unit_price, 2) # Round to 2 decimal places
|
||||
|
||||
|
||||
class InvoiceData(BaseModel):
|
||||
"""
|
||||
Represents the complete structured data for an invoice.
|
||||
"""
|
||||
client_name: Optional[str] = Field(None, description="Name of the client.")
|
||||
client_address: Optional[str] = Field(None, description="Billing address of the client.")
|
||||
invoice_number: Optional[str] = Field(None, description="Unique invoice identification number.")
|
||||
invoice_date: Optional[str] = Field(None, description="Date the invoice was issued (YYYY-MM-DD format).")
|
||||
due_date: Optional[str] = Field(None, description="Date the payment is due (YYYY-MM-DD format).")
|
||||
items: List[InvoiceItem] = Field(default_factory=list, description="List of invoice items.")
|
||||
subtotal: Optional[float] = Field(None, description="Sum of all item totals.")
|
||||
tax_rate: Optional[float] = Field(0.08, ge=0, le=1, description="Tax rate applicable (e.g., 0.08 for 8%).")
|
||||
tax_amount: Optional[float] = Field(None, description="Calculated tax amount.")
|
||||
grand_total: Optional[float] = Field(None, description="Total amount including subtotal and tax.")
|
||||
notes: Optional[str] = Field(None, description="Any additional notes for the invoice.")
|
||||
|
||||
def model_post_init(self, __context: Any) -> None:
|
||||
"""
|
||||
Pydantic V2 post-initialization hook. Calculates subtotal, tax_amount, and grand_total.
|
||||
These are primarily for consistency and display, the autofill_invoice_data
|
||||
utility will ensure robust calculation.
|
||||
"""
|
||||
# Ensure item totals are calculated before summing for subtotal
|
||||
for item in self.items:
|
||||
if item.total is None: # This should ideally be handled by InvoiceItem's post_init
|
||||
item.total = round(item.quantity * item.unit_price, 2)
|
||||
|
||||
self.subtotal = round(sum(item.total for item in self.items if item.total is not None), 2) if self.items else 0.0
|
||||
# Ensure tax_rate is not None before multiplication
|
||||
effective_tax_rate = self.tax_rate if self.tax_rate is not None else 0.0
|
||||
self.tax_amount = round(self.subtotal * effective_tax_rate, 2)
|
||||
self.grand_total = round(self.subtotal + self.tax_amount, 2)
|
||||
199
backend/services/llm_service.py
Normal file
199
backend/services/llm_service.py
Normal file
@@ -0,0 +1,199 @@
|
||||
# backend/services/llm_service.py
|
||||
|
||||
import torch
|
||||
from transformers import Qwen2AudioForConditionalGeneration, AutoProcessor
|
||||
import librosa
|
||||
import re
|
||||
import json
|
||||
import os
|
||||
from typing import Dict, Any, Optional
|
||||
from backend.config import QWEN2_AUDIO_MODEL_NAME, MINIO_AUDIO_BUCKET
|
||||
from backend.models import InvoiceData, InvoiceItem
|
||||
from backend.core.utils import clear_gpu_memory, autofill_invoice_data
|
||||
from backend.services.storage_service import minio_storage_service # Import the MinIO service
|
||||
|
||||
# Global variables to hold model instances for efficiency
|
||||
qwen2_audio_processor = None
|
||||
qwen2_audio_model = None
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
def load_qwen2_audio_model() -> str:
|
||||
"""
|
||||
Loads the Qwen2-Audio model and processor.
|
||||
Uses 8-bit quantization for GPU to reduce VRAM usage.
|
||||
"""
|
||||
global qwen2_audio_processor, qwen2_audio_model, device
|
||||
|
||||
if qwen2_audio_model is None or qwen2_audio_processor is None:
|
||||
try:
|
||||
print(f"Loading Qwen2-Audio model from {QWEN2_AUDIO_MODEL_NAME} for device: {device}")
|
||||
qwen2_audio_processor = AutoProcessor.from_pretrained(QWEN2_AUDIO_MODEL_NAME, trust_remote_code=True)
|
||||
|
||||
if device == "cuda":
|
||||
qwen2_audio_model = Qwen2AudioForConditionalGeneration.from_pretrained(
|
||||
QWEN2_AUDIO_MODEL_NAME,
|
||||
load_in_8bit=True,
|
||||
device_map="auto",
|
||||
torch_dtype=torch.float16,
|
||||
trust_remote_code=True
|
||||
)
|
||||
print("Attempting to load model with 8-bit quantization.")
|
||||
else:
|
||||
qwen2_audio_model = Qwen2AudioForConditionalGeneration.from_pretrained(
|
||||
QWEN2_AUDIO_MODEL_NAME,
|
||||
device_map="auto",
|
||||
torch_dtype=torch.float32,
|
||||
trust_remote_code=True
|
||||
)
|
||||
print("Loading model for CPU in float32.")
|
||||
|
||||
qwen2_audio_model.eval()
|
||||
status_message = f"Qwen2-Audio model loaded on {device} successfully!"
|
||||
print(status_message)
|
||||
return status_message
|
||||
except Exception as e:
|
||||
error_message = f"Error loading Qwen2-Audio model: {e}"
|
||||
print(error_message)
|
||||
qwen2_audio_processor = None
|
||||
qwen2_audio_model = None
|
||||
clear_gpu_memory()
|
||||
return f"Failed to load Qwen2-Audio model: {e}"
|
||||
else:
|
||||
status_message = "Qwen2-Audio model already loaded."
|
||||
print(status_message)
|
||||
return status_message
|
||||
|
||||
def create_qwen_invoice_prompt(audio_data_path: str, prompt_text: str = "") -> Dict:
|
||||
"""
|
||||
Creates the prompt for the Qwen2-Audio LLM, instructing it to extract structured invoice data.
|
||||
The output format is explicitly requested as JSON.
|
||||
|
||||
Args:
|
||||
audio_data_path (str): The local file path to the audio data.
|
||||
This function assumes the audio is already downloaded locally.
|
||||
prompt_text (str): Optional pre-provided transcript or additional prompt text.
|
||||
|
||||
Returns:
|
||||
Dict: Inputs prepared for the Qwen2-Audio model.
|
||||
"""
|
||||
if qwen2_audio_processor is None:
|
||||
raise RuntimeError("Qwen2-Audio processor is not loaded. Please load the model first.")
|
||||
|
||||
try:
|
||||
# librosa expects a file path or file-like object. We are passing a local path.
|
||||
audio, _ = librosa.load(audio_data_path, sr=qwen2_audio_processor.feature_extractor.sampling_rate)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Error loading audio file from {audio_data_path}: {e}")
|
||||
|
||||
# Define the desired JSON schema for the LLM. Updated to match Pydantic model closely.
|
||||
json_schema = """
|
||||
{
|
||||
"client_name": "string (e.g., John Doe, ACME Corp)",
|
||||
"client_address": "string (e.g., 123 Main St, Anytown, CA)",
|
||||
"invoice_number": "string (optional, e.g., INV-2025-001)",
|
||||
"invoice_date": "YYYY-MM-DD (optional, defaults to today if not specified)",
|
||||
"due_date": "YYYY-MM-DD (optional, defaults to invoice_date + 30 days if not specified)",
|
||||
"items": [
|
||||
{
|
||||
"description": "string (e.g., Laptop, Consulting Services)",
|
||||
"quantity": "float (e.g., 1.0, 2.5)",
|
||||
"unit_price": "float (e.g., 1200.00, 75.50)"
|
||||
}
|
||||
],
|
||||
"notes": "string (optional, any additional notes)"
|
||||
}
|
||||
"""
|
||||
|
||||
# Craft the prompt for the LLM, instructing it to act as an invoice data extractor.
|
||||
messages = [
|
||||
{"role": "system", "content": "You are an AI assistant that extracts structured invoice data from spoken requests. Your output MUST be a valid JSON object matching the provided schema. Do not include any other text or explanation outside the JSON. Extract all available details, and use best guesses for missing information (like current date for invoice_date)."},
|
||||
{"role": "user", "content": f"Here is the invoice request. Please extract the details into a JSON object. Ensure item totals are calculated by quantity * unit_price. If no explicit invoice or due date, use today's date and 30 days from today respectively. If a field is not mentioned, omit it or set it to null. Here's the schema: {json_schema}\n\nInvoice request: {prompt_text}"}
|
||||
]
|
||||
|
||||
# Qwen2-Audio expects a single audio input along with a list of messages
|
||||
# The `audio` parameter is specifically for the audio input, not part of the messages content.
|
||||
inputs = qwen2_audio_processor(audio=audio, text=messages, return_tensors="pt")
|
||||
inputs = {k: v.to(device) for k, v in inputs.items()} # Move inputs to correct device
|
||||
|
||||
return inputs
|
||||
|
||||
def process_llm_output(output_ids: torch.Tensor, original_input_len: int) -> str:
|
||||
"""Decodes the LLM output and extracts the JSON string."""
|
||||
if qwen2_audio_processor is None:
|
||||
raise RuntimeError("Qwen2-Audio processor is not loaded.")
|
||||
|
||||
# Decode the generated tokens
|
||||
# Ensure to skip the original input tokens by slicing `output_ids`
|
||||
generated_text = qwen2_audio_processor.decode(output_ids[0, original_input_len:], skip_special_tokens=True)
|
||||
return generated_text
|
||||
|
||||
async def extract_and_validate_invoice_data(audio_object_name: str, transcript_text: str = "") -> InvoiceData:
|
||||
"""
|
||||
Downloads an audio file from MinIO, processes it with the LLM to extract invoice data,
|
||||
validates it with Pydantic, and autofills missing information.
|
||||
|
||||
Args:
|
||||
audio_object_name (str): The object name of the audio file in the MinIO audio bucket.
|
||||
transcript_text (str): Optional pre-provided transcript of the audio.
|
||||
|
||||
Returns:
|
||||
InvoiceData: The validated and autofilled Pydantic model representing the invoice.
|
||||
"""
|
||||
if qwen2_audio_model is None or qwen2_audio_processor is None:
|
||||
raise RuntimeError("Qwen2-Audio model and/or processor not loaded. Call load_qwen2_audio_model() first.")
|
||||
|
||||
temp_audio_file_path = f"/tmp/{audio_object_name.split('/')[-1]}" # Use /tmp for temporary storage
|
||||
|
||||
try:
|
||||
# 1. Download audio from MinIO to a temporary local file
|
||||
audio_data_bytesio = minio_storage_service.download_file(MINIO_AUDIO_BUCKET, audio_object_name)
|
||||
with open(temp_audio_file_path, "wb") as f:
|
||||
f.write(audio_data_bytesio.read())
|
||||
print(f"Downloaded audio from MinIO to temporary path: {temp_audio_file_path}")
|
||||
|
||||
# 2. Create prompt and get inputs for the LLM using the local file
|
||||
llm_inputs = create_qwen_invoice_prompt(temp_audio_file_path, transcript_text)
|
||||
original_input_len = llm_inputs['input_ids'].shape[1]
|
||||
|
||||
# 3. Generate response from LLM
|
||||
with torch.no_grad():
|
||||
generated_ids = qwen2_audio_model.generate(**llm_inputs, max_new_tokens=2048)
|
||||
|
||||
# 4. Process LLM output
|
||||
llm_raw_output = process_llm_output(generated_ids, original_input_len)
|
||||
print(f"LLM Raw Output:\n{llm_raw_output}")
|
||||
|
||||
# 5. Extract JSON part using regex
|
||||
json_match = re.search(r"```json\s*(\{.*\})\s*```", llm_raw_output, re.DOTALL)
|
||||
if not json_match:
|
||||
# Fallback: try to find the JSON directly if it's not wrapped in markdown
|
||||
json_match = re.search(r"(\{.*\})", llm_raw_output, re.DOTALL)
|
||||
|
||||
if not json_match:
|
||||
raise ValueError("Could not extract a valid JSON object from LLM output.")
|
||||
|
||||
json_str = json_match.group(1)
|
||||
|
||||
try:
|
||||
extracted_data = json.loads(json_str)
|
||||
print(f"Extracted JSON Data: {extracted_data}")
|
||||
# Validate and convert to Pydantic model
|
||||
invoice_data = InvoiceData.model_validate(extracted_data)
|
||||
print(f"Validated Invoice Data (before autofill): {invoice_data.model_dump_json(indent=2)}")
|
||||
# Autofill missing details using the comprehensive function from utils
|
||||
invoice_data = autofill_invoice_data(invoice_data)
|
||||
print(f"Final Invoice Data (after autofill): {invoice_data.model_dump_json(indent=2)}")
|
||||
return invoice_data
|
||||
except json.JSONDecodeError as e:
|
||||
raise ValueError(f"Failed to decode JSON from LLM output: {e}. Raw output: {json_str}")
|
||||
except Exception as e: # Catch Pydantic ValidationError here and re-raise as ValueError
|
||||
if "validation" in str(e).lower() and "model" in str(e).lower(): # Simple check for Pydantic validation error
|
||||
raise ValueError(f"Pydantic validation failed for extracted data: {e}. Data: {json_str}")
|
||||
else:
|
||||
raise RuntimeError(f"An unexpected error occurred during data extraction or validation: {e}")
|
||||
finally:
|
||||
# Clean up the temporary audio file
|
||||
if os.path.exists(temp_audio_file_path):
|
||||
os.remove(temp_audio_file_path)
|
||||
print(f"Cleaned up temporary audio file: {temp_audio_file_path}")
|
||||
|
||||
144
backend/services/pdf_service.py
Normal file
144
backend/services/pdf_service.py
Normal file
@@ -0,0 +1,144 @@
|
||||
# backend/services/pdf_service.py
|
||||
|
||||
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle
|
||||
from reportlab.lib.styles import getSampleStyleSheet
|
||||
from reportlab.lib.units import inch
|
||||
from reportlab.lib import colors
|
||||
from reportlab.lib.pagesizes import A4
|
||||
# from reportlab.pdfgen import canvas # Not directly used with SimpleDocTemplate for main content
|
||||
import os
|
||||
from io import BytesIO
|
||||
from datetime import datetime
|
||||
|
||||
from backend.models import InvoiceData
|
||||
from backend.config import MINIO_PDF_BUCKET
|
||||
from backend.services.storage_service import minio_storage_service # Import the MinIO service
|
||||
|
||||
def generate_invoice_pdf(invoice: InvoiceData) -> str:
|
||||
"""
|
||||
Generates a PDF invoice from the InvoiceData Pydantic model
|
||||
and uploads it to the MinIO PDF bucket.
|
||||
|
||||
Args:
|
||||
invoice (InvoiceData): The Pydantic model containing invoice data.
|
||||
|
||||
Returns:
|
||||
str: The MinIO object name (path within the bucket) of the generated PDF.
|
||||
"""
|
||||
if not invoice.invoice_number:
|
||||
# Generate a simple invoice number if not provided
|
||||
invoice_number = f"INV-{invoice.invoice_date.replace('-', '') if invoice.invoice_date else datetime.now().strftime('%Y%m%d')}-{datetime.now().strftime('%H%M%S')}"
|
||||
invoice.invoice_number = invoice_number
|
||||
else:
|
||||
invoice_number = invoice.invoice_number
|
||||
|
||||
# Sanitize client name for filename
|
||||
sanitized_client_name = "".join(c for c in invoice.client_name if c.isalnum() or c in [' ', '_']).replace(' ', '_') if invoice.client_name else "unknown_client"
|
||||
# Create a unique filename for the PDF in MinIO
|
||||
pdf_object_name = f"invoice_{sanitized_client_name}_{invoice_number}_{datetime.now().strftime('%Y%m%d%H%M%S')}.pdf"
|
||||
|
||||
# Use BytesIO to create the PDF in memory
|
||||
buffer = BytesIO()
|
||||
doc = SimpleDocTemplate(buffer, pagesize=A4,
|
||||
rightMargin=inch, leftMargin=inch,
|
||||
topMargin=inch, bottomMargin=inch)
|
||||
elements = []
|
||||
styles = getSampleStyleSheet()
|
||||
|
||||
# --- Header ---
|
||||
elements.append(Paragraph("<b>INVOICE</b>", styles['h1']))
|
||||
elements.append(Spacer(1, 0.2 * inch))
|
||||
elements.append(Paragraph(f"<b>Invoice #:</b> {invoice_number}", styles['Normal']))
|
||||
elements.append(Paragraph(f"<b>Date:</b> {invoice.invoice_date if invoice.invoice_date else 'N/A'}", styles['Normal']))
|
||||
elements.append(Paragraph(f"<b>Due Date:</b> {invoice.due_date if invoice.due_date else 'N/A'}", styles['Normal']))
|
||||
elements.append(Spacer(1, 0.4 * inch))
|
||||
|
||||
# --- Client Information ---
|
||||
elements.append(Paragraph("<b>Bill To:</b>", styles['h3']))
|
||||
if invoice.client_name:
|
||||
elements.append(Paragraph(invoice.client_name, styles['Normal']))
|
||||
if invoice.client_address:
|
||||
elements.append(Paragraph(invoice.client_address, styles['Normal']))
|
||||
elements.append(Spacer(1, 0.4 * inch))
|
||||
|
||||
# --- Items Table ---
|
||||
if invoice.items:
|
||||
data = [['Description', 'Quantity', 'Unit Price', 'Total']]
|
||||
for item in invoice.items:
|
||||
# Ensure item totals are calculated if not present (though autofill_invoice_data handles this)
|
||||
item_total = item.total if item.total is not None else (item.quantity * item.unit_price if item.quantity is not None and item.unit_price is not None else 0.0)
|
||||
data.append([
|
||||
item.description,
|
||||
f"{item.quantity:.2f}",
|
||||
f"${item.unit_price:.2f}",
|
||||
f"${item_total:.2f}"
|
||||
])
|
||||
|
||||
table_style = TableStyle([
|
||||
('BACKGROUND', (0, 0), (-1, 0), colors.HexColor('#ADD8E6')), # Light blue header
|
||||
('TEXTCOLOR', (0, 0), (-1, 0), colors.black),
|
||||
('ALIGN', (0, 0), (-1, -1), 'LEFT'),
|
||||
('FONTNAME', (0, 0), (-1, 0), 'Helvetica-Bold'),
|
||||
('BOTTOMPADDING', (0, 0), (-1, 0), 12),
|
||||
('BACKGROUND', (0, 1), (-1, -1), colors.HexColor('#F0F8FF')), # Alice blue rows
|
||||
('GRID', (0, 0), (-1, -1), 1, colors.HexColor('#C0C0C0')), # Silver grid lines
|
||||
('BOX', (0, 0), (-1, -1), 1, colors.HexColor('#808080')), # Grey box border
|
||||
('ALIGN', (1, 0), (-1, -1), 'RIGHT'), # Align Qty, Unit Price, Total to right
|
||||
])
|
||||
|
||||
item_table = Table(data, colWidths=[3*inch, 0.8*inch, 1*inch, 1*inch])
|
||||
item_table.setStyle(table_style)
|
||||
elements.append(item_table)
|
||||
elements.append(Spacer(1, 0.2 * inch))
|
||||
|
||||
# --- Totals ---
|
||||
# Ensure all totals are float and not None before formatting
|
||||
subtotal = invoice.subtotal if invoice.subtotal is not None else 0.0
|
||||
tax_amount = invoice.tax_amount if invoice.tax_amount is not None else 0.0
|
||||
grand_total = invoice.grand_total if invoice.grand_total is not None else 0.0
|
||||
tax_rate_display = invoice.tax_rate * 100 if invoice.tax_rate is not None else 0.0
|
||||
|
||||
totals_data = [
|
||||
['Subtotal:', f"${subtotal:.2f}"],
|
||||
[f"Tax ({tax_rate_display:.0f}%):", f"${tax_amount:.2f}"],
|
||||
['Grand Total:', f"${grand_total:.2f}"]
|
||||
]
|
||||
totals_table = Table(totals_data, colWidths=[4.8*inch, 1.2*inch])
|
||||
totals_table.setStyle(TableStyle([
|
||||
('ALIGN', (0, 0), (-1, -1), 'RIGHT'),
|
||||
('FONTNAME', (0, -1), (-1, -1), 'Helvetica-Bold'),
|
||||
('TOPPADDING', (0, 0), (-1, -1), 4),
|
||||
('BOTTOMPADDING', (0, 0), (-1, -1), 4),
|
||||
]))
|
||||
elements.append(totals_table)
|
||||
elements.append(Spacer(1, 0.3 * inch))
|
||||
|
||||
# --- Notes ---
|
||||
if invoice.notes:
|
||||
elements.append(Paragraph("<b>Notes:</b>", styles['h3']))
|
||||
elements.append(Paragraph(invoice.notes, styles['Normal']))
|
||||
elements.append(Spacer(1, 0.2 * inch))
|
||||
|
||||
# --- Footer (Optional) ---
|
||||
elements.append(Spacer(1, 0.5 * inch))
|
||||
elements.append(Paragraph("<i>Thank you for your business!</i>", styles['Italic']))
|
||||
|
||||
# Build PDF in memory
|
||||
try:
|
||||
doc.build(elements)
|
||||
buffer.seek(0) # Reset buffer position to the beginning
|
||||
|
||||
# Upload the PDF from memory to MinIO
|
||||
minio_storage_service.upload_file(
|
||||
bucket_name=MINIO_PDF_BUCKET,
|
||||
object_name=pdf_object_name,
|
||||
data=buffer,
|
||||
length=buffer.tell(), # Get the current position (which is the file size)
|
||||
content_type="application/pdf"
|
||||
)
|
||||
print(f"Invoice PDF generated and uploaded to MinIO: {MINIO_PDF_BUCKET}/{pdf_object_name}")
|
||||
return pdf_object_name
|
||||
except Exception as e:
|
||||
print(f"Error generating or uploading PDF: {e}")
|
||||
raise RuntimeError(f"Failed to generate or upload PDF: {e}")
|
||||
|
||||
145
backend/services/storage_service.py
Normal file
145
backend/services/storage_service.py
Normal file
@@ -0,0 +1,145 @@
|
||||
# backend/services/storage_service.py
|
||||
|
||||
import os
|
||||
from io import BytesIO
|
||||
from minio import Minio
|
||||
from minio.error import S3Error
|
||||
from backend.config import MINIO_ENDPOINT, MINIO_ACCESS_KEY, MINIO_SECRET_KEY, MINIO_SECURE, MINIO_AUDIO_BUCKET, MINIO_PDF_BUCKET
|
||||
|
||||
class MinIOStorageService:
|
||||
"""
|
||||
Service class for interacting with MinIO (S3-compatible) storage.
|
||||
Handles uploading, downloading, and checking existence of objects in specified buckets.
|
||||
"""
|
||||
def __init__(self):
|
||||
"""
|
||||
Initializes the MinIO client.
|
||||
Ensures the necessary buckets exist on startup.
|
||||
"""
|
||||
try:
|
||||
self.client = Minio(
|
||||
endpoint=MINIO_ENDPOINT,
|
||||
access_key=MINIO_ACCESS_KEY,
|
||||
secret_key=MINIO_SECRET_KEY,
|
||||
secure=MINIO_SECURE,
|
||||
# Set a higher timeout for potentially large file uploads/downloads
|
||||
# or slow network conditions. Default is usually sufficient but good for large files.
|
||||
# Timeout in seconds. Can be None for no timeout.
|
||||
http_client=None # Use default http client, or provide custom
|
||||
)
|
||||
print(f"MinIO client initialized for endpoint: {MINIO_ENDPOINT}, Secure: {MINIO_SECURE}")
|
||||
self._ensure_buckets_exist()
|
||||
except Exception as e:
|
||||
print(f"ERROR: Failed to initialize MinIO client: {e}")
|
||||
raise
|
||||
|
||||
def _ensure_buckets_exist(self):
|
||||
"""
|
||||
Ensures that the audio and PDF buckets exist. Creates them if they don't.
|
||||
"""
|
||||
buckets_to_create = [MINIO_AUDIO_BUCKET, MINIO_PDF_BUCKET]
|
||||
for bucket_name in buckets_to_create:
|
||||
try:
|
||||
if not self.client.bucket_exists(bucket_name):
|
||||
self.client.make_bucket(bucket_name)
|
||||
print(f"MinIO bucket '{bucket_name}' created successfully.")
|
||||
else:
|
||||
print(f"MinIO bucket '{bucket_name}' already exists.")
|
||||
except S3Error as e:
|
||||
print(f"ERROR: S3 Error ensuring bucket '{bucket_name}': {e}")
|
||||
raise
|
||||
except Exception as e:
|
||||
print(f"ERROR: Unexpected error ensuring bucket '{bucket_name}': {e}")
|
||||
raise
|
||||
|
||||
def upload_file(self, bucket_name: str, object_name: str, data: BytesIO, length: int, content_type: str = "application/octet-stream") -> str:
|
||||
"""
|
||||
Uploads a file (BytesIO object) to a specified MinIO bucket.
|
||||
|
||||
Args:
|
||||
bucket_name (str): The name of the bucket.
|
||||
object_name (str): The desired name of the object in the bucket.
|
||||
data (BytesIO): The file data as a BytesIO object.
|
||||
length (int): The length of the data in bytes.
|
||||
content_type (str): The MIME type of the file.
|
||||
|
||||
Returns:
|
||||
str: The full path of the uploaded object (bucket_name/object_name).
|
||||
"""
|
||||
try:
|
||||
self.client.put_object(
|
||||
bucket_name=bucket_name,
|
||||
object_name=object_name,
|
||||
data=data,
|
||||
length=length,
|
||||
content_type=content_type
|
||||
)
|
||||
print(f"Successfully uploaded {object_name} to bucket {bucket_name}")
|
||||
# Construct a downloadable URL (for MinIO, this usually implies access via its API/proxy)
|
||||
# For direct public access, you'd need presigned URLs or public buckets configured externally.
|
||||
# Here, we return a logical path that your FastAPI can use to retrieve.
|
||||
return f"{bucket_name}/{object_name}"
|
||||
except S3Error as e:
|
||||
print(f"ERROR: S3 Error uploading {object_name} to {bucket_name}: {e}")
|
||||
raise
|
||||
except Exception as e:
|
||||
print(f"ERROR: Unexpected error uploading {object_name}: {e}")
|
||||
raise
|
||||
|
||||
def download_file(self, bucket_name: str, object_name: str) -> BytesIO:
|
||||
"""
|
||||
Downloads a file from a specified MinIO bucket.
|
||||
|
||||
Args:
|
||||
bucket_name (str): The name of the bucket.
|
||||
object_name (str): The name of the object to download.
|
||||
|
||||
Returns:
|
||||
BytesIO: The file data as a BytesIO object.
|
||||
"""
|
||||
try:
|
||||
response = self.client.get_object(bucket_name, object_name)
|
||||
file_data = BytesIO(response.read())
|
||||
file_data.seek(0) # Reset stream position to the beginning
|
||||
response.close()
|
||||
response.release_conn()
|
||||
print(f"Successfully downloaded {object_name} from bucket {bucket_name}")
|
||||
return file_data
|
||||
except S3Error as e:
|
||||
if e.code == "NoSuchKey":
|
||||
print(f"ERROR: Object '{object_name}' not found in bucket '{bucket_name}'.")
|
||||
raise FileNotFoundError(f"Object '{object_name}' not found in bucket '{bucket_name}'.")
|
||||
else:
|
||||
print(f"ERROR: S3 Error downloading {object_name} from {bucket_name}: {e}")
|
||||
raise
|
||||
except Exception as e:
|
||||
print(f"ERROR: Unexpected error downloading {object_name}: {e}")
|
||||
raise
|
||||
|
||||
def object_exists(self, bucket_name: str, object_name: str) -> bool:
|
||||
"""
|
||||
Checks if an object exists in a specified MinIO bucket.
|
||||
|
||||
Args:
|
||||
bucket_name (str): The name of the bucket.
|
||||
object_name (str): The name of the object to check.
|
||||
|
||||
Returns:
|
||||
bool: True if the object exists, False otherwise.
|
||||
"""
|
||||
try:
|
||||
self.client.stat_object(bucket_name, object_name)
|
||||
return True
|
||||
except S3Error as e:
|
||||
if e.code == "NoSuchKey":
|
||||
return False
|
||||
else:
|
||||
print(f"ERROR: S3 Error checking existence of {object_name} in {bucket_name}: {e}")
|
||||
raise
|
||||
except Exception as e:
|
||||
print(f"ERROR: Unexpected error checking existence of {object_name}: {e}")
|
||||
raise
|
||||
|
||||
# Instantiate the service globally or pass it around via dependency injection in FastAPI
|
||||
# For simplicity, we'll instantiate it here. FastAPI's Depends can be used for DI.
|
||||
minio_storage_service = MinIOStorageService()
|
||||
Reference in New Issue
Block a user