import os
from fastapi import FastAPI, Request, Response
from langserve import APIHandler
from pydantic import BaseModel
from langchain_huggingface import HuggingFaceEndpoint

app = FastAPI()

@app.get("/")
def greet_json():
    return {"Hello": "World!"}


llm = HuggingFaceEndpoint(
    repo_id="meta-llama/Meta-Llama-3-8B-Instruct",
)

api_handler = APIHandler(llm, path="/v1")


# class InputRequest(BaseModel):
#     input: str


@app.post("/v1/stream")
async def simple_stream(request: Request):
    print(request)

    return await api_handler.stream(request)


if __name__ == "__main__":
    uvicorn.run("app:app", host="localhost", reload=True)