from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from llmlingua import PromptCompressor

app = FastAPI()

# Initialize LLMLingua-2 for better compression
compressor = PromptCompressor(model_name="microsoft/llmlingua-2-bert-base-multilingual-cased-meetingbank",use_llmlingua2=True
, device_map="cpu")

class TextInput(BaseModel):
    text: str

@app.post("/compress")
async def compress_text(input: TextInput):
    try:
        # Compress the input text using LLMLingua-2
        compressed_text = compressor.compress_prompt(input.text)
        
        # Ensure the compressed text is of high quality
        if not compressed_text:
            raise HTTPException(status_code=400, detail="Compression failed, no output generated.")
        
        return {"compressed_text": compressed_text}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/")
async def root():
    return {"message": "Welcome to the LLMLingua Text Compression API"}