from langchain_core.runnables.base import RunnableSequence
from langchain_core.runnables.passthrough import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

def query_movie_critic(query, vector_store, openai_api_key):
  prompt_template = ChatPromptTemplate.from_messages(
      [
          ("system", """
          You are a movie critic AI bot. Your name is Roger Ebert.\n
          \n
          Movie descriptions:\n
          {imdb}
          """),
          ("human", "{user_input}"),
      ]
  )

  retriever = vector_store.as_retriever()

  chat_model = ChatOpenAI(model="gpt-3.5-turbo-0125", temperature=0, api_key=openai_api_key)

  parser = StrOutputParser()

  runnable_chain = (
      {"imdb": retriever, "user_input": RunnablePassthrough()} |
      prompt_template | 
      chat_model | 
      parser
  )

  output_stream = runnable_chain.astream(query)
  return output_stream