Update app.py

app.py

@@ -14,16 +14,48 @@ from bs4 import BeautifulSoup
 from huggingface_hub import snapshot_download
 from llama_cpp import Llama
 
+
+# ------------------ Helpers ------------------
+def env_int(name: str, default: int) -> int:
+    try:
+        return int(os.getenv(name, "").strip() or default)
+    except Exception:
+        return default
+
+
+def env_float(name: str, default: float) -> float:
+    try:
+        return float(os.getenv(name, "").strip() or default)
+    except Exception:
+        return default
+
+
+def env_list(name: str) -> list[str]:
+    raw = os.getenv(name, "").strip()
+    return [x.strip() for x in raw.split(",") if x.strip()]
+
+
 # ------------------ Config ------------------
 MODEL_REPO = os.getenv("MODEL_REPO", "Qwen/Qwen2.5-7B-Instruct-GGUF")
-# Si te falta RAM en CPU Basic: exporta MODEL_PATTERN=qwen2.5-7b-instruct-q3_k_m-*.gguf
+# Si te falta RAM en CPU Basic, usa q3_k_m:
+# export MODEL_PATTERN=qwen2.5-7b-instruct-q3_k_m-*.gguf
 MODEL_PATTERN = os.getenv("MODEL_PATTERN", "qwen2.5-7b-instruct-q4_k_m-*.gguf")
 
 # Carpeta de modelos en /data (escribible en Docker Spaces)
 MODELS_DIR = Path(os.getenv("MODELS_DIR", "/data/models"))
 MODELS_DIR.mkdir(parents=True, exist_ok=True)
 
-N_THREADS = max(1, (os.cpu_count() or 2) - 1)
+# Rendimiento (overrides por variables de entorno)
+CPU_COUNT = os.cpu_count() or 2
+N_THREADS = env_int("N_THREADS", max(1, CPU_COUNT - 1))
+N_BATCH = env_int("N_BATCH", 64)          # bajar si vas muy justo de RAM
+N_CTX = env_int("N_CTX", 2048)            # 2048 va bien en CPU basic
+
+# Decodificación / longitud por defecto
+DEF_TEMPERATURE = env_float("LLM_TEMPERATURE", 0.6)
+DEF_TOP_P = env_float("LLM_TOP_P", 0.95)
+DEF_MAX_TOKENS = env_int("LLM_MAX_TOKENS", 160)          # longitud típica
+MAX_TOKENS_CAP = env_int("LLM_MAX_TOKENS_CAP", 320)      # tope duro
 
 SYSTEM_DEFAULT = textwrap.dedent("""\
 Eres Astrohunters-Guide, un asistente en español.
@@ -32,14 +64,19 @@ Eres Astrohunters-Guide, un asistente en español.
 - Si te paso una URL, lees su contenido y lo usas como contexto.
 """)
 
-ALLOWED_ORIGINS = [
-    # agrega tu dominio(s) aquí
-    "https://pruebas.nataliacoronel.com",
-    "https://*.nataliacoronel.com",
-    # durante pruebas puedes permitir todo, pero es menos seguro:
-    os.getenv("ALLOW_ALL_ORIGINS", "") and "*",
-]
-ALLOWED_ORIGINS = [o for o in ALLOWED_ORIGINS if o]
+# CORS
+# Opciones:
+#  - ALLOW_ALL_ORIGINS=1   (menos seguro, útil en pruebas)
+#  - CORS_ORIGINS="https://dominio1,https://dominio2"
+allow_all = os.getenv("ALLOW_ALL_ORIGINS", "").strip() in ("1", "true", "yes")
+CORS_ORIGINS = env_list("CORS_ORIGINS")
+if not CORS_ORIGINS:
+    # defaults cómodos para tu caso
+    CORS_ORIGINS = [
+        "https://pruebas.nataliacoronel.com",
+        "https://*.nataliacoronel.com",
+    ]
+
 
 # ------------------ Descarga del modelo ------------------
 print(f"[Boot] Descargando {MODEL_REPO} patrón {MODEL_PATTERN} en {MODELS_DIR} ...")
@@ -54,17 +91,22 @@ if not candidates:
 MODEL_PATH = candidates[0]
 print(f"[Boot] Usando shard: {MODEL_PATH}")
 
+
 # ------------------ Carga LLaMA.cpp ------------------
 llm = Llama(
     model_path=MODEL_PATH,
-    n_ctx=4096,
+    n_ctx=N_CTX,
     n_threads=N_THREADS,
-    n_batch=256,
-    n_gpu_layers=0,
+    n_batch=N_BATCH,
+    n_gpu_layers=env_int("N_GPU_LAYERS", 0),
     verbose=False,
 )
 
+
+# ------------------ Utilidades ------------------
 def fetch_url_text(url: str, max_chars: int = 6000) -> str:
+    if not url:
+        return ""
     try:
         r = requests.get(url, timeout=15)
         r.raise_for_status()
@@ -76,40 +118,85 @@ def fetch_url_text(url: str, max_chars: int = 6000) -> str:
     except Exception as e:
         return f"[No se pudo cargar {url}: {e}]"
 
-def run_llm(messages, temperature=0.6, top_p=0.95, max_tokens=768) -> str:
+
+def clamp_tokens(requested: Optional[int]) -> int:
+    if requested is None or requested <= 0:
+        return DEF_MAX_TOKENS
+    return max(1, min(requested, MAX_TOKENS_CAP))
+
+
+def run_llm(
+    messages,
+    temperature: Optional[float] = None,
+    top_p: Optional[float] = None,
+    max_tokens: Optional[int] = None,
+) -> str:
     out = llm.create_chat_completion(
         messages=messages,
-        temperature=temperature,
-        top_p=top_p,
-        max_tokens=max_tokens,
+        temperature=DEF_TEMPERATURE if temperature is None else float(temperature),
+        top_p=DEF_TOP_P if top_p is None else float(top_p),
+        max_tokens=clamp_tokens(max_tokens),
         stream=False,
     )
     return out["choices"][0]["message"]["content"].strip()
 
+
 # ------------------ FastAPI ------------------
 app = FastAPI(title="Astrohunters LLM API", docs_url="/docs", redoc_url=None)
 
-if ALLOWED_ORIGINS:
+if allow_all:
     app.add_middleware(
         CORSMiddleware,
-        allow_origins=ALLOWED_ORIGINS,
+        allow_origins=["*"],
+        allow_credentials=False,
+        allow_methods=["*"],
+        allow_headers=["*"],
+    )
+else:
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=CORS_ORIGINS,
         allow_credentials=True,
         allow_methods=["*"],
         allow_headers=["*"],
     )
 
+
+# --------- Esquemas de entrada ---------
 class PredictIn(BaseModel):
     prompt: str
     system: Optional[str] = None
+    max_tokens: Optional[int] = None
+    temperature: Optional[float] = None
+    top_p: Optional[float] = None
+
 
 class PredictURLIn(BaseModel):
     prompt: str
     url: Optional[str] = None
     system: Optional[str] = None
+    max_tokens: Optional[int] = None
+    temperature: Optional[float] = None
+    top_p: Optional[float] = None
+
 
+# --------- Endpoints ---------
 @app.get("/healthz")
 def healthz():
-    return {"ok": True, "model": os.path.basename(MODEL_PATH), "threads": N_THREADS}
+    return {
+        "ok": True,
+        "model": os.path.basename(MODEL_PATH),
+        "threads": N_THREADS,
+        "n_batch": N_BATCH,
+        "n_ctx": N_CTX,
+        "defaults": {
+            "temperature": DEF_TEMPERATURE,
+            "top_p": DEF_TOP_P,
+            "max_tokens": DEF_MAX_TOKENS,
+            "max_tokens_cap": MAX_TOKENS_CAP,
+        },
+    }
+
 
 @app.get("/")
 def root():
@@ -118,15 +205,22 @@ def root():
         "endpoints": ["/healthz", "/run_predict", "/run_predict_with_url", "/docs"],
     }
 
+
 @app.post("/run_predict")
 def run_predict(body: PredictIn):
     messages = [
         {"role": "system", "content": body.system or SYSTEM_DEFAULT},
         {"role": "user", "content": body.prompt},
     ]
-    reply = run_llm(messages, max_tokens=512)
+    reply = run_llm(
+        messages,
+        temperature=body.temperature,
+        top_p=body.top_p,
+        max_tokens=body.max_tokens,
+    )
     return {"reply": reply}
 
+
 @app.post("/run_predict_with_url")
 def run_predict_with_url(body: PredictURLIn):
     web_ctx = fetch_url_text(body.url) if body.url else ""
@@ -135,9 +229,15 @@ def run_predict_with_url(body: PredictURLIn):
         {"role": "system", "content": body.system or SYSTEM_DEFAULT},
         {"role": "user", "content": user_msg},
     ]
-    reply = run_llm(messages, max_tokens=700)
+    reply = run_llm(
+        messages,
+        temperature=body.temperature,
+        top_p=body.top_p,
+        max_tokens=body.max_tokens,
+    )
     return {"reply": reply}
 
+
 if __name__ == "__main__":
-    import uvicorn, os
+    import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", "7860")))