import os
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from llama_cpp import Llama

app = FastAPI()

# AIの脳みそを直接読み込む
print("🧠 脳みそ（モデル）を読み込み中...")
llm = Llama(model_path="./model/qwen2.5-coder-1.5b-instruct-q4_k_m.gguf", n_ctx=10000, n_threads=2)
print("✨ 脳みその読み込み完了！")

# ブラウザで見に来た時の確認用画面（これでもうエラー画面にはなりません！）
@app.get("/")
def read_root():
    return {"status": "AI is awake and ready!", "message": "マスター、通信は正常です。AIは起動しています。"}

# マスター専用の通信受付口（Termuxからの通信を直接脳みそに届ける）
@app.post("/v1/completions")
async def create_completion(request: Request):
    try:
        data = await request.json()
        prompt = data.get("prompt", "")
        max_tokens = data.get("max_tokens", 1024)
        
        # 脳みそに直接考えさせる
        output = llm(prompt, max_tokens=max_tokens, stop=["<|im_end|>"])
        
        return JSONResponse(content={"choices": [{"text": output["choices"][0]["text"]}]})
    except Exception as e:
        return JSONResponse(status_code=500, content={"error": str(e)})

if __name__ == "__main__":
    import uvicorn
    # どんな複雑な通信でも絶対に受け入れる最強設定で起動
    uvicorn.run(app, host="0.0.0.0", port=7860)