app.py

import os
import logging
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from langchain_community.llms import Ollama
from langchain.callbacks.manager import CallbackManager
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI()
MODEL_NAME = 'tinyllama'

def get_llm():
    callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])
    return Ollama(model=MODEL_NAME, callback_manager=callback_manager)

class Question(BaseModel):
    text: str

@app.get("/")
def read_root():
    return {"Hello": f"Welcome to {MODEL_NAME} FastAPI"}

@app.post("/generate")
def generate_response(q: Question):
    try:
        llm = get_llm()
        response = llm.invoke(q.text)
        return {"response": response}
    except Exception as e:
        logger.error(f"Error generating response: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.on_event("startup")
async def startup_event():
    logger.info(f"Starting up with model: {MODEL_NAME}")

@app.on_event("shutdown")
async def shutdown_event():
    logger.info("Shutting down")