Spaces:
Sleeping
Sleeping
File size: 2,342 Bytes
6cc3191 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | import json
import re
from pathlib import Path
from typing import Optional
import pymupdf4llm
from groq import Groq
from loguru import logger
from pydantic import BaseModel, field_validator
PROMPT_TEMPLATE = (Path(__file__).parent / "prompts" / "invoice.prompt.txt").read_text(encoding="utf-8")
class Invoice(BaseModel):
proveedor: str = "Desconocido"
nif_proveedor: Optional[str] = None
cliente: Optional[str] = None
numero_factura: Optional[str] = None
fecha: Optional[str] = None
fecha_vencimiento: Optional[str] = None
concepto: str = "Sin descripción"
categoria: str = "Otros"
subtotal: Optional[float] = None
iva_porcentaje: Optional[float] = None
iva_importe: Optional[float] = None
total: Optional[float] = None
moneda: str = "EUR"
@field_validator("categoria")
@classmethod
def validate_categoria(cls, v: str) -> str:
valid = {"Servicios", "Suministros", "Transporte", "Software", "Material", "Alquiler", "Otros"}
return v if v in valid else "Otros"
def extract_text_from_pdf(pdf_path: str) -> str:
try:
pages = pymupdf4llm.to_markdown(pdf_path, page_chunks=True, show_progress=False)
return "\n\n".join(p["text"] for p in pages)
except Exception as e:
logger.error(f"Error leyendo {pdf_path}: {e}")
return ""
def extract_invoice_data(pdf_path: str, client: Groq, model: str = "llama-3.1-8b-instant") -> Invoice:
text = extract_text_from_pdf(pdf_path)
if not text.strip():
logger.warning(f"Sin texto en {pdf_path}")
return Invoice(proveedor=Path(pdf_path).stem)
prompt = PROMPT_TEMPLATE.format(text=text[:4000])
try:
response = client.chat.completions.create(
model=model,
messages=[{"role": "user", "content": prompt}],
temperature=0,
max_tokens=512,
)
raw = response.choices[0].message.content.strip()
# Extract JSON block if wrapped in markdown
match = re.search(r"\{.*\}", raw, re.DOTALL)
if match:
raw = match.group(0)
data = json.loads(raw)
return Invoice(**data)
except (json.JSONDecodeError, Exception) as e:
logger.error(f"Error extrayendo datos de {pdf_path}: {e}")
return Invoice(proveedor=Path(pdf_path).stem)
|