Update app.py

app.py

@@ -3,122 +3,133 @@ from fastapi.responses import StreamingResponse
 from pydantic import BaseModel, Extra
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
-import time, uuid, json, asyncio, requests
-from bs4 import BeautifulSoup
-from typing import Optional, List, Union
-import re
-
-# ==============================================================================
-# ===                BLOC DE CONFIGURATION DE L'AGENT                        ===
-# ==============================================================================
-
-MODEL_CONFIG = {
-    # L'identifiant du modèle depuis le Hub Hugging Face.
-    "MODEL_ID": "Qwen/Qwen2.5-Coder-0.5B-Instruct",
-
-    # Paramètres de génération
-    "MAX_TOKENS_ANSWER": 1500,
-    "TEMPERATURE": 0.2,
-}
-
-# --- NOUVEAU CERVEAU DE L'AGENT (SYSTEM PROMPT TEMPLATE) ---
-# On crée un template avec un emplacement {user_instructions}
-SYSTEM_PROMPT_TEMPLATE = """
-This work will be evaluated by a committee of the smartest thinkers to determine if you deserve to be classified as a reliable model. Be worthy of winning and bring out your best, now and in the future,Understanding the intent of user_input & respond in user language only.
-
-The user's specific instructions for this task are provided below. You must follow them precisely.
----
-USER INSTRUCTIONS:
-{user_instructions}
----
-"""
-
-# ==============================================================================
-# ===          FIN DU BLOC DE CONFIGURATION - NE PAS MODIFIER CI-DESSOUS     ===
-# ==============================================================================
-
-# --- Configuration & Model Loading ---
+import time
+import uuid
+import json
+from typing import Optional, List, Union, Dict, Any
+
+# --- Configuration ---
+MODEL_ID = "Qwen/Qwen2.5-Coder-0.5B-Instruct"
 DEVICE = "cpu"
-print(f"Loading model: {MODEL_CONFIG['MODEL_ID']}")
-model = AutoModelForCausalLM.from_pretrained(MODEL_CONFIG['MODEL_ID'], torch_dtype=torch.bfloat16, device_map=DEVICE)
-tokenizer = AutoTokenizer.from_pretrained(MODEL_CONFIG['MODEL_ID'], padding_side='left')
-tokenizer.pad_token = tokenizer.eos_token
-print("Model and tokenizer loaded successfully.")
 
+# --- 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()
 
-# --- Pydantic Models ---
-class ContentPart(BaseModel): type: str; text: str
-class ChatMessage(BaseModel): role: str; content: Union[str, List[ContentPart]]
+# --- 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]
+    
+    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]
+
+# --- Définition des API ---
 
-# --- API Endpoints ---
 @app.get("/models", response_model=ModelList)
 async def list_models():
-    return ModelList(data=[ModelData(id=MODEL_CONFIG['MODEL_ID'])])
+    """Répond à la requête GET /models pour satisfaire l'extension."""
+    return ModelList(data=[ModelData(id=MODEL_ID)])
 
 @app.post("/chat/completions")
 async def create_chat_completion(request: ChatCompletionRequest):
+    """Endpoint principal qui gère la génération de texte en streaming."""
+    
+    # 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 not found."}
-
-    async def stream_direct_response():
+            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)
+    
+    # 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()}"
-
-        def stream_chunk(content: str):
-            chunk = {"id": response_id, "object": "chat.completion.chunk", "created": int(time.time()), "model": MODEL_CONFIG['MODEL_ID'], "choices": [{"index": 0, "delta": {"content": content}, "finish_reason": None}]}
-            return f"data: {json.dumps(chunk)}\n\n"
-
-        # --- LOGIQUE DYNAMIQUE ---
-        # 1. On injecte l'input de l'utilisateur dans le template du system prompt
-        final_system_prompt = SYSTEM_PROMPT_TEMPLATE.format(user_instructions=user_prompt)
-        
-        # 2. On crée le message pour le modèle. Le rôle 'user' devient un simple déclencheur.
-        messages = [
-            {'role': 'system', 'content': final_system_prompt},
-            {'role': 'user', 'content': "Based on the detailed instructions I provided in the system prompt, generate the required response."}
-        ]
         
-        # On prépare les données pour le modèle
-        formatted_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        inputs = tokenizer(formatted_prompt, return_tensors="pt", padding=True).to(DEVICE)
-        
-        # On génère la réponse
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=MODEL_CONFIG['MAX_TOKENS_ANSWER'],
-            do_sample=True,
-            temperature=MODEL_CONFIG['TEMPERATURE'],
-            top_k=50,
-            top_p=0.95,
-            eos_token_id=tokenizer.eos_token_id
-        )
-        response_text = tokenizer.decode(outputs[0][len(inputs['input_ids'][0]):], skip_special_tokens=True)
-
-        # On streame la réponse finale
+        # On envoie la réponse caractère par caractère, au format attendu
         for char in response_text:
-            yield stream_chunk(char)
-            await asyncio.sleep(0.005)
+            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
         
-        # --- Fin du stream ---
-        final_chunk = {"id": response_id, "object": "chat.completion.chunk", "created": int(time.time()), "model": MODEL_CONFIG['MODEL_ID'], "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}
+        # 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"
 
-    return StreamingResponse(stream_direct_response(), media_type="text/event-stream")
+    # 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": "Dynamic Context Agent is online", "model_id": MODEL_CONFIG['MODEL_ID']}
+    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