import os
from fastapi import FastAPI
from pydantic import BaseModel
from llama_cpp import Llama
from huggingface_hub import hf_hub_download

# Nome do modelo no Hugging Face Hub
REPO_ID = "darkshop/olama"
FILENAME = "Llama-1B-Instruct.gguf"

# Caminho local do modelo
MODEL_FILE = os.path.join(os.getcwd(), FILENAME)

# Baixa o modelo se não existir
if not os.path.exists(MODEL_FILE):
    print(f"Baixando {FILENAME} do Hub...")
    MODEL_FILE = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)

# Inicializa o Llama
llm = Llama(model_path=MODEL_FILE)

# Cria a API FastAPI
app = FastAPI(title="Darkshop Llama API")

class Prompt(BaseModel):
    text: str
    max_tokens: int = 200

@app.post("/gerar")
def gerar_texto(prompt: Prompt):
    result = llm(prompt.text, max_tokens=prompt.max_tokens)
    return {"texto": result['text']}

@app.get("/")
def root():
    return {"mensagem": "API do modelo darkshop/olama ativa!"}

# Para rodar localmente:
# uvicorn app:app --host 0.0.0.0 --port 7860