from langchain_openai import AzureChatOpenAI
from langfuse.callback import CallbackHandler
from langchain.prompts import PromptTemplate
from langchain.prompts.chat import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    AIMessagePromptTemplate,
    HumanMessagePromptTemplate,
)

llm = AzureChatOpenAI(
        azure_endpoint=os.environ.get("AZURE_OPENAI_API_BASE"),
        azure_deployment="gpt35-turbo",
        openai_api_version="2023-03-15-preview",
        openai_api_key=os.environ.get("AZURE_OPENAI_API_KEY"),
    )

class OpenAIRequest(BaseModel):
    system_message: str = "You are a helpful assistant"
    query: str
    temperature: float = 0
    model_name: CHAT_MODELS = CHAT_MODELS.gpt35_turbo
    output_parser: Optional[Literal["list", "object"]] = None
    object_schema: Optional[dict] = (
        None  # Example of how to setup pydantic model (MUST NAME CLASS OutputClassGPT): https://python.langchain.com/docs/modules/model_io/output_parsers/pydantic
    )
    source: str
    user: str = None
    project_name: str = None # will be used for tracing in langfuse
    metadata: Optional[dict] = None
    version: Optional[str] = None

@llm_app.post("/generate/chat")
def generate_chat(request: OpenAIRequest):

system_message_prompt = SystemMessagePromptTemplate.from_template(
            "you are a helpful assistant"
        )
prompt = PromptTemplate(
                template=request.prompt,
                input_variables=[],
            )
human_message_prompt = HumanMessagePromptTemplate(prompt=prompt)
chat_prompt = ChatPromptTemplate.from_messages(
            [system_message_prompt, human_message_prompt]
        )

formatted_prompt = chat_prompt.format_prompt().to_messages()

trace = langfuse.trace(
name=request.project_name,
user_id=request.user,
tags=[request.source],
metadata=request.metadata if request.metadata else {},
version=request.version if request.version else "1",
input=formatted_prompt,
)
langfuse_handler = trace.get_langchain_handler()
resp = chatgpt.invoke(formatted_prompt, config={"callbacks": [langfuse_handler]})
result = resp.content
trace.update(output=result)