# app/main.py """(c) 2025, Elastic Co. Author: Adhish Thite Script Flow: -------------- 1. Set up the project root and update the Python path. 2. Import necessary modules and configure logging. 3. Initialize the output directory. 4. Query KB and News articles from BigQuery. 5. Obtain an Elasticsearch client. 6. Create the raw data index and store KB and News articles. 7. Delete and recreate the vector index for embedding storage. 8. Process KB documents in batches: - Embed each document batch. - Bulk insert embedded documents into Elasticsearch. - Collect any errors and enforce a delay to avoid rate limits. 9. Process News documents similarly in batches. 10. Log the total processed chunks and any errors. 11. Conclude the application run. """ import sys from pathlib import Path # Add the project root directory to the Python path to allow local imports project_root = str(Path(__file__).parent.parent) sys.path.insert(0, project_root) import time from typing import List, Dict, Any from tqdm import tqdm # Progress bar for monitoring batch processing from elasticsearch.helpers import bulk # Bulk helper for Elasticsearch insertion # Import project-specific configurations and utility functions from app.config.logging_config import setup_logger from app.config.settings import ( OUTPUT_DIR, ES_INDEX_NAME, ES_VECTOR_INDEX_NAME, BATCH_SIZE, ) from app.utils.helpers import init, batch_documents from app.services.bigquery import query_bigquery, query_news_articles from app.services.elasticsearch import ( get_elasticsearch_client, create_elastic_index, insert_dataframe_to_elasticsearch, create_vector_index, ) from app.services.embeddings import process_batch, EMBEDDING_MODEL # Set up the logger using the project's logging configuration logger = setup_logger(__name__) if __name__ == "__main__": # Log the start of the application logger.info("🚀 Starting main application") # Initialize the output directory (e.g., create necessary folders) init(OUTPUT_DIR) # ------------------------------- # Querying Data from BigQuery # ------------------------------- # Process Knowledge Base (KB) Articles logger.info("📚 Processing KB Articles...") kb_results = query_bigquery() # Process News Articles logger.info("📰 Processing News Articles...") news_results = query_news_articles() # ------------------------------- # Elasticsearch Setup # ------------------------------- # Get the Elasticsearch client instance es_client = get_elasticsearch_client() # Create a regular index to store raw data and insert queried documents create_elastic_index(es_client, ES_INDEX_NAME) insert_dataframe_to_elasticsearch(es_client, ES_INDEX_NAME, kb_results, doc_type="kb") insert_dataframe_to_elasticsearch(es_client, ES_INDEX_NAME, news_results, doc_type="news") logger.info("🔄 Processing and embedding documents...") # ------------------------------- # Vector Index Management # ------------------------------- # Delete the existing vector index if it exists to ensure a clean slate if es_client.indices.exists(index=ES_VECTOR_INDEX_NAME): es_client.indices.delete(index=ES_VECTOR_INDEX_NAME) logger.info(f"🗑️ Deleted existing vector index: {ES_VECTOR_INDEX_NAME}") # Create a new vector index to store embedded documents create_vector_index(es_client) logger.info(f"✨ Created new vector index: {ES_VECTOR_INDEX_NAME}") # Initialize error collection and chunk counter all_error_chunks: List[Dict[str, Any]] = [] total_chunks = 0 # ------------------------------- # Processing KB Documents # ------------------------------- # Convert the KB articles from DataFrame to list of dictionary records kb_documents = kb_results.to_dict("records") total_kb_batches = (len(kb_documents) + BATCH_SIZE - 1) // BATCH_SIZE logger.info(f"📚 Processing {len(kb_documents)} KB articles...") # Process documents in batches using a progress bar for visibility for batch in tqdm( list(batch_documents(kb_documents, BATCH_SIZE)), total=total_kb_batches, unit="batch" ): # Process each batch: generate embeddings and capture errors batch_embedded_docs, error_chunks, chunks_count = process_batch( batch, EMBEDDING_MODEL, es_client, source_type="kb" ) total_chunks += chunks_count # If there are embedded documents, bulk insert them into Elasticsearch if batch_embedded_docs: success, failed = bulk( es_client, batch_embedded_docs, chunk_size=500, request_timeout=60, raise_on_error=False, ) logger.info(f"📦 KB Bulk insert: {success} succeeded, {len(failed)} failed") # Collect any error chunks from processing if error_chunks: all_error_chunks.extend(error_chunks) # Pause for 5 seconds to avoid rate limiting and potential overload time.sleep(5) # ------------------------------- # Processing News Documents # ------------------------------- # Convert the News articles from DataFrame to list of dictionary records news_documents = news_results.to_dict("records") total_news_batches = (len(news_documents) + BATCH_SIZE - 1) // BATCH_SIZE logger.info(f"📰 Processing {len(news_documents)} News articles...") # Process News documents in batches for batch in tqdm( list(batch_documents(news_documents, BATCH_SIZE)), total=total_news_batches, unit="batch" ): # Process each batch: generate embeddings and capture errors batch_embedded_docs, error_chunks, chunks_count = process_batch( batch, EMBEDDING_MODEL, es_client, source_type="news" ) total_chunks += chunks_count # Bulk insert embedded documents if available if batch_embedded_docs: success, failed = bulk( es_client, batch_embedded_docs, chunk_size=500, request_timeout=60, raise_on_error=False, ) logger.info(f"📦 News Bulk insert: {success} succeeded, {len(failed)} failed") # Collect any error chunks from processing if error_chunks: all_error_chunks.extend(error_chunks) # Pause for 5 seconds to avoid rate limits and overload time.sleep(5) # ------------------------------- # Final Logging and Completion # ------------------------------- # Log the total number of document chunks processed and any errors encountered logger.info(f"📊 Total chunks processed: {total_chunks}") logger.info(f"⚠️ Total errors: {len(all_error_chunks)}") # Indicate that the application has completed successfully logger.info("✅ Application completed successfully")