notebook fmt (#12498)

templates/elastic-query-generator/elastic_query_generator/chain.py

@@ -16,10 +16,7 @@ ELASTIC_PASSWORD = "..."
 CLOUD_ID = "..."
 
 # Create the client instance
-db = Elasticsearch(
-    cloud_id=CLOUD_ID,
-    basic_auth=("elastic", ELASTIC_PASSWORD)
-)
+db = Elasticsearch(cloud_id=CLOUD_ID, basic_auth=("elastic", ELASTIC_PASSWORD))
 
 # Specify indices to include
 # If you want to use on your own indices, you will need to change this.
@@ -29,13 +26,20 @@ INCLUDE_INDICES = ["customers"]
 
 _model = ChatOpenAI(temperature=0, model="gpt-4")
 
-chain = {
-    "input": lambda x: x["input"],
-    # This line only get index info for "customers" index.
-    # If you are running this on your own data, you will want to change.
-    "indices_info": lambda _: get_indices_infos(db, include_indices=INCLUDE_INDICES),
-    "top_k": lambda x: x.get("top_k", 5),
-} | DSL_PROMPT | _model | SimpleJsonOutputParser()
+chain = (
+    {
+        "input": lambda x: x["input"],
+        # This line only get index info for "customers" index.
+        # If you are running this on your own data, you will want to change.
+        "indices_info": lambda _: get_indices_infos(
+            db, include_indices=INCLUDE_INDICES
+        ),
+        "top_k": lambda x: x.get("top_k", 5),
+    }
+    | DSL_PROMPT
+    | _model
+    | SimpleJsonOutputParser()
+)
 
 
 # Nicely typed inputs for playground

templates/elastic-query-generator/elastic_query_generator/elastic_index_info.py

@@ -2,9 +2,7 @@ from typing import List
 
 
 def _list_indices(database, include_indices=None, ignore_indices=None) -> List[str]:
-    all_indices = [
-        index["index"] for index in database.cat.indices(format="json")
-    ]
+    all_indices = [index["index"] for index in database.cat.indices(format="json")]
 
     if include_indices:
         all_indices = [i for i in all_indices if i in include_indices]
@@ -15,15 +13,13 @@ def _list_indices(database, include_indices=None, ignore_indices=None) -> List[s
 
 
 def get_indices_infos(
-        database,
-        sample_documents_in_index_info=5,
-        include_indices=None,
-        ignore_indices=None
+    database,
+    sample_documents_in_index_info=5,
+    include_indices=None,
+    ignore_indices=None,
 ) -> str:
     indices = _list_indices(
-        database,
-        include_indices=include_indices,
-        ignore_indices=ignore_indices
+        database, include_indices=include_indices, ignore_indices=ignore_indices
     )
     mappings = database.indices.get_mapping(index=",".join(indices))
     if sample_documents_in_index_info > 0:
@@ -40,4 +36,4 @@ def get_indices_infos(
             "Mapping for index {}:\n{}".format(index, mappings[index]["mappings"])
             for index in mappings
         ]
-    )
+    )

templates/elastic-query-generator/ingest.py

@@ -10,17 +10,14 @@ ELASTIC_PASSWORD = "..."
 CLOUD_ID = "..."
 
 # Create the client instance
-db = Elasticsearch(
-    cloud_id=CLOUD_ID,
-    basic_auth=("elastic", ELASTIC_PASSWORD)
-)
+db = Elasticsearch(cloud_id=CLOUD_ID, basic_auth=("elastic", ELASTIC_PASSWORD))
 
 customers = [
     {"firstname": "Jennifer", "lastname": "Walters"},
-    {"firstname": "Monica","lastname":"Rambeau"},
-    {"firstname": "Carol","lastname":"Danvers"},
-    {"firstname": "Wanda","lastname":"Maximoff"},
-    {"firstname": "Jennifer","lastname":"Takeda"},
+    {"firstname": "Monica", "lastname": "Rambeau"},
+    {"firstname": "Carol", "lastname": "Danvers"},
+    {"firstname": "Wanda", "lastname": "Maximoff"},
+    {"firstname": "Jennifer", "lastname": "Takeda"},
 ]
 for i, customer in enumerate(customers):
     db.create(index="customers", document=customer, id=i)