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

app = FastAPI()

# Initialize LLMLingua
compressor = PromptCompressor(model_name="openai-community/gpt2", device_map="cpu")

class TextInput(BaseModel):
    text: str

class TextOutput(BaseModel):
    compressed_text: str

@app.post("/compress", response_model=TextOutput)
async def compress_text(input: TextInput):
    try:
        compressed_text = compressor.compress_prompt(input.text)
        return TextOutput(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"}