Update app.py

app.py

@@ -1,16 +1,21 @@
 import os
 import secrets
-import html
-import asyncio
 import logging
+import asyncio
+import html
+
 from dataclasses import dataclass
+from datetime import datetime
 from typing import Any, Optional, Tuple
 
 import gradio as gr
 from transformers import pipeline
+from dotenv import load_dotenv
 from pydantic import BaseModel
 
 # ----------------- Configuration & Models -----------------
+load_dotenv()
+
 
 @dataclass
 class Config:
@@ -35,14 +40,29 @@ class APIResponse(BaseModel):
 
 
 # ----------------- Enhanced Logger -----------------
+def setup_logger() -> logging.Logger:
+    cfg = Config()
+    log_level = getattr(logging, cfg.LOG_LEVEL.upper(), logging.INFO)
+
+    # Avoid adding duplicate handlers if called multiple times
+    logger = logging.getLogger("gemma_saas")
+    if not logger.handlers:
+        logger.setLevel(log_level)
 
-def setup_logger():
-    logging.basicConfig(
-        level=getattr(logging, Config().LOG_LEVEL),
-        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
-        handlers=[logging.FileHandler('gemma_saas.log'), logging.StreamHandler()]
-    )
-    return logging.getLogger(__name__)
+        formatter = logging.Formatter(
+            "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+        )
+
+        file_handler = logging.FileHandler("gemma_saas.log")
+        file_handler.setFormatter(formatter)
+
+        stream_handler = logging.StreamHandler()
+        stream_handler.setFormatter(formatter)
+
+        logger.addHandler(file_handler)
+        logger.addHandler(stream_handler)
+
+    return logger
 
 
 logger = setup_logger()
@@ -55,22 +75,26 @@ class ModelManager:
         self.pipeline = None
         self.model_loaded = False
 
-    async def initialize(self):
+    async def initialize(self) -> None:
         if not self.config.HF_TOKEN:
-            logger.error("Token do Hugging Face não encontrado. O carregamento do modelo irá falhar.")
+            logger.error(
+                "Token do Hugging Face não encontrado. O carregamento do modelo irá falhar."
+            )
             return
 
         try:
             logger.info(f"A carregar o modelo: {self.config.MODEL_NAME}...")
-            loop = asyncio.get_running_loop()
+
+            loop = asyncio.get_event_loop()
 
             def load_pipeline():
-                # Use `use_auth_token` (aplicável em muitas versões do transformers)
+                # Nota: dependendo da versão do transformers/huggingface-hub, 
+                # o parâmetro pode ser use_auth_token ou token. Ajuste se necessário.
                 return pipeline(
-                    task="text-generation",
+                    "text-generation",
                     model=self.config.MODEL_NAME,
-                    device_map="auto",
                     model_kwargs={"torch_dtype": "auto"},
+                    device_map="auto",
                     use_auth_token=self.config.HF_TOKEN,
                 )
 
@@ -78,22 +102,34 @@ class ModelManager:
             self.model_loaded = True
             logger.info("✅ Modelo carregado com sucesso!")
         except Exception as e:
-            logger.error(f"❌ Erro ao carregar o modelo: {e}")
-            self.model_loaded = False
+            logger.error(f"❌ Erro ao carregar o modelo: {e}", exc_info=True)
 
     async def generate(self, request: GenerationRequest) -> Tuple[bool, str, int]:
         if not self.model_loaded or self.pipeline is None:
-            return False, "❌ O modelo não está disponível. Por favor, verifique os logs do servidor.", 0
-
-        if not request.prompt.strip():
-            return False, "⚠️ O prompt não pode estar vazio.", 0
+            return (
+                False,
+                "❌ O modelo não está disponível. Por favor, verifique os logs do servidor.",
+                0,
+            )
 
         try:
-            loop = asyncio.get_running_loop()
+            if not request.prompt.strip():
+                return False, "⚠️ O prompt não pode estar vazio.", 0
+
+            loop = asyncio.get_event_loop()
+            messages = [{"role": "user", "content": request.prompt.strip()}]
 
             def do_generation():
-                # Para a maioria dos modelos de geração textual, passamos o prompt diretamente
-                prompt_text = request.prompt.strip()
+                # A pipeline que usava apply_chat_template é específica de alguns tokenizers;
+                # mantemos o uso mas com fallback simples se não existir.
+                tokenizer = getattr(self.pipeline, "tokenizer", None)
+                if tokenizer and hasattr(tokenizer, "apply_chat_template"):
+                    prompt_text = tokenizer.apply_chat_template(
+                        messages, tokenize=False, add_generation_prompt=True
+                    )
+                else:
+                    # Fallback simples
+                    prompt_text = request.prompt.strip()
 
                 outputs = self.pipeline(
                     prompt_text,
@@ -104,20 +140,17 @@ class ModelManager:
                     top_p=request.top_p,
                 )
 
-                # A saída típica é uma lista com dicionários contendo 'generated_text'
                 generated_text = outputs[0].get("generated_text", "")
+                # Se usamos prompt_text, retirar a parte inicial que corresponde ao prompt
+                if generated_text.startswith(prompt_text):
+                    generated_text = generated_text[len(prompt_text) :]
 
-                # Contagem aproximada de tokens (usa o tokenizer do pipeline se disponível)
                 tokens_used = 0
-                try:
-                    tokenizer = getattr(self.pipeline, "tokenizer", None)
-                    if tokenizer is not None:
-                        # Evitar adicionar special tokens na contagem
-                        tokens_used = len(tokenizer.encode(generated_text, add_special_tokens=False))
-                    else:
-                        tokens_used = len(generated_text.split())
-                except Exception:
-                    tokens_used = len(generated_text.split())
+                if tokenizer and hasattr(tokenizer, "encode"):
+                    try:
+                        tokens_used = len(tokenizer.encode(generated_text))
+                    except Exception:
+                        tokens_used = 0
 
                 return generated_text, tokens_used
 
@@ -125,7 +158,7 @@ class ModelManager:
             return True, generated_text, tokens_used
 
         except Exception as e:
-            logger.error(f"Erro na geração: {e}")
+            logger.error(f"Erro na geração: {e}", exc_info=True)
             return False, f"❌ A geração falhou: {str(e)}", 0
 
 
@@ -141,16 +174,17 @@ class GemmaService:
     async def generate_text(self, api_key: str, prompt: str, **kwargs) -> APIResponse:
         if not api_key or not api_key.startswith("gsk-"):
             return APIResponse(success=False, error="Chave de API inválida ou ausente.")
-
         try:
             request = GenerationRequest(prompt=prompt, **kwargs)
             success, text, tokens_used = await self.model_manager.generate(request)
             if success:
-                return APIResponse(success=True, data={"generated_text": text, "tokens_used": tokens_used})
+                return APIResponse(
+                    success=True, data={"generated_text": text, "tokens_used": tokens_used}
+                )
             else:
                 return APIResponse(success=False, error=text)
         except Exception as e:
-            logger.error(f"Erro de serviço durante a geração de texto: {e}")
+            logger.error(f"Erro de serviço durante a geração de texto: {e}", exc_info=True)
             return APIResponse(success=False, error="Ocorreu um erro interno no serviço.")
 
 
@@ -159,7 +193,7 @@ class GradioInterface:
     def __init__(self, service: GemmaService):
         self.service = service
 
-    def create_custom_css(self):
+    def create_custom_css(self) -> str:
         return """
         :root {
             --dark-bg: #0a0a0a; --panel-bg: #1a1a1a; --border-color: #333;
@@ -188,16 +222,29 @@ class GradioInterface:
         .gr-slider { color: var(--text-light); }
         """
 
-    def create_interface(self):
+    async def create_interface(self) -> gr.Blocks:
         with gr.Blocks(css=self.create_custom_css(), theme=None) as app:
             with gr.Row(elem_id="main_layout", equal_height=False):
                 with gr.Column(scale=2):
                     with gr.Column(elem_id="left_panel"):
-                        output_display = gr.Markdown(elem_id="output_display", value="<p style='color: #a0a0a0;'>A sua resposta aparecerá aqui...</p>")
+                        output_display = gr.Markdown(
+                            elem_id="output_display",
+                            value="<p style='color: #a0a0a0;'>A sua resposta aparecerá aqui...</p>",
+                        )
                         with gr.Column(elem_id="input_area"):
-                            api_key_input = gr.Textbox(label="A Sua Chave de API", placeholder="Cole a sua chave gsk-... aqui", type="password", elem_id="api_key_input")
+                            api_key_input = gr.Textbox(
+                                label="A Sua Chave de API",
+                                placeholder="Cole a sua chave gsk-... aqui",
+                                type="password",
+                                elem_id="api_key_input",
+                            )
                             with gr.Row():
-                                prompt_input = gr.Textbox(show_label=False, placeholder="Digite a sua mensagem...", elem_id="prompt_input", scale=10)
+                                prompt_input = gr.Textbox(
+                                    show_label=False,
+                                    placeholder="Digite a sua mensagem...",
+                                    elem_id="prompt_input",
+                                    scale=10,
+                                )
                                 send_button = gr.Button("➤ Enviar", elem_id="send_button", scale=2)
 
                 with gr.Column(scale=1):
@@ -206,60 +253,79 @@ class GradioInterface:
                         key_button = gr.Button("✨ Gerar Nova Chave", elem_id="generate_button")
 
                         with gr.Accordion("Parâmetros Avançados", open=False):
-                            temp_slider = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperatura")
-                            max_tokens_slider = gr.Slider(minimum=64, maximum=self.service.config.MAX_TOKENS, value=512, step=64, label="Max Tokens")
-                            top_k_slider = gr.Slider(minimum=1, maximum=100, value=50, step=1, label="Top-K")
-                            top_p_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-P")
+                            temp_slider = gr.Slider(
+                                minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperatura"
+                            )
+                            max_tokens_slider = gr.Slider(
+                                minimum=64, maximum=self.service.config.MAX_TOKENS, value=512, step=64, label="Max Tokens"
+                            )
+                            top_k_slider = gr.Slider(
+                                minimum=1, maximum=100, value=50, step=1, label="Top-K"
+                            )
+                            top_p_slider = gr.Slider(
+                                minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-P"
+                            )
 
                         gr.Markdown("### Como Usar a API")
-                        api_example_display = gr.HTML("<p style='color: #a0a0a0;'>Clique em 'Gerar Nova Chave' para ver um exemplo de código.</p>")
+                        api_example_display = gr.HTML(
+                            "<p style='color: #a0a0a0;'>Clique em 'Gerar Nova Chave' para ver um exemplo de código.</p>"
+                        )
 
             def handle_key_generation():
+                # Gera chave e mostra exemplo de payload
                 key = f"gsk-{secrets.token_urlsafe(24).replace('_', '').replace('-', '')}"
                 code_html = f"""
-                <div class=\"code-snippet\"> 
-                    <div><span class=\"keyword\">import</span> requests</div>
+                <div class="code-snippet">
+                    <div><span class="keyword">import</span> requests</div>
                     <div>&nbsp;</div>
-                    <div>url = <span class=\"string\">\"https://SEU_SPACE.hf.space/run/generate\"</span></div>
+                    <div>url = <span class="string">"https://SEU_SPACE.hf.space/run/generate"</span></div>
                     <div>payload = {{</div>
-                    <div>&nbsp;&nbsp;&nbsp;&nbsp;<span class=\"string\">\"api_key\"</span>: <span class=\"string\">\"{key}\"</span>,</div>
-                    <div>&nbsp;&nbsp;&nbsp;&nbsp;<span class=\"string\">\"prompt\"</span>: <span class=\"string\">\"Escreva um haikai sobre o universo\"</span>,</div>
-                    <div>&nbsp;&nbsp;&nbsp;&nbsp;<span class=\"string\">\"max_tokens\"</span>: <span class=\"number\">50</span></div>
+                    <div>&nbsp;&nbsp;&nbsp;&nbsp;<span class="string">"api_key"</span>: <span class="string">"{key}"</span>,</div>
+                    <div>&nbsp;&nbsp;&nbsp;&nbsp;<span class="string">"prompt"</span>: <span class="string">"Escreva um haikai sobre o universo"</span>,</div>
+                    <div>&nbsp;&nbsp;&nbsp;&nbsp;<span class="string">"max_tokens"</span>: <span class="number">50</span></div>
                     <div>}}</div>
                     <div>&nbsp;</div>
                     <div>response = requests.post(url, json=payload)</div>
-                    <div><span class=\"keyword\">print</span>(response.json())</div>
+                    <div><span class="keyword">print</span>(response.json())</div>
                 </div>
                 """
-
-                return gr.Textbox.update(value=key, interactive=True), api_example_display.update(value=code_html)
+                return key, gr.update(value=code_html)
 
             async def handle_generation(api_key, prompt, temp, max_tokens, top_k, top_p, btn):
+                # Validações básicas
                 if not api_key:
-                    yield "<p style='color: #FFCC00;'>Por favor, insira a sua chave de API para começar.</p>", gr.Button.update(value="➤ Enviar", interactive=True)
+                    yield (
+                        "<p style='color: #FFCC00;'>Por favor, insira a sua chave de API para começar.</p>",
+                        gr.update(value="➤ Enviar", interactive=True),
+                    )
                     return
+
                 if not prompt:
-                    yield "<p style='color: #FFCC00;'>Por favor, digite um prompt.</p>", gr.Button.update(value="➤ Enviar", interactive=True)
+                    yield (
+                        "<p style='color: #FFCC00;'>Por favor, digite um prompt.</p>",
+                        gr.update(value="➤ Enviar", interactive=True),
+                    )
                     return
 
-                yield "<p style='color: #a0a0a0;'>A gerar resposta...</p>", gr.Button.update(value="A gerar...", interactive=False)
+                # Indicador de carregamento
+                yield "<p style='color: #a0a0a0;'>A gerar resposta...</p>", gr.update(value="A gerar...", interactive=False)
 
-                # chama o serviço de geração
                 response = await self.service.generate_text(
                     api_key=api_key,
                     prompt=prompt,
-                    temperature=float(temp),
+                    temperature=temp,
                     max_tokens=int(max_tokens),
                     top_k=int(top_k),
-                    top_p=float(top_p),
+                    top_p=top_p,
                 )
 
                 if response.success:
                     formatted_text = html.escape(response.data["generated_text"]).replace("\n", "<br>")
-                    yield formatted_text, gr.Button.update(value="➤ Enviar", interactive=True)
+                    yield formatted_text, gr.update(value="➤ Enviar", interactive=True)
                 else:
-                    yield f"<p style='color: #FF4500;'>{response.error}</p>", gr.Button.update(value="➤ Enviar", interactive=True)
+                    yield f"<p style='color: #FF4500;'>{response.error}</p>", gr.update(value="➤ Enviar", interactive=True)
 
+            # Registar handlers
             send_button.click(
                 handle_generation,
                 inputs=[api_key_input, prompt_input, temp_slider, max_tokens_slider, top_k_slider, top_p_slider, send_button],
@@ -269,25 +335,30 @@ class GradioInterface:
 
             key_button.click(handle_key_generation, outputs=[api_key_input, api_example_display])
 
+            # Carregar exemplo inicial
+            app.load(
+                lambda: gr.update(value="<p style='color: #a0a0a0;'>Clique em 'Gerar Nova Chave' para ver um exemplo de código.</p>"),
+                [],  # sem inputs
+                [api_example_display],
+            )
+
         return app
 
 
 # ----------------- Main Application -----------------
-
-def main():
+async def main():
     try:
         service = GemmaService()
-        # inicializa o modelo (bloqueante, mas necessário antes de lançar a UI)
-        asyncio.run(service.initialize())
+        await service.initialize()
 
         interface = GradioInterface(service)
-        app = interface.create_interface()
+        app = await interface.create_interface()
 
-        # Lança a aplicação Gradio (bloqueia até terminar)
-        app.launch(server_name="0.0.0.0", server_port=7860, share=False, debug=False, show_error=True)
+        # Ligar a interface Gradio
+        app.launch(server_name="0.0.0.0", server_port=7860, share=False, debug=False)
     except Exception as e:
         logger.critical(f"Falha ao iniciar a aplicação: {e}", exc_info=True)
 
 
 if __name__ == "__main__":
-    main()
+    asyncio.run(main())