Update app.py

app.py

@@ -1,100 +1,151 @@
 import os
-import json
 import time
-from typing import List, Optional
+import json
+import uuid
 import gradio as gr
-from fastapi import FastAPI, Request, HTTPException
+from fastapi import FastAPI, Request
 from fastapi.responses import StreamingResponse, JSONResponse
 from pydantic import BaseModel
-from huggingface_hub import InferenceClient
+from typing import List, Optional
+
+from vllm.engine.async_llm_engine import AsyncLLMEngine
+from vllm.engine.arg_utils import AsyncEngineArgs
+from vllm.sampling_params import SamplingParams
+from transformers import AutoTokenizer
+
+# --- Configuración del Modelo ---
+MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen3.8-27B")
+
+# Argumentos del motor vLLM
+engine_args = AsyncEngineArgs(
+    model=MODEL_NAME,
+    trust_remote_code=True,
+    gpu_memory_utilization=0.90,
+    max_model_len=4096,
+    tensor_parallel_size=1
+)
 
-# Configuración del modelo
-MODEL_ID = "Qwen/Qwen3.8-27B"  # Cambiar al modelo deseado
-HF_TOKEN = os.getenv("HF_TOKEN")
+# Inicializar motor y tokenizador
+engine = AsyncLLMEngine.from_engine_args(engine_args)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
 
-client = InferenceClient(model=MODEL_ID, token=HF_TOKEN)
-app = FastAPI(title="OpenAI Compatible HF Space API")
+app = FastAPI(title="vLLM Gradio + OpenAI Server")
 
-# --- Esquemas Pydantic para OpenAI API ---
+# --- Esquemas Pydantic OpenAI ---
 class ChatMessage(BaseModel):
     role: str
     content: str
 
 class ChatCompletionRequest(BaseModel):
-    model: Optional[str] = MODEL_ID
+    model: Optional[str] = MODEL_NAME
     messages: List[ChatMessage]
     temperature: Optional[float] = 0.7
+    top_p: Optional[float] = 0.95
     max_tokens: Optional[int] = 2048
     stream: Optional[bool] = False
 
-# --- Endpoints OpenAI ---
+# --- Endpoints OpenAI (/v1) ---
 @app.get("/v1/models")
 async def list_models():
     return {
         "object": "list",
-        "data": [{"id": MODEL_ID, "object": "model", "owned_by": "huggingface"}]
+        "data": [{"id": MODEL_NAME, "object": "model", "owned_by": "vllm"}]
     }
 
 @app.post("/v1/chat/completions")
 async def chat_completions(req: ChatCompletionRequest):
-    formatted_messages = [{"role": m.role, "content": m.content} for m in req.messages]
+    raw_messages = [{"role": m.role, "content": m.content} for m in req.messages]
+    
+    # Aplicar plantilla de chat del modelo
+    prompt = tokenizer.apply_chat_template(raw_messages, tokenize=False, add_generation_prompt=True)
+    
+    sampling_params = SamplingParams(
+        temperature=req.temperature if req.temperature > 0 else 0.0,
+        top_p=req.top_p,
+        max_tokens=req.max_tokens or 2048
+    )
+    
+    request_id = f"chatcmpl-{uuid.uuid4().hex}"
 
     if req.stream:
         async def stream_generator():
-            stream = client.chat_completion(
-                messages=formatted_messages,
-                temperature=req.temperature,
-                max_tokens=req.max_tokens,
-                stream=True
-            )
-            for chunk in stream:
-                delta = chunk.choices[0].delta.content or ""
-                payload = {
-                    "id": f"chatcmpl-{int(time.time())}",
+            prev_text = ""
+            results_generator = engine.generate(prompt, sampling_params, request_id)
+            async for request_output in results_generator:
+                text = request_output.outputs[0].text
+                delta = text[len(prev_text):]
+                prev_text = text
+                
+                chunk_data = {
+                    "id": request_id,
                     "object": "chat.completion.chunk",
                     "created": int(time.time()),
-                    "model": MODEL_ID,
+                    "model": MODEL_NAME,
                     "choices": [{"index": 0, "delta": {"content": delta}, "finish_reason": None}]
                 }
-                yield f"data: {json.dumps(payload)}\n\n"
+                yield f"data: {json.dumps(chunk_data)}\n\n"
             yield "data: [DONE]\n\n"
 
         return StreamingResponse(stream_generator(), media_type="text/event-stream")
     else:
-        response = client.chat_completion(
-            messages=formatted_messages,
-            temperature=req.temperature,
-            max_tokens=req.max_tokens,
-            stream=False
-        )
+        results_generator = engine.generate(prompt, sampling_params, request_id)
+        final_output = None
+        async for request_output in results_generator:
+            final_output = request_output
+            
+        response_text = final_output.outputs[0].text if final_output else ""
         return {
-            "id": f"chatcmpl-{int(time.time())}",
+            "id": request_id,
             "object": "chat.completion",
             "created": int(time.time()),
-            "model": MODEL_ID,
+            "model": MODEL_NAME,
             "choices": [{
                 "index": 0,
-                "message": {"role": "assistant", "content": response.choices[0].message.content},
+                "message": {"role": "assistant", "content": response_text},
                 "finish_reason": "stop"
             }]
         }
 
-# --- Interfaz Web Gradio ---
-def gradio_chat(message, history):
+# --- Interfaz de Gradio ---
+async def gradio_chat(message, history, system_prompt, temperature, max_tokens):
     messages = []
+    if system_prompt:
+        messages.append({"role": "system", "content": system_prompt})
+        
     for user_msg, bot_msg in history:
         messages.append({"role": "user", "content": user_msg})
         messages.append({"role": "assistant", "content": bot_msg})
+        
     messages.append({"role": "user", "content": message})
+    
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    sampling_params = SamplingParams(
+        temperature=temperature if temperature > 0 else 0.0,
+        max_tokens=max_tokens
+    )
+    
+    request_id = f"gradio-{uuid.uuid4().hex}"
+    results_generator = engine.generate(prompt, sampling_params, request_id)
+    
+    async for request_output in results_generator:
+        yield request_output.outputs[0].text
 
-    response = ""
-    for chunk in client.chat_completion(messages=messages, stream=True, max_tokens=2048):
-        response += chunk.choices[0].delta.content or ""
-        yield response
+with gr.Blocks(title="vLLM + OpenAI Compatible Server") as demo:
+    gr.Markdown(f"# 🚀 vLLM LLM Engine\n**Modelo cargado:** `{MODEL_NAME}`\n\n**Base URL para Cursor / OpenCode:** `https://huggingface.co/spaces/acsaco.hf.space/v1`")
+    
+    with gr.Accordion("Configuración Avanzada", open=False):
+        system_prompt = gr.Textbox(label="System Prompt", value="Eres un asistente experto en programación.")
+        temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.7, step=0.05, label="Temperature")
+        max_tokens = gr.Slider(minimum=128, maximum=4096, value=2048, step=128, label="Max Tokens")
+        
+    gr.ChatInterface(
+        fn=gradio_chat,
+        additional_inputs=[system_prompt, temperature, max_tokens]
+    )
 
-with gr.Blocks(title="AI Chat & API Hub") as demo:
-    gr.Markdown(f"# 🤖 Interfaz de Chat & API\n**Modelo Activo:** `{MODEL_ID}`\n\nConéctate desde Cursor/OpenCode usando:\n`https:huggingface.co/spaces/acsaco.hf.space/v1`")
-    gr.ChatInterface(fn=gradio_chat)
+# Montar interfaz web Gradio dentro del servidor FastAPI
+app = gr.mount_gradio_app(app, demo, path="/")
 
-# Montar Gradio dentro de FastAPI
-app = gr.mount_gradio_app(app, demo, path="/")
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)