Update app.py

app.py

@@ -1,53 +1,135 @@
 from fastapi import FastAPI, Request
-from fastapi.responses import JSONResponse
+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
 
-# --- Création de l'application API de diagnostic ---
+# --- 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
+)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+print("Modèle et tokenizer chargés avec succès sur le CPU.")
+
+# --- Création de l'application API ---
 app = FastAPI()
 
-# --- Définition des API de diagnostic ---
+# --- Modèles de données pour accepter la structure complexe de l'extension ---
+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")
+# --- Définition des API ---
+
+@app.get("/models", response_model=ModelList)
 async def list_models():
-    """Répond à la requête GET /models pour que l'extension soit satisfaite."""
-    print("--- REQUETE RECUE SUR /models ---")
-    return {"object": "list", "data": [{"id": "deepseek-ai/deepseek-coder-1.3b-instruct", "object": "model"}]}
+    """Répond à la requête GET /models pour satisfaire l'extension."""
+    return ModelList(data=[ModelData(id=MODEL_ID)])
 
 @app.post("/chat/completions")
-async def debug_chat_completion(request: Request):
-    """
-    Endpoint qui ne fait qu'une seule chose :
-    afficher le contenu exact de la requête envoyée par VS Code.
-    """
-    print("\n\n" + "="*50)
-    print("=== REQUETE POST RECUE SUR /chat/completions ===")
-    
-    # On affiche les headers de la requête
-    print("\n--- HEADERS ---")
-    for name, value in request.headers.items():
-        print(f"{name}: {value}")
+async def create_chat_completion(request: ChatCompletionRequest):
+    """Endpoint principal qui gère la génération de texte en streaming."""
     
-    # On essaie de lire et d'afficher le corps (body) de la requête
-    try:
-        body = await request.json()
-        print("\n--- BODY (JSON) ---")
-        # On utilise json.dumps pour un affichage propre
-        print(json.dumps(body, indent=2))
-    except Exception as e:
-        print(f"\n--- ERREUR LORS DE LA LECTURE DU BODY ---")
-        print(f"L'erreur est : {e}")
-        # On essaie de lire le corps comme du texte brut si le JSON échoue
-        body_raw = await request.body()
-        print("\n--- BODY (BRUT) ---")
-        print(body_raw.decode('utf-8', errors='ignore'))
-
-    print("="*50 + "\n\n")
+    # On extrait le prompt de l'utilisateur de la structure complexe
+    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é."}
+
+    # Préparation pour le modèle DeepSeek
+    messages_for_model = [{'role': 'user', 'content': user_prompt}]
+    inputs = tokenizer.apply_chat_template(messages_for_model, add_generation_prompt=True, return_tensors="pt").to(DEVICE)
     
-    # On renvoie une erreur 422 volontairement pour pouvoir retenter facilement.
-    # Le but n'est pas que ça marche, mais de voir les logs.
-    error_message = {"error": "Mode diagnostic actif. Vérifiez les logs du Space."}
-    return JSONResponse(status_code=422, content=error_message)
+    # Génération de la réponse complète
+    outputs = model.generate(inputs, max_new_tokens=250, do_sample=True, temperature=0.2, 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)
+
+    # Fonction génératrice pour le streaming
+    async def stream_generator():
+        response_id = f"chatcmpl-{uuid.uuid4()}"
+        
+        # On envoie la réponse caractère par caractère, au format attendu
+        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.01) # Petite pause pour simuler un flux
+        
+        # On envoie le chunk final de fin
+        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"
+        
+        # On envoie le signal [DONE]
+        yield "data: [DONE]\n\n"
+
+    # Si l'extension demande un stream, on renvoie le générateur
+    if request.stream:
+        return StreamingResponse(stream_generator(), media_type="text/event-stream")
+    else:
+        # Code de secours si le stream n'est pas demandé (peu probable)
+        return {"choices": [{"message": {"role": "assistant", "content": response_text}}]}
 
 @app.get("/")
 def root():
-    return {"status": "API en mode diagnostic. Vérifiez les logs."}
+    return {"status": "API compatible OpenAI en ligne (avec streaming)", "model_id": MODEL_ID}
+
+# On a besoin de asyncio pour la pause dans le stream
+import asyncio