from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import subprocess
import logging

app = FastAPI()

# 这里设置日志级别为 DEBUG
logging.basicConfig(level=logging.DEBUG)

class PromptRequest(BaseModel):
    prompt: str

@app.on_event("startup")
async def startup_event():
    logging.debug("Starting application...")

@app.post("/generate")
async def generate_response(request: PromptRequest):
    try:
        logging.debug(f"Received prompt: {request.prompt}")
        result = subprocess.check_output(["ollama", "run", request.prompt], text=True)
        logging.debug(f"Model response: {result.strip()}")
        return {"response": result.strip()}
    except subprocess.CalledProcessError as e:
        logging.error(f"Error calling ollama: {e}")
        raise HTTPException(status_code=500, detail=f"Error: {e}")

@app.get("/")
def read_root():
    logging.debug("Received GET request on root")
    return {"message": "Ollama API is running"}

@app.get("/health")
def health_check():
    logging.debug("Health check received")
    return {"status": "ok"}