Improvements to the index creation script

scripts/custom_models/create_index.py

+# Usage
+# 1. Install requirements:
+# pip install requests langchain langchain_text_splitter
+# 2. Run the script:
+# GLAB_TOKEN=<api_token> python3 scripts/custom_models/create_index.py --version_tag="v17.0.0"
+
 import argparse
 import glob
 import os
 import datetime
 import re
 import sqlite3
+import sys
 import requests
 import json
 from zipfile import ZipFile
 from langchain.docstore.document import Document
 from langchain_text_splitters import MarkdownHeaderTextSplitter
 import tempfile
+import logging
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 
 # Function to parse command-line arguments
 def parse_arguments():
     parser = argparse.ArgumentParser(description="Generate and upload GitLab docs index.")
-    parser.add_argument("--version_tag", help="GitLab version tag to include in the URL (e.g., v17.1.0-ee)")
-    parser.add_argument("upload_url", help="URL to upload the database")
-    parser.add_argument("private_token", help="GitLab personal access token")
+    parser.add_argument("--project_id", help="GitLab project ID", default=278964)
+    parser.add_argument("--version_tag", help="GitLab version tag to include in the URL (e.g., v17.1.0-ee)",
+                        default='master')
+    parser.add_argument("--base_url", help="URL to gitlab instance", default="https://gitlab.com")
     return parser.parse_args()
 
+
+def execution_error(error_message):
+    logger.error(error_message)
+    sys.exit(1)
+
+
 # Function to fetch documents from GitLab
-def fetch_documents(version_tag=None):
-    if version_tag:
-        docs_url = f"https://gitlab.com/gitlab-org/gitlab/-/archive/{version_tag}/gitlab-{version_tag}.zip?path=doc"
-    else:
-        print("No version tag provided. Defaulting to fetching from master.")
-        docs_url = f"https://gitlab.com/gitlab-org/gitlab/-/archive/master/gitlab-master.zip?path=doc"
+def fetch_documents(version_tag):
+    docs_url = f"https://gitlab.com/gitlab-org/gitlab/-/archive/{version_tag}/gitlab-{version_tag}.zip?path=doc"
 
     response = requests.get(docs_url)
 
@@ -40,23 +55,44 @@ def fetch_documents(version_tag=None):
         # Find the directory that was extracted
         extracted_dirs = [os.path.join(tmpdirname, name) for name in os.listdir(tmpdirname) if os.path.isdir(os.path.join(tmpdirname, name))]
         if not extracted_dirs:
-            print("No directory found after extraction. Exiting.")
-            return None
+            execution_error("No directory found after extraction. Exiting.")
 
-        print("Documents are fetched.")
+        logger.info("Documents are fetched.")
         extracted_dir = extracted_dirs[0]
-        print(f"Extracted documents to {extracted_dir}")
+        logger.info(f"Extracted documents to {extracted_dir}")
         return extracted_dir
     else:
-        print(f"Failed to download documents. Status code: {response.status_code}")
-        return None
+        execution_error(f"Failed to download documents. Status code: {response.status_code}")
+
+
+def upload_url(base_url, project_id, version_tag):
+    return f"{base_url}/api/v4/projects/{project_id}/packages/generic/gitlab-duo-local-documentation-index/{version_tag}/docs.db"
+
+
+def build_row_corpus(row):
+    corpus = row['content']
+    # Remove the preamble
+    preamble_start = corpus.find('---')
+    if preamble_start != -1:
+        preamble_end = corpus.find('---', preamble_start + 1)
+        corpus = corpus[preamble_end + 2:-1]
+    if not corpus:
+        return ''
+    # Attach the titles to the corpus, these can still be useful
+    corpus = ''.join(row['metadata'].get(f"Header{i}", '') for i in range(1, 6)) + ' ' + corpus
+    # Stemming could be helpful, but it is already applied by the sqlite
+    # Remove punctuation and set to lowercase, this should reduce the size of the corpus and allow
+    # the query to be a bit more robust
+    corpus = corpus.lower()
+    corpus = re.sub(r'[^\w\s]', '', corpus)
+    return corpus
+
 
 # Function to process documents and create the database
 def create_database(path, output_path):
     files = glob.glob(os.path.join(path, "doc/**/*.md"), recursive=True)
     if not files:
-        print("No markdown files found. Exiting.")
-        return
+        execution_error("No markdown files found")
 
     documents = []
 
@@ -86,25 +122,6 @@ def create_database(path, output_path):
             metadata = {**chunk.metadata, **d.metadata}
             rows_to_insert.append({"content": chunk.page_content, "metadata": metadata})
 
-    # Process each row to yield better results
-    def build_row_corpus(row):
-        corpus = row['content']
-        # Remove the preamble
-        preamble_start = corpus.find('---')
-        if preamble_start != -1:
-            preamble_end = corpus.find('---', preamble_start + 1)
-            corpus = corpus[preamble_end + 2:-1]
-        if not corpus:
-            return ''
-        # Attach the titles to the corpus, these can still be useful
-        corpus = ''.join(row['metadata'].get(f"Header{i}", '') for i in range(1, 6)) + ' ' + corpus
-        # Stemming could be helpful, but it is already applied by the sqlite
-        # Remove punctuation and set to lowercase, this should reduce the size of the corpus and allow
-        # the query to be a bit more robust
-        corpus = corpus.lower()
-        corpus = re.sub(r'[^\w\s]', '', corpus)
-        return corpus
-
     for r in rows_to_insert:
         r['processed'] = build_row_corpus(r)
     # sql_tuples = [(r['processed'], r['content'], r['metadata']['filename']) for r in rows_to_insert if r['processed']]
@@ -117,7 +134,8 @@ def create_database(path, output_path):
     conn.commit()
     conn.close()
 
-# Function to upload the database file to GitLab model registry
+
+# Function to upload the database file to GitLab package registry
 def upload_to_gitlab(upload_url, file_path, private_token):
     headers = {"Authorization": f"Bearer {private_token}"}
 
@@ -126,31 +144,35 @@ def upload_to_gitlab(upload_url, file_path, private_token):
         response = requests.put(upload_url, headers=headers, files=files)
 
     if response.status_code in {200, 201}:
-        print("Database uploaded successfully.")
+        logger.info("Database uploaded successfully.")
     else:
-        print(f"Upload failed with status code: {response.status_code}, response: {response.content}")
+        logger.error(f"Upload failed with status code: {response.status_code}, response: {response.content}")
+
 
-# Main function
-def main():
+if __name__ == "__main__":
     args = parse_arguments()
 
+    private_token = os.environ['GLAB_TOKEN']
+
+    if not private_token:
+        execution_error("Private token must be set.")
+
     # Fetch documents based on version tag (if provided)
     docs_path = fetch_documents(version_tag=args.version_tag)
     if not docs_path:
-        print("Fetching documents failed. Exiting.")
-        return
+        execution_error("Fetching documents failed")
 
     # Create database
     timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
     output_path = f"{docs_path}/created_index_docs_{timestamp}.db"
     create_database(docs_path, output_path)
-    print(f"Database created at {output_path}")
-
+    logger.info(f"Database created at {output_path}")
     # Upload to GitLab
-    if os.path.exists(output_path):
-        upload_to_gitlab(args.upload_url, output_path, args.private_token)
-    else:
-        print("Database file not found. Upload skipped.")
+    if not os.path.exists(output_path):
+        execution_error("Database file not found.")
 
-if __name__ == "__main__":
-    main()
+    url = upload_url(args.base_url, args.project_id, args.version_tag)
+
+    logger.info(f"Uploading to {url}")
+
+    upload_to_gitlab(url, output_path, private_token)