from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from llama_cpp import Llama
from typing import List, Dict

app = FastAPI(title="Gemma 4 API Server")

# Инициализация модели со строгим лимитом на 2 потока CPU (для стабильности бесплатного тарифа)
print("Загрузка модели Gemma 4...")
llm = Llama(
    model_path="/code/model.gguf",
    n_ctx=2048,
    n_threads=2
)
print("Модель успешно загружена!")

class Message(BaseModel):
    role: str
    content: str

class ChatCompletionRequest(BaseModel):
    model: str = "local"
    messages: List[Dict[str, str]]
    temperature: float = 0.7
    max_tokens: int = 500

@app.post("/v1/chat/completions")
async def chat_completions(request: ChatCompletionRequest):
    try:
        # Формируем промпт из истории сообщений
        prompt = ""
        for msg in request.messages:
            role = msg.get("role", "user")
            content = msg.get("content", "")
            prompt += f"<|im_start|>{role}\n{content}<|im_end|>\n"
        prompt += "<|im_start|>assistant\n"

        # Генерация ответа
        output = llm(
            prompt,
            max_tokens=request.max_tokens,
            temperature=request.temperature,
            stop=["<|im_end|>", "<|im_start|>"]
        )
        
        text_response = output["choices"][0]["text"].strip()

        # Формируем ответ, идентичный OpenAI API структуре
        return {
            "choices": [
                {
                    "message": {
                        "role": "assistant",
                        "content": text_response
                    },
                    "finish_reason": "stop",
                    "index": 0
                }
            ]
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/")
def read_root():
    return {"status": "Gemma 4 API is running"}