Skip to content

Fix indexing notebook #161

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
MichalSteczko wants to merge 1 commit into
base: main
Choose a base branch
from
+148 −15
Open

Fix indexing notebook #161

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
163 changes: 148 additions & 15 deletions ai4rag/assets_generator/notebook_templates/maas_indexing_template.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,19 +10,55 @@
{
"cell_type": "markdown",
"metadata": {},
"source": "## Pattern {PATTERN_NAME} Index Building Content\n\nThis notebook demonstrates how to process documents and build a vector store index for RAG applications. It covers document discovery, text extraction, chunking, and uploading embeddings to a vector database using OpenShift AI Models-as-a-Service (MaaS).\n\n### 📋 Contents \nThis notebook contains the following sections:\n\n- **[Setup](#Setup)**\n - [Install packages](#Install-packages)\n - [Import required libraries](#Import-required-libraries)\n - [Configure S3 credentials](#Configure-S3-credentials)\n - [Prepare S3 client](#Prepare-S3-client)\n- **[Process input documents](#Process-input-documents)**\n - [Documents discovery](#Documents-discovery)\n - [Text extraction](#Text-extraction)\n- **[Upload documents content into vector store](#Upload-documents-content-into-vector-store)**\n - [Prepare MaaS client](#Prepare-MaaS-client)\n - [Prepare chunker](#Prepare-chunker)\n - [Initialize vector store](#Initialize-vector-store)\n - [Upload chunks to vector store](#Upload-chunks-to-vector-store)\n - [Retrieve chunks for sample question](#Retrieve-chunks-for-sample-question)\n- **[Summary](#Summary)**"
"source": [
"## Pattern {PATTERN_NAME} Index Building Content\n",
"\n",
"This notebook demonstrates how to process documents and build a vector store index for RAG applications. It covers document discovery, text extraction, chunking, and uploading embeddings to a vector database using OpenShift AI Models-as-a-Service (MaaS).\n",
"\n",
"### 📋 Contents \n",
"This notebook contains the following sections:\n",
"\n",
"- **[Setup](#Setup)**\n",
" - [Install packages](#Install-packages)\n",
" - [Import required libraries](#Import-required-libraries)\n",
" - [Configure S3 credentials](#Configure-S3-credentials)\n",
" - [Prepare S3 client](#Prepare-S3-client)\n",
"- **[Process input documents](#Process-input-documents)**\n",
" - [Documents discovery](#Documents-discovery)\n",
" - [Text extraction](#Text-extraction)\n",
"- **[Upload documents content into vector store](#Upload-documents-content-into-vector-store)**\n",
" - [Prepare MaaS client](#Prepare-MaaS-client)\n",
" - [Prepare chunker](#Prepare-chunker)\n",
" - [Initialize vector store](#Initialize-vector-store)\n",
" - [Upload chunks to vector store](#Upload-chunks-to-vector-store)\n",
" - [Retrieve chunks for sample question](#Retrieve-chunks-for-sample-question)\n",
"- **[Summary](#Summary)**"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": "---\n\n## Setup\n\nThis section sets up the notebook environment by installing required packages, importing libraries, and configuring access to S3 storage.\n\n### Install packages\n\nInstall all required Python packages for document processing and RAG operations:\n- **ai4rag**: The AutoRAG framework for building RAG applications"
"source": [
"---\n",
"\n",
"## Setup\n",
"\n",
"This section sets up the notebook environment by installing required packages, importing libraries, and configuring access to S3 storage.\n",
"\n",
"### Install packages\n",
"\n",
"Install all required Python packages for document processing and RAG operations:\n",
"- **ai4rag**: The AutoRAG framework for building RAG applications"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": "%pip install 'ai4rag[text-extraction]~={AI4RAG_VERSION}' | tail -n 1"
"source": [
"%pip install 'ai4rag[text-extraction]~={AI4RAG_VERSION}' | tail -n 1"
]
},
{
"cell_type": "markdown",
Expand Down Expand Up @@ -135,7 +171,23 @@
"execution_count": null,
"metadata": {},
"outputs": [],
"source": "import getpass\nimport json\nimport logging\nimport os\nimport warnings\nfrom pathlib import Path\n\nwarnings.filterwarnings(\"ignore\")\n\nfor logger_name in (\n \"httpx\",\n \"documents-discovery\",\n \"text-extraction\",\n):\n logging.getLogger(logger_name).propagate = False"
"source": [
"import getpass\n",
"import json\n",
"import logging\n",
"import os\n",
"import warnings\n",
"from pathlib import Path\n",
"\n",
"warnings.filterwarnings(\"ignore\")\n",
"\n",
"for logger_name in (\n",
" \"httpx\",\n",
" \"documents-discovery\",\n",
" \"text-extraction\",\n",
"):\n",
" logging.getLogger(logger_name).propagate = False"
]
},
{
"cell_type": "markdown",
Expand Down Expand Up @@ -206,7 +258,14 @@
{
"cell_type": "markdown",
"metadata": {},
"source": "The data processing pipeline prepares documents for the RAG system in multiple steps. Each step produces outputs stored under `step_outputs/`. \n\n| Step | Function | Purpose |\n|------|----------|---------| \n| 1 | **`discover_documents`** | List documents in the bucket, prioritize benchmark-referenced docs, apply a size cap, and write a JSON manifest (no content download). |\n| 2 | **`extract_text`** | Download the listed documents from S3 and extract text to DoclingDocument JSON files using Docling. |"
"source": [
"The data processing pipeline prepares documents for the RAG system in multiple steps. Each step produces outputs stored under `step_outputs/`. \n",
"\n",
"| Step | Function | Purpose |\n",
"|------|----------|---------| \n",
"| 1 | **`discover_documents`** | List documents in the bucket, prioritize benchmark-referenced docs, apply a size cap, and write a JSON manifest (no content download). |\n",
"| 2 | **`extract_text`** | Download the listed documents from S3 and extract text to DoclingDocument JSON files using Docling. |"
]
},
{
"cell_type": "code",
Expand Down Expand Up @@ -237,7 +296,16 @@
"execution_count": null,
"metadata": {},
"outputs": [],
"source": "result = discover_documents(\n bucket_name=input_data_bucket_name,\n prefix=input_data_key,\n s3_client=s3_client,\n)\nresult.save(step_output_dir / \"discovered_documents\")\n\nprint(json.dumps(result.to_dict(), indent=4, ensure_ascii=False))"
"source": [
"result = discover_documents(\n",
" bucket_name=input_data_bucket_name,\n",
" prefix=input_data_key,\n",
" s3_client=s3_client,\n",
")\n",
"result.save(step_output_dir / \"discovered_documents\")\n",
"\n",
"print(json.dumps(result.to_dict(), indent=4, ensure_ascii=False))"
]
},
{
"cell_type": "markdown",
Expand All @@ -254,8 +322,17 @@
"metadata": {},
"outputs": [],
"source": [
"extracted_text_dir = step_output_dir / \"extracted_text\"extraction_result = extract_text( documents=result.to_dict()[\"documents\"], bucket=result.bucket, output_dir=extracted_text_dir, docling_artifacts_path=os.getenv(\"DOCLING_ARTIFACTS_PATH\"),\n",
")print( f\"Extracted {{extraction_result.processed_count}}/{{extraction_result.total_documents}} documents \" f\"({{extraction_result.error_count}} errors)\")"
"extracted_text_dir = step_output_dir / \"extracted_text\"\n",
"extraction_result = extract_text(\n",
" documents=result.to_dict()[\"documents\"],\n",
" bucket=result.bucket, output_dir=extracted_text_dir,\n",
" docling_artifacts_path=os.getenv(\"DOCLING_ARTIFACTS_PATH\"),\n",
")\n",
"\n",
"print(\n",
" f\"Extracted {{extraction_result.processed_count}}/{{extraction_result.total_documents}} documents \"\n",
" f\"({{extraction_result.error_count}} errors)\"\n",
")"
]
},
{
Expand All @@ -272,14 +349,31 @@
{
"cell_type": "markdown",
"metadata": {},
"source": "### Prepare MaaS client\n\nOpenShift AI Models-as-a-Service (MaaS) serves every model from a single OpenAI-compatible endpoint. This section builds one client from the MaaS base URL and API key; the embedding model reuses it.\n\n**Prerequisites:**\n- `MAAS_API_KEY`: Your authentication key for the MaaS API\n- `MAAS_BASE_URL`: The complete OpenAI-compatible endpoint URL, used verbatim (e.g. `https://<host>/v1`)\n\n&#x1F4A1; **Tip**: In OpenShift AI Workbench, you can add these as environment variables or data connections to avoid entering them manually each time."
"source": [
"### Prepare MaaS client\n",
"\n",
"OpenShift AI Models-as-a-Service (MaaS) serves every model from a single OpenAI-compatible endpoint. This section builds one client from the MaaS base URL and API key; the embedding model reuses it.\n",
"\n",
"**Prerequisites:**\n",
"- `MAAS_API_KEY`: Your authentication key for the MaaS API\n",
"- `MAAS_BASE_URL`: The complete OpenAI-compatible endpoint URL, used verbatim (e.g. `https://<host>/v1`)\n",
"\n",
"&#x1F4A1; **Tip**: In OpenShift AI Workbench, you can add these as environment variables or data connections to avoid entering them manually each time."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": "from ai4rag.utils.clients.maas_client import create_maas_client\n\nMAAS_API_KEY = os.getenv(\"MAAS_API_KEY\") or getpass.getpass(\"Please enter 'MAAS_API_KEY': \")\nMAAS_BASE_URL = os.getenv(\"MAAS_BASE_URL\") or getpass.getpass(\"Please enter 'MAAS_BASE_URL': \")\n\nclient = create_maas_client(base_url=MAAS_BASE_URL, api_key=MAAS_API_KEY)"
"source": [
"from ai4rag.utils.clients.maas_client import create_maas_client\n",
"\n",
"MAAS_API_KEY = os.getenv(\"MAAS_API_KEY\") or getpass.getpass(\"Please enter 'MAAS_API_KEY': \")\n",
"MAAS_BASE_URL = os.getenv(\"MAAS_BASE_URL\") or getpass.getpass(\"Please enter 'MAAS_BASE_URL': \")\n",
"\n",
"client = create_maas_client(base_url=MAAS_BASE_URL, api_key=MAAS_API_KEY)"
]
},
{
"cell_type": "markdown",
Expand Down Expand Up @@ -316,14 +410,44 @@
{
"cell_type": "markdown",
"metadata": {},
"source": "### Initialize vector store\n\nThe vector store holds document embeddings and enables semantic search. The backend is selected from the pattern's provider and configured entirely from environment variables, so no connection details or secrets are stored in this notebook.\n\n**Provider:** `{PROVIDER_TYPE}`\n\n&#x1F4CC; **Action**: Set the environment variables below for the selected provider before running the next cell.\n\n{REQUIRED_ENV_VARS}\n\n&#x1F4A1; **Tip**: In OpenShift AI Workbench, add these as environment variables or data connections so you don't have to set them manually each session."
"source": [
"### Initialize vector store\n",
"\n",
"The vector store holds document embeddings and enables semantic search. The backend is selected from the pattern's provider and configured entirely from environment variables, so no connection details or secrets are stored in this notebook.\n",
"\n",
"**Provider:** `{PROVIDER_TYPE}`\n",
"\n",
"&#x1F4CC; **Action**: Set the environment variables below for the selected provider before running the next cell.\n",
"\n",
"{REQUIRED_ENV_VARS}\n",
"\n",
"&#x1F4A1; **Tip**: In OpenShift AI Workbench, add these as environment variables or data connections so you don't have to set them manually each session."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": "from ai4rag.rag.embedding.openai_model import OpenAIEmbeddingModel, OpenAIEmbeddingParams\nfrom ai4rag.rag.vector_store import get_vector_store, get_vector_store_config\n\nembedding_model_id = \"{EMBEDDING_MODEL_ID}\"\nparams = OpenAIEmbeddingParams(**{EMBEDDING_PARAMS})\n\nembedding_model = OpenAIEmbeddingModel(client=client, model_id=embedding_model_id, params=params)\n\nprovider_type = \"{PROVIDER_TYPE}\"\ncollection_name = \"{COLLECTION_NAME}\"\n\nvector_store_config = get_vector_store_config(provider_type)\nvector_store = get_vector_store(\n embedding_model=embedding_model,\n config=vector_store_config,\n collection_name=collection_name,\n)"
"source": [
"from ai4rag.rag.embedding.openai_model import OpenAIEmbeddingModel, OpenAIEmbeddingParams\n",
"from ai4rag.rag.vector_store import get_vector_store, get_vector_store_config\n",
"\n",
"embedding_model_id = \"{EMBEDDING_MODEL_ID}\"\n",
"params = OpenAIEmbeddingParams(**{EMBEDDING_PARAMS})\n",
"\n",
"embedding_model = OpenAIEmbeddingModel(client=client, model_id=embedding_model_id, params=params)\n",
"\n",
"provider_type = \"{PROVIDER_TYPE}\"\n",
"collection_name = \"{COLLECTION_NAME}\"\n",
"\n",
"vector_store_config = get_vector_store_config(provider_type)\n",
"vector_store = get_vector_store(\n",
" embedding_model=embedding_model,\n",
" config=vector_store_config,\n",
" collection_name=collection_name,\n",
")"
]
},
{
"cell_type": "markdown",
Expand Down Expand Up @@ -364,11 +488,20 @@
]
},
{
"metadata": {},
"cell_type": "code",
"outputs": [],
"execution_count": null,
"source": "from dataclasses import asdict\nfrom pprint import pprint\n\nsample_question = input()\n\nresults = vector_store.search(query=sample_question, k=5)\nfor chunk in results:\n pprint(asdict(chunk), indent=4)"
"metadata": {},
"outputs": [],
"source": [
"from dataclasses import asdict\n",
"from pprint import pprint\n",
"\n",
"sample_question = input()\n",
"\n",
"results = vector_store.search(query=sample_question, k=5)\n",
"for chunk in results:\n",
" pprint(asdict(chunk), indent=4)"
]
},
{
"cell_type": "markdown",
Expand Down
Loading