from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Optional
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
import uvicorn

app = FastAPI(title="ChemLLM CPU OpenAI API")

print("Завантаження GGUF моделі...")
model_path = hf_hub_download(
    repo_id="RichardErkhov/AI4Chem___ChemLLM-7B-Chat-1_5-DPO-gguf", 
    filename="ChemLLM-7B-Chat-1_5-DPO.Q4_K_M.gguf"
)

llm = Llama(model_path=model_path, n_ctx=2048, n_threads=4)
print("Модель успішно завантажена на CPU!")

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

class ChatCompletionRequest(BaseModel):
    model: str
    messages: List[ChatMessage]
    temperature: Optional[float] = 0.7
    max_tokens: Optional[int] = 256

@app.post("/v1/chat/completions")
async def chat_completions(request: ChatCompletionRequest):
    try:
        full_prompt = ""
        for msg in request.messages:
            if msg.role == "user":
                full_prompt += f"<|User|>:{msg.content}"
            elif msg.role == "assistant":
                full_prompt += f"<|Bot|>:{msg.content}"
        full_prompt += "<|Bot|>:"

        # Виклик генерації на CPU
        output = llm(
            full_prompt,
            max_tokens=request.max_tokens,
            temperature=request.temperature,
            stop=["<|User|>", "<|Bot|>", "\n\n"]
        )
        
        response_text = output["choices"][0]["text"].strip()

        return {
            "id": "chatcmpl-chem-cpu",
            "object": "chat.completion",
            "model": request.model,
            "choices": [{
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": response_text
                },
                "finish_reason": "stop"
            }]
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/")
def health():
    return {"status": "healthy", "hardware": "CPU"}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)