53 lines
1.8 KiB
Python
53 lines
1.8 KiB
Python
from __future__ import annotations
|
|
|
|
import base64
|
|
from pathlib import Path
|
|
|
|
import requests
|
|
|
|
from backend.app.config import settings
|
|
|
|
|
|
class MistralClient:
|
|
endpoint = "https://api.mistral.ai/v1/chat/completions"
|
|
model = "pixtral-12b-2409"
|
|
|
|
def ocr_image(self, image_path: Path) -> str:
|
|
if not settings.mistral_api_key:
|
|
raise RuntimeError("MISTRAL_API_KEY no esta configurada")
|
|
data_url = self._image_data_url(image_path)
|
|
response = requests.post(
|
|
self.endpoint,
|
|
headers={
|
|
"Content-Type": "application/json",
|
|
"Authorization": f"Bearer {settings.mistral_api_key}",
|
|
},
|
|
json={
|
|
"model": self.model,
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"type": "text",
|
|
"text": "Extract all text from this image and format it as Markdown. If there are tables, represent them using Markdown table syntax. Do not include commentary.",
|
|
},
|
|
{"type": "image_url", "image_url": data_url},
|
|
],
|
|
}
|
|
],
|
|
"max_tokens": 3000,
|
|
},
|
|
timeout=180,
|
|
)
|
|
if response.status_code == 401:
|
|
raise RuntimeError("API Key de Mistral invalida")
|
|
if not response.ok:
|
|
raise RuntimeError(f"Mistral error {response.status_code}: {response.text}")
|
|
data = response.json()
|
|
return data["choices"][0]["message"]["content"]
|
|
|
|
def _image_data_url(self, image_path: Path) -> str:
|
|
encoded = base64.b64encode(image_path.read_bytes()).decode("ascii")
|
|
return f"data:image/jpeg;base64,{encoded}"
|