from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from langchain_huggingface import ChatHuggingFace, HuggingFacePipeline
import uvicorn
import os

# Define request model
class PromptRequest(BaseModel):
    prompt: str

app = FastAPI(title="Local LLM Server")

# Global variable for model
model = None

@app.on_event("startup")
async def startup_event():
    global model
    print("Initializing TinyLlama... This will auto-download the model if not present.")
    
    # Initialize the pipeline as requested
    llm = HuggingFacePipeline.from_model_id(
        model_id="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
        task="text-generation",
        pipeline_kwargs=dict(
            temperature=0.7,
            max_new_tokens=512,
            do_sample=True 
        )
    )
    
    model = ChatHuggingFace(llm=llm)
    print("Model loaded and ready to serve!")

@app.post("/generate")
async def generate_text(request: PromptRequest):
    if not model:
        raise HTTPException(status_code=503, detail="Model is loading, please wait...")
    
    try:
        # Invoke the model
        result = model.invoke(request.prompt)
        # Return the content of the AIMessage
        return {"response": result.content}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    print("Starting server on http://127.0.0.1:8000")
    uvicorn.run(app, host="127.0.0.1", port=8000)