up by default, but run should work

Signed-off-by: Adrian Cole <[email protected]>
elastic · Feb 20, 2025 · c7327da · c7327da
1 parent 7f82332
commit c7327da
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 38 deletions.
diff --git a/example-apps/chatbot-rag-app/README.md b/example-apps/chatbot-rag-app/README.md
@@ -59,25 +59,13 @@ working Python environment.
 
 **Double-check you have a `.env` file with all your variables set first!**
 
-#### Create your Elasticsearch index
-
-First, ingest the data into elasticsearch:
-```bash
-docker compose run -T --rm --pull always create-index
-```
-
-*Note*: This may take several minutes to complete
-
-#### Run the application
-
-Now, run the app, which listens on http://localhost:4000
 ```bash
-docker compose run --rm --pull always api-frontend
+docker compose up --pull always --force-recreate
 ```
 
-#### Cleanup when finished
+*Note*: The first run may take several minutes to become available.
 
-When you are done, clean up the services above like this:
+Clean up when finished, like this:
 
 ```bash
 docker compose down
@@ -203,9 +191,8 @@ See [Langchain documentation][loader-docs] for more ways to load documents.
 To build the app from source instead of using published images, pass the
 `--build` flag to Docker Compose instead of `--pull always`
 
-For example, to run the create-index service from source:
 ```bash
-docker compose run --rm --build create-index
+docker compose up --build --force-recreate
 ```
 
 ---

diff --git a/example-apps/chatbot-rag-app/data/index_data.py b/example-apps/chatbot-rag-app/data/index_data.py
@@ -49,6 +49,7 @@ def install_elser():
         es.ml.put_trained_model(
             model_id=ELSER_MODEL, input={"field_names": ["text_field"]}
         )
+
     while True:
         status = es.ml.get_trained_models(
             model_id=ELSER_MODEL, include="definition_status"
@@ -57,19 +58,31 @@ def install_elser():
             break
         time.sleep(1)
 
-    # Step 1: Ensure ELSER_MODEL is deployed
-    try:
-        es.ml.start_trained_model_deployment(
-            model_id=ELSER_MODEL, wait_for="fully_allocated"
-        )
-        print(f'"{ELSER_MODEL}" model is deployed')
-    except BadRequestError:
-        # This error means it already exists
-        pass
+    # Step 1: Ensure ELSER_MODEL is fully allocated
+    if not is_elser_fully_allocated():
+        try:
+            es.ml.start_trained_model_deployment(
+                model_id=ELSER_MODEL, wait_for="fully_allocated"
+            )
+            print(f'"{ELSER_MODEL}" model is deployed')
+        except BadRequestError:
+            pass
+
+    while True:
+        if is_elser_fully_allocated():
+            break
+        time.sleep(1)
 
     print(f'"{ELSER_MODEL}" model is ready')
 
 
+def is_elser_fully_allocated():
+    stats = es.ml.get_trained_models_stats(model_id=ELSER_MODEL)
+    deployment_stats = stats["trained_model_stats"][0].get("deployment_stats", {})
+    allocation_status = deployment_stats.get("allocation_status", {})
+    return allocation_status.get("state") == "fully_allocated"
+
+
 def main():
     install_elser()
 
@@ -96,7 +109,7 @@ def main():
 
     print(f"Split {len(workplace_docs)} documents into {len(docs)} chunks")
 
-    print(f"Creating Elasticsearch sparse vector store in {ELASTICSEARCH_URL}")
+    print(f"Creating Elasticsearch sparse vector store for {ELASTICSEARCH_URL}")
 
     store = ElasticsearchStore(
         es_connection=es,
@@ -110,6 +123,13 @@ def main():
     #
     # Once elastic/elasticsearch#107077 is fixed, we can use bulk_kwargs to
     # adjust the timeout.
+
+    print(f"Adding documents to index {INDEX}")
+
+    spinner = Halo(text="Processing bulk operation", spinner="dots")
+    if stdout.isatty():
+        spinner.start()
+
     try:
         es.indices.delete(index=INDEX, ignore_unavailable=True)
         store.add_documents(list(docs))
@@ -124,6 +144,11 @@ def main():
         es.indices.delete(index=INDEX, ignore_unavailable=True)
         store.add_documents(list(docs))
 
+    if stdout.isatty():
+        spinner.stop()
+
+    print(f"Documents added to index {INDEX}")
+
 
 def await_ml_tasks(max_timeout=600, interval=5):
     """
@@ -141,23 +166,14 @@ def await_ml_tasks(max_timeout=600, interval=5):
     if not ml_tasks:
         return  # likely a lost race on tasks
 
-    spinner = Halo(text="Awaiting ML tasks", spinner="dots")
-    if stdout.isatty():
-        spinner.start()
-    else:
-        print(f"Awaiting {len(ml_tasks)} ML tasks")
+    print(f"Awaiting {len(ml_tasks)} ML tasks")
 
     while time.time() - start_time < max_timeout:
         ml_tasks = get_ml_tasks()
         if len(ml_tasks) == 0:
             break
         time.sleep(interval)
 
-    if stdout.isatty():
-        spinner.stop()
-    else:
-        print(f"ML tasks complete")
-
     if ml_tasks:
         raise TimeoutError(
             f"Timeout reached. ML tasks are still running: {', '.join(ml_tasks)}"
@@ -176,4 +192,4 @@ def get_ml_tasks():
 
 # Unless we run through flask, we can miss critical settings or telemetry signals.
 if __name__ == "__main__":
-    main()
+    raise RuntimeError("Run via the parent directory: 'flask create-index'")
diff --git a/example-apps/chatbot-rag-app/docker-compose.yml b/example-apps/chatbot-rag-app/docker-compose.yml
@@ -10,6 +10,7 @@ services:
     env_file:
       - .env
     # Add settings that allow `docker compose run` to use tty and accept Ctrl+C
+    tty: true
     stdin_open: true
     command: flask create-index
     volumes:
@@ -29,6 +30,7 @@ services:
     env_file:
       - .env
     # Add settings that allow `docker compose run` to use tty and accept Ctrl+C
+    tty: true
     stdin_open: true
     volumes:
       # VertexAI uses a file for GOOGLE_APPLICATION_CREDENTIALS, not an API key