pedidosClientes/gestion_pedidos/services/ocr_service.py

"""
Servicio de OCR usando GPT-4 Vision API
"""
import base64
import json
from pathlib import Path
from typing import Dict, List, Optional
from openai import OpenAI
from django.conf import settings
from PIL import Image
import io


class OCRService:
    """Servicio para procesar imágenes y PDFs con GPT-4 Vision"""
    
    def __init__(self):
        self.client = OpenAI(api_key=settings.OPENAI_API_KEY)
    
    def _encode_image(self, image_path: Path) -> str:
        """Codifica una imagen en base64"""
        with open(image_path, "rb") as image_file:
            return base64.b64encode(image_file.read()).decode('utf-8')
    
    def _encode_pil_image(self, image: Image.Image) -> str:
        """Codifica una imagen PIL en base64"""
        buffered = io.BytesIO()
        image.save(buffered, format="PNG")
        return base64.b64encode(buffered.getvalue()).decode('utf-8')
    
    def process_pdf_pedido_cliente(self, pdf_path: Path) -> Dict:
        """
        Procesa un PDF de pedido de cliente y extrae la información
        
        Returns:
            Dict con: numero_pedido, matricula, fecha_cita, referencias (lista)
        """
        from pdf2image import convert_from_path
        
        # Convertir PDF a imágenes
        images = convert_from_path(pdf_path, dpi=200)
        
        if not images:
            raise ValueError("No se pudieron extraer imágenes del PDF")
        
        # Procesar la primera página (o todas si es necesario)
        image = images[0]
        base64_image = self._encode_pil_image(image)
        
        prompt = """
        Analiza este documento de pedido de cliente de recambios. Extrae la siguiente información en formato JSON:
        
        {
            "numero_pedido": "número único del pedido",
            "matricula": "matrícula del vehículo",
            "fecha_cita": "fecha de la cita en formato YYYY-MM-DD o YYYY-MM-DD HH:MM",
            "referencias": [
                {
                    "referencia": "código de la referencia",
                    "denominacion": "descripción/nombre de la pieza",
                    "unidades": número de unidades
                }
            ]
        }
        
        Si algún campo no está disponible, usa null. La fecha debe estar en formato ISO si es posible.
        """
        
        response = self.client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": prompt},
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/png;base64,{base64_image}"
                            }
                        }
                    ]
                }
            ],
            max_tokens=2000
        )
        
        content = response.choices[0].message.content
        
        # Extraer JSON de la respuesta
        try:
            # Buscar JSON en la respuesta
            json_start = content.find('{')
            json_end = content.rfind('}') + 1
            if json_start != -1 and json_end > json_start:
                json_str = content[json_start:json_end]
                return json.loads(json_str)
            else:
                raise ValueError("No se encontró JSON en la respuesta")
        except json.JSONDecodeError as e:
            raise ValueError(f"Error al parsear JSON: {e}. Respuesta: {content}")
    
    def process_albaran(self, image_path: Path) -> Dict:
        """
        Procesa un albarán y extrae la información
        
        Returns:
            Dict con: proveedor (nombre o número), numero_albaran, fecha_albaran, 
                     referencias (lista con precios e impuestos)
        """
        base64_image = self._encode_image(image_path)
        
        prompt = """
        Analiza este albarán de proveedor. Extrae la siguiente información en formato JSON:
        
        {
            "proveedor": {
                "nombre": "nombre del proveedor",
                "numero": "número de proveedor si está visible"
            },
            "numero_albaran": "número del albarán",
            "fecha_albaran": "fecha del albarán en formato YYYY-MM-DD",
            "referencias": [
                {
                    "referencia": "código de la referencia",
                    "denominacion": "descripción de la pieza",
                    "unidades": número de unidades,
                    "precio_unitario": precio por unidad (número decimal),
                    "impuesto_tipo": "21", "10", "7", "4", "3" o "0" según el porcentaje de IVA,
                    "impuesto_valor": valor del impuesto (número decimal)
                }
            ],
            "totales": {
                "base_imponible": total base imponible,
                "iva_21": total IVA al 21%,
                "iva_10": total IVA al 10%,
                "iva_7": total IVA al 7%,
                "iva_4": total IVA al 4%,
                "iva_3": total IVA al 3%,
                "total": total general
            }
        }
        
        IMPORTANTE: 
        - Si hay múltiples tipos de impuestos, identifica qué referencias tienen cada tipo
        - Si el impuesto no está claro por referencia, intenta calcularlo basándote en los totales
        - Si algún campo no está disponible, usa null
        """
        
        response = self.client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": prompt},
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/png;base64,{base64_image}"
                            }
                        }
                    ]
                }
            ],
            max_tokens=3000
        )
        
        content = response.choices[0].message.content
        
        try:
            json_start = content.find('{')
            json_end = content.rfind('}') + 1
            if json_start != -1 and json_end > json_start:
                json_str = content[json_start:json_end]
                return json.loads(json_str)
            else:
                raise ValueError("No se encontró JSON en la respuesta")
        except json.JSONDecodeError as e:
            raise ValueError(f"Error al parsear JSON: {e}. Respuesta: {content}")
    
    def process_image(self, image_path: Path, tipo: str = 'albaran') -> Dict:
        """
        Procesa una imagen (albarán o pedido)
        
        Args:
            image_path: Ruta a la imagen
            tipo: 'albaran' o 'pedido_cliente'
        """
        if tipo == 'albaran':
            return self.process_albaran(image_path)
        elif tipo == 'pedido_cliente':
            return self.process_pdf_pedido_cliente(image_path)
        else:
            raise ValueError(f"Tipo no soportado: {tipo}")