caiiofc/llm-agent-api

1from fastapi import FastAPI2from pydantic import BaseModel3from llama_cpp import Llama4from huggingface_hub import hf_hub_download5import os6import psutil7import multiprocessing8 9app = FastAPI(title="LLM Agent API", version="1.0.0")10 11class ChatRequest(BaseModel):12    message: str13    max_tokens: int = 10014    temperature: float = 0.715 16class ChatResponse(BaseModel):17    response: str18 19class LocalLLMAgent:20    def __init__(self):21        # Download do modelo se não existir22        model_path = "./llama-2-7b-chat.Q4_K_M.gguf"23        24        if not os.path.exists(model_path):25            print("📥 Baixando modelo Llama-2-7B-Chat (Q4_K_M)...", flush=True)26            print("   Isso pode levar alguns minutos...", flush=True)27            model_path = hf_hub_download(28                repo_id="TheBloke/Llama-2-7B-Chat-GGUF",29                filename="llama-2-7b-chat.Q4_K_M.gguf",30                local_dir="./"31            )32            print("✅ Modelo baixado com sucesso!", flush=True)33        else:34            print("📁 Modelo já existe, carregando...", flush=True)35        36        # Configura para usar todas as CPUs disponíveis37        n_threads = multiprocessing.cpu_count()38        print(f"🔧 Configurando llama-cpp-python:", flush=True)39        print(f"   - CPUs disponíveis: {n_threads}", flush=True)40        print(f"   - Threads: {n_threads}", flush=True)41        print(f"   - Contexto: 2048 tokens", flush=True)42        43        print("🚀 Inicializando modelo...", flush=True)44        self.llm = Llama(45            model_path=model_path,46            chat_format="llama-2",47            n_ctx=2048,48            n_threads=n_threads,49            n_threads_batch=n_threads,50            verbose=False51        )52        print(f"✅ Modelo carregado! Usando {n_threads} threads", flush=True)53        self.messages = [54            {"role": "system", "content": "Responda sempre em português brasileiro de forma natural e conversacional."}55        ]56    57    def chat(self, message: str, max_tokens: int = 100, temperature: float = 0.75) -> str:58        print(f"💬 Nova mensagem: {message[:50]}{'...' if len(message) > 50 else ''}")59        print(f"   Parâmetros: max_tokens={max_tokens}, temperature={temperature}")60        61        self.messages.append({"role": "user", "content": message})62        63        response = self.llm.create_chat_completion(64            messages=self.messages,65            max_tokens=max_tokens,66            temperature=temperature67        )68        69        assistant_message = response['choices'][0]['message']['content']70        self.messages.append({"role": "assistant", "content": assistant_message})71        72        print(f"✅ Resposta gerada ({len(assistant_message)} chars)")73        return assistant_message74 75# Inicializa o agente globalmente76agent = None77 78@app.on_event("startup")79async def startup_event():80    print("=== INICIANDO LLM AGENT API ===", flush=True)81    print(f"CPUs disponíveis: {multiprocessing.cpu_count()}", flush=True)82    print(f"Memória total: {round(psutil.virtual_memory().total / (1024**3), 2)} GB", flush=True)83    84    global agent85    agent = LocalLLMAgent()86    87    print("✅ API pronta para uso!", flush=True)88    print("Endpoints disponíveis:", flush=True)89    print("  - POST /chat", flush=True)90    print("  - GET /health", flush=True)91    print("  - GET /system", flush=True)92 93@app.post("/chat", response_model=ChatResponse)94async def chat_endpoint(request: ChatRequest):95    if agent is None:96        return ChatResponse(response="Modelo ainda carregando, tente novamente.")97    response = agent.chat(request.message, request.max_tokens, request.temperature)98    return ChatResponse(response=response)99 100@app.get("/health")101async def health_check():102    return {"status": "healthy"}103 104@app.get("/system")105async def system_info():106    cpu_count = multiprocessing.cpu_count()107    cpu_percent = psutil.cpu_percent(interval=1, percpu=True)108    memory = psutil.virtual_memory()109    110    return {111        "cpu_cores": cpu_count,112        "cpu_usage_per_core": cpu_percent,113        "cpu_usage_total": psutil.cpu_percent(interval=1),114        "memory_total_gb": round(memory.total / (1024**3), 2),115        "memory_used_gb": round(memory.used / (1024**3), 2),116        "memory_percent": memory.percent117    }118 119# Removido - uvicorn será executado pelo Dockerfile