src/document_classifier.py

from enum import Enum

SYSTEM_PROMPT = """Tu dois classer des documents selon leur type. Les types possibles sont :
- Bon de livraison : un bon de libraison pour de la nourriture pour animaux d'élevage.
- Rapport d'analyse fourrage : un document décrivant les informations nutritionnelles et composition d'un fourrage.

Répond uniquement avec "livraison" ou "analyse". Ne réponds jamais autre chose.
"""

def classify_document (client, document_source):
  model = "mistral-small-latest" # Mistral Medium 3.1

  messages = [
    {
      "role": "system",
      "content": SYSTEM_PROMPT,
    },
    {
      "role": "user",
      "content": [
          {
              "type": "text",
              "text": "Quel est le type de ce document ?"
          },
          document_source
      ]
    }
  ]

  chat_response = client.chat.complete(
    model=model,
    messages=messages,
    temperature=0.0,
  )

  if 'livraison' in chat_response.choices[0].message.content.lower():
    return "livraison"
  elif 'analyse' in chat_response.choices[0].message.content.lower():
    return "analyse"
  else:
    raise ValueError("Document type could not be classified as 'livraison' or 'analyse'")