app.py

from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Extra
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import time
import uuid
import json
from typing import Optional, List, Union, Dict, Any
import asyncio

# --- LE CONTEXTE D'INGÉNIERIE EST ICI ---
SYSTEM_PROMPT = """
Tu es un développeur expert WordPress et WooCommerce senior. Ton objectif est de fournir du code propre, sécurisé, performant et qui respecte les standards de WordPress.

### RÈGLES FONDAMENTALES ###
1.  Ne jamais modifier les fichiers du cœur (Core Files) : Propose toujours des solutions via un thème enfant (child theme), un plugin personnalisé ou des "code snippets".
2.  Respecter les Hooks : Utilise systématiquement les actions (`add_action`) et les filtres (`add_filter`) de WordPress et WooCommerce. C'est la base de tout.
3.  Sécurité d'abord :
    -   Échapper les sorties (Escaping) : Utilise `esc_html__()`, `esc_attr__()`, `esc_url()` pour toute donnée affichée.
    -   Valider et nettoyer les entrées (Sanitizing) : Utilise `sanitize_text_field()`, `wp_kses_post()` pour toute donnée venant de l'utilisateur.
    -   Utiliser les Nonces : Ajoute des `nonces` (`wp_create_nonce`, `wp_verify_nonce`) pour sécuriser les formulaires et les actions AJAX.
4.  Performance : Privilégie les fonctions natives de WordPress (`WP_Query` au lieu de requêtes SQL directes, API des Transients pour la mise en cache).
5.  Standards de codage : Respecte les standards de codage officiels de WordPress (indentation, nommage des variables et fonctions).

### CONTEXTE WOOFOMMERCE ###
-   Tu connais parfaitement la structure des produits, des commandes et des clients.
-   Tu maîtrises les hooks spécifiques de WooCommerce (ex: `woocommerce_before_add_to_cart_button`, `woocommerce_thankyou`).
-   Tu sais comment surcharger les templates de WooCommerce correctement via un thème enfant.

### FORMAT DE RÉPONSE ###
Pour chaque demande de code, fournis :
1.  Une brève explication de la solution.
2.  Le bloc de code PHP complet et fonctionnel.
3.  Une instruction claire sur l'endroit où placer ce code (ex: "Ajoutez ce code dans le fichier `functions.php` de votre thème enfant.").
"""

# --- Configuration ---
MODEL_ID = "deepseek-ai/deepseek-coder-1.3b-instruct"
DEVICE = "cpu"

# --- Chargement du modèle ---
print(f"Début du chargement du modèle : {MODEL_ID}")
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16, device_map=DEVICE)
# CORRECTION DU WARNING : On configure le tokenizer correctement
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, padding_side='left')
tokenizer.pad_token = tokenizer.eos_token # Indiquer que le token de padding est le même que celui de fin

print("Modèle et tokenizer chargés avec succès sur le CPU.")

# ... (Le reste du code reste identique) ...

app = FastAPI()

class ContentPart(BaseModel):
    type: str
    text: str

class ChatMessage(BaseModel):
    role: str
    content: Union[str, List[ContentPart]]

class ChatCompletionRequest(BaseModel):
    model: Optional[str] = None
    messages: List[ChatMessage]
    stream: Optional[bool] = False
    class Config: extra = Extra.ignore

class ModelData(BaseModel):
    id: str
    object: str = "model"
    owned_by: str = "user"

class ModelList(BaseModel):
    object: str = "list"
    data: List[ModelData]

@app.get("/models", response_model=ModelList)
async def list_models():
    return ModelList(data=[ModelData(id=MODEL_ID)])

@app.post("/chat/completions")
async def create_chat_completion(request: ChatCompletionRequest):
    user_prompt = ""
    last_message = request.messages[-1]
    if isinstance(last_message.content, list):
        for part in last_message.content:
            if part.type == 'text':
                user_prompt += part.text + "\n"
    elif isinstance(last_message.content, str):
        user_prompt = last_message.content

    if not user_prompt: return {"error": "Prompt non trouvé."}

    # INJECTION DU SYSTEM PROMPT
    messages_for_model = [
        {'role': 'system', 'content': SYSTEM_PROMPT},
        {'role': 'user', 'content': user_prompt}
    ]
    
    # CORRECTION DU WARNING : On passe l'attention_mask
    inputs = tokenizer.apply_chat_template(messages_for_model, add_generation_prompt=True, return_tensors="pt").to(DEVICE)
    
    # Génération de la réponse complète
    outputs = model.generate(
        inputs,
        attention_mask=inputs.attention_mask, # On ajoute l'attention_mask ici
        max_new_tokens=500, # Augmenté pour des réponses plus longues
        do_sample=True,
        temperature=0.1,
        top_k=50,
        top_p=0.95,
        num_return_sequences=1,
        eos_token_id=tokenizer.eos_token_id
    )
    response_text = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)

    async def stream_generator():
        response_id = f"chatcmpl-{uuid.uuid4()}"
        for char in response_text:
            chunk = {"id": response_id, "object": "chat.completion.chunk", "created": int(time.time()), "model": MODEL_ID, "choices": [{"index": 0, "delta": {"content": char}, "finish_reason": None}]}
            yield f"data: {json.dumps(chunk)}\n\n"
            await asyncio.sleep(0.005)
        final_chunk = {"id": response_id, "object": "chat.completion.chunk", "created": int(time.time()), "model": MODEL_ID, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}
        yield f"data: {json.dumps(final_chunk)}\n\n"
        yield "data: [DONE]\n\n"

    if request.stream:
        return StreamingResponse(stream_generator(), media_type="text/event-stream")
    else:
        return {"choices": [{"message": {"role": "assistant", "content": response_text}}]}

@app.get("/")
def root():
    return {"status": "Agent spécialisé WordPress/WooCommerce en ligne", "model_id": MODEL_ID}