Integrations LangChain question answering with self query retriever

!python3 -m pip install -qU lark elasticsearch langchain langchain-elasticsearch openai

from langchain.schema import Document
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain_elasticsearch import ElasticsearchStore
from langchain.llms import OpenAI
from langchain.retrievers.self_query.base import SelfQueryRetriever
from langchain.chains.query_constructor.base import AttributeInfo
from getpass import getpass

docs = [
    Document(
        page_content="A bunch of scientists bring back dinosaurs and mayhem breaks loose",
        metadata={
            "year": 1993,
            "rating": 7.7,
            "genre": "science fiction",
            "director": "Steven Spielberg",
            "title": "Jurassic Park",
        },
    ),
    Document(
        page_content="Leo DiCaprio gets lost in a dream within a dream within a dream within a ...",
        metadata={
            "year": 2010,
            "director": "Christopher Nolan",
            "rating": 8.2,
            "title": "Inception",
        },
    ),
    Document(
        page_content="A psychologist / detective gets lost in a series of dreams within dreams within dreams and Inception reused the idea",
        metadata={
            "year": 2006,
            "director": "Satoshi Kon",
            "rating": 8.6,
            "title": "Paprika",
        },
    ),
    Document(
        page_content="A bunch of normal-sized women are supremely wholesome and some men pine after them",
        metadata={
            "year": 2019,
            "director": "Greta Gerwig",
            "rating": 8.3,
            "title": "Little Women",
        },
    ),
    Document(
        page_content="Toys come alive and have a blast doing so",
        metadata={
            "year": 1995,
            "genre": "animated",
            "director": "John Lasseter",
            "rating": 8.3,
            "title": "Toy Story",
        },
    ),
    Document(
        page_content="Three men walk into the Zone, three men walk out of the Zone",
        metadata={
            "year": 1979,
            "rating": 9.9,
            "director": "Andrei Tarkovsky",
            "genre": "science fiction",
            "rating": 9.9,
            "title": "Stalker",
        },
    ),
]

# https://www.elastic.co/search-labs/tutorials/install-elasticsearch/elastic-cloud#finding-your-cloud-id
ELASTIC_CLOUD_ID = getpass("Elastic Cloud ID: ")

# https://www.elastic.co/search-labs/tutorials/install-elasticsearch/elastic-cloud#creating-an-api-key
ELASTIC_API_KEY = getpass("Elastic Api Key: ")

# https://platform.openai.com/api-keys
OPENAI_API_KEY = getpass("OpenAI API key: ")

embeddings = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY)


vectorstore = ElasticsearchStore.from_documents(
    docs,
    embeddings,
    index_name="elasticsearch-self-query-demo",
    es_cloud_id=ELASTIC_CLOUD_ID,
    es_api_key=ELASTIC_API_KEY,
)

# Add details about metadata fields
metadata_field_info = [
    AttributeInfo(
        name="genre",
        description="The genre of the movie. Can be either 'science fiction' or 'animated'.",
        type="string or list[string]",
    ),
    AttributeInfo(
        name="year",
        description="The year the movie was released",
        type="integer",
    ),
    AttributeInfo(
        name="director",
        description="The name of the movie director",
        type="string",
    ),
    AttributeInfo(
        name="rating", description="A 1-10 rating for the movie", type="float"
    ),
]

document_content_description = "Brief summary of a movie"

# Set up openAI llm with sampling temperature 0
llm = OpenAI(temperature=0, openai_api_key=OPENAI_API_KEY)

# instantiate retriever
retriever = SelfQueryRetriever.from_llm(
    llm, vectorstore, document_content_description, metadata_field_info, verbose=True
)

from langchain.chat_models import ChatOpenAI
from langchain.schema.runnable import RunnableParallel, RunnablePassthrough
from langchain.prompts import ChatPromptTemplate, PromptTemplate
from langchain.schema import format_document

LLM_CONTEXT_PROMPT = ChatPromptTemplate.from_template(
    """
Use the following context movies that matched the user question. Use the movies below only to answer the user's question.

If you don't know the answer, just say that you don't know, don't try to make up an answer.

----
{context}
----
Question: {question}
Answer:
"""
)

DOCUMENT_PROMPT = PromptTemplate.from_template(
    """
---
title: {title}                                                                                   
year: {year}  
director: {director}    
---
"""
)


def _combine_documents(
    docs, document_prompt=DOCUMENT_PROMPT, document_separator="\n\n"
):
    doc_strings = [format_document(doc, document_prompt) for doc in docs]
    return document_separator.join(doc_strings)


_context = RunnableParallel(
    context=retriever | _combine_documents,
    question=RunnablePassthrough(),
)

chain = _context | LLM_CONTEXT_PROMPT | llm

chain.invoke(
    "What movies are about dreams and was released after the year 2007 but before 2012?"
)