Azure Function to vectorise PDFs and store in Qdrant container app with OpenAI and Python

"""
An Azure function app which:

- is triggered when a new PDF file is added to the blob container 'docupload'
- reads the PDF file and turns to vectors
- stores the vectors in Azure Qdrant

Components:

- Function app 
- Blob store 
- Qdrant container 

Prerequisites:
- Install VS Code Azure extension 
- Read getting started documentation at https://shorturl.at/59jYg
"""

import azure.functions as func
import logging
import fitz
import openai
import qdrant_client.models as models
import tiktoken
from langchain.text_splitter import RecursiveCharacterTextSplitter
from qdrant_client import QdrantClient
from qdrant_client.http.models import *
from qdrant_client.fastembed_common import *


app = func.FunctionApp()

@app.blob_trigger(arg_name="myblob", path="docupload",
                               connection="saaicdupsertvectorspoc01_STORAGE") 
def aicdfaupsertvectorspoc(myblob: func.InputStream):
    # 1. Read document
    blob_name: str = myblob.name

    logging.info(f"Python blob trigger function processed blob"
                f"Name: {myblob.name}"
                f"Blob Size: {myblob.length} bytes")
    ''
    try:
        document = fitz.open(stream=myblob.read(), filetype="pdf")
        logging.info(f'PDF read successfully: {document}')
    except:
        print("The PDF could not be read.")

    # 2. Vectorise document and upload to Qdrant
    def tiktoken_len(text: str) -> int:
        tokenizer = tiktoken.get_encoding("p50k_base")
        tokens = tokenizer.encode(text, disallowed_special=())

        return len(tokens)

    def data_upload(qdrant_index_name: str, document) -> None:
        settings = {
            "url": "https://ca-qdrant-poc.azurecontainerapps.io", # The URL to your container app
            "host": "ca-qdrant-poc.azurecontainerapps.io",
            "port": "6333",
            "openai_api_key": "", # Enter your OpenAI API key
            "openai_embedding_model": "text-embedding-ada-002"
        }

        whole_text = []
        for page in document:
            text = page.get_text()
            text = text.replace("\n", " ")
            text = text.replace("\\xc2\\xa3", "£")
            text = text.replace("\\xe2\\x80\\x93", "-")
            whole_text.append(text)

        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=500,
            chunk_overlap=100,
            length_function=tiktoken_len,
            separators=["\n\n", "\n", " ", ""],
        )

        chunks = []

        for record in whole_text:
            text_temp = text_splitter.split_text(record)
            chunks.extend([{"text": text_temp[i]} for i in range(len(text_temp))])

            try:
                client = QdrantClient(url=settings["url"],
                                      port=None)

                collection_names = []
                collections = client.get_collections()

                for i in range(len(collections.collections)):
                    collection_names.append(collections.collections[i].name)

                if qdrant_index_name in collection_names:
                    client.get_collection(collection_name=qdrant_index_name)
                else:
                    client.create_collection(
                        collection_name=qdrant_index_name,
                        vectors_config=models.VectorParams(
                            distance=models.Distance.COSINE, size=1536
                        ),
                    )
            except Exception as e:
                logging.error("Unable to connect to QdrantClient")
                logging.error(f"Error message: {str(e)}")

        for id, observation in enumerate(chunks):
            text = observation["text"]

            try:
                openai.api_key = settings["openai_api_key"]
                res = openai.Embedding.create(
                    input=text, engine=settings["openai_embedding_model"]
                )
            except openai.AuthenticationError:
                logging.error("Invalid API key")
            except openai.APIConnectionError:
                logging.error(
                    "Issue connecting to open ai service. Check network and configuration settings"
                )
            except openai.RateLimitError:
                logging.error("You have exceeded your predefined rate limits")

            client.upsert(
                collection_name=qdrant_index_name,
                points=[
                    models.PointStruct(
                        id=id,
                        payload={"text": text},
                        vector=res.data[0].embedding,
                    )
                ],
            )
            logging.info("Text uploaded")

        logging.info("Embeddings upserted")

    file_index = blob_name \
        .strip() \
        .lower() \
        .replace(" ", "_") \
        .replace("docupload/", "") \
        .replace(".pdfblob", "") \
        .replace(".pdf", "")
    
    logging.info(f"File index: {file_index}")

    data_upload(qdrant_index_name=file_index, document=document)