langchain/templates/elastic-query-generator/elastic_query_generator/chain.py
Erick Friis ebf998acb6
Templates (#12294)
Co-authored-by: Harrison Chase <hw.chase.17@gmail.com>
Co-authored-by: Lance Martin <lance@langchain.dev>
Co-authored-by: Jacob Lee <jacoblee93@gmail.com>
2023-10-25 18:47:42 -07:00

26 lines
767 B
Python

import os
from langchain.chat_models import ChatOpenAI
from langchain.output_parsers.json import SimpleJsonOutputParser
from elasticsearch import Elasticsearch
from pathlib import Path
from .prompts import DSL_PROMPT
from .elastic_index_info import get_indices_infos
es_host = os.environ["ELASTIC_SEARCH_SERVER"]
es_password = os.environ["ELASTIC_PASSWORD"]
db = Elasticsearch(
es_host,
http_auth=('elastic', es_password),
ca_certs=Path(__file__).parents[1] / 'http_ca.crt' # Replace with your actual path
)
_model = ChatOpenAI(temperature=0, model="gpt-4")
chain = {
"input": lambda x: x["input"],
"indices_info": lambda _: get_indices_infos(db),
"top_k": lambda x: x.get("top_k", 5),
} | DSL_PROMPT | _model | SimpleJsonOutputParser()