Update app.py

app.py

@@ -1,7 +1,7 @@
 import os
 import re
-import requests
 import json
+import requests
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from google.auth import exceptions
@@ -18,28 +18,35 @@ GCS_BUCKET_NAME = os.getenv("GCS_BUCKET_NAME")
 GOOGLE_APPLICATION_CREDENTIALS_JSON = os.getenv("GOOGLE_APPLICATION_CREDENTIALS_JSON")
 HF_API_TOKEN = os.getenv("HF_API_TOKEN")
 
+def sanitize_bucket_name(bucket_name):
+    """Corrige un nombre de bucket inválido para ajustarse a las reglas de Google Cloud Storage."""
+    bucket_name = bucket_name.lower()
+    bucket_name = re.sub(r"[^a-z0-9-\.]", "-", bucket_name)
+    bucket_name = bucket_name.strip("-.")
+    if len(bucket_name) > 63:
+        bucket_name = bucket_name[:63]
+    if not re.match(r"^[a-z0-9]", bucket_name):
+        bucket_name = "a" + bucket_name
+    if not re.match(r"[a-z0-9]$", bucket_name):
+        bucket_name = bucket_name + "a"
+    return bucket_name
+
 def validate_bucket_name(bucket_name):
-    if not isinstance(bucket_name, str):
-        raise ValueError("Bucket name must be a string.")
-    if len(bucket_name) < 3 or len(bucket_name) > 63:
-        raise ValueError("Bucket name must be between 3 and 63 characters long.")
+    """Valida si el nombre de bucket cumple con las reglas de Google Cloud Storage."""
     if not re.match(r"^[a-z0-9][a-z0-9\-\.]*[a-z0-9]$", bucket_name):
-        raise ValueError(
-            f"Invalid bucket name '{bucket_name}'. Bucket names must:"
-            " - Use only lowercase letters, numbers, hyphens (-), and periods (.)"
-            " - Start and end with a letter or number."
-        )
-    if "--" in bucket_name or ".." in bucket_name or ".-" in bucket_name or "-." in bucket_name:
-        raise ValueError(
-            f"Invalid bucket name '{bucket_name}'. Bucket names cannot contain consecutive periods, hyphens, or use '.-' or '-.'"
-        )
+        raise ValueError(f"Nombre de bucket inválido: '{bucket_name}'. Debe cumplir con las reglas de GCS.")
     return bucket_name
 
 try:
+    # Sanitizar y validar el nombre del bucket
+    GCS_BUCKET_NAME = sanitize_bucket_name(GCS_BUCKET_NAME)
     GCS_BUCKET_NAME = validate_bucket_name(GCS_BUCKET_NAME)
+
+    # Cargar credenciales de Google Cloud Storage
     credentials_info = json.loads(GOOGLE_APPLICATION_CREDENTIALS_JSON)
     storage_client = storage.Client.from_service_account_info(credentials_info)
     bucket = storage_client.bucket(GCS_BUCKET_NAME)
+
 except (exceptions.DefaultCredentialsError, json.JSONDecodeError, KeyError, ValueError) as e:
     print(f"Error al cargar credenciales o bucket: {e}")
     exit(1)
@@ -61,7 +68,7 @@ class GCSStreamHandler:
     def stream_file_from_gcs(self, blob_name):
         blob = self.bucket.blob(blob_name)
         if not blob.exists():
-            raise HTTPException(status_code=404, detail=f"File '{blob_name}' not found in GCS.")
+            raise HTTPException(status_code=404, detail=f"Archivo '{blob_name}' no encontrado en GCS.")
         return blob.download_as_bytes()
 
     def upload_file_to_gcs(self, blob_name, data_stream):
@@ -103,7 +110,7 @@ def download_model_from_huggingface(model_name):
                 blob_name = f"{model_name}/{filename}"
                 blob = bucket.blob(blob_name)
                 blob.upload_from_file(BytesIO(response.content))
-        except Exception as e:
+        except Exception:
             pass
 
 @app.post("/predict/")
@@ -127,6 +134,8 @@ async def predict(request: DownloadModelRequest):
         model = AutoModelForCausalLM.from_pretrained(BytesIO(config_stream))
         tokenizer = AutoTokenizer.from_pretrained(BytesIO(tokenizer_stream))
         pipeline_task = request.pipeline_task
+        if pipeline_task not in ["text-generation", "sentiment-analysis", "translation", "fill-mask", "question-answering"]:
+            raise HTTPException(status_code=400, detail="Tarea no soportada")
         pipeline_ = pipeline(pipeline_task, model=model, tokenizer=tokenizer)
         input_text = request.input_text
         result = pipeline_(input_text)