Azure Cognitive Search as a vector database for OpenAI embeddings

embeddings_url = "https://cdn.openai.com/API/examples/data/vector_database_wikipedia_articles_embedded.zip"

# The file is ~700 MB so this will take some time
wget.download(embeddings_url)

article_df = pd.read_csv('../../data/vector_database_wikipedia_articles_embedded.csv')  
  
# Read vectors from strings back into a list using json.loads  
article_df["title_vector"] = article_df.title_vector.apply(json.loads)  
article_df["content_vector"] = article_df.content_vector.apply(json.loads)  
article_df['vector_id'] = article_df['vector_id'].apply(str)  
article_df.head()  

# Configure a search index
index_client = SearchIndexClient(
    endpoint=search_service_endpoint, credential=credential)
fields = [
    SimpleField(name="id", type=SearchFieldDataType.String),
    SimpleField(name="vector_id", type=SearchFieldDataType.String, key=True),
    SimpleField(name="url", type=SearchFieldDataType.String),
    SearchableField(name="title", type=SearchFieldDataType.String),
    SearchableField(name="text", type=SearchFieldDataType.String),
    SearchField(name="title_vector", type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
                searchable=True, vector_search_dimensions=1536, vector_search_configuration="my-vector-config"),
    SearchField(name="content_vector", type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
                searchable=True, vector_search_dimensions=1536, vector_search_configuration="my-vector-config"),
]

# Configure the vector search configuration
vector_search = VectorSearch(
    algorithm_configurations=[
        HnswVectorSearchAlgorithmConfiguration(
            name="my-vector-config",
            kind="hnsw",
            parameters={
                "m": 4,
                "efConstruction": 400,
                "efSearch": 500,
                "metric": "cosine"
            }
        )
    ]
)

# Optional: configure semantic reranking by passing your title, keywords, and content fields
semantic_config = SemanticConfiguration(
    name="my-semantic-config",
    prioritized_fields=PrioritizedFields(
        title_field=SemanticField(field_name="title"),
        prioritized_keywords_fields=[SemanticField(field_name="url")],
        prioritized_content_fields=[SemanticField(field_name="text")]
    )
)
# Create the semantic settings with the configuration
semantic_settings = SemanticSettings(configurations=[semantic_config])

# Create the index 
index = SearchIndex(name=index_name, fields=fields,
                    vector_search=vector_search, semantic_settings=semantic_settings)
result = index_client.create_or_update_index(index)
print(f'{result.name} created')

# Convert the 'id' and 'vector_id' columns to string so one of them can serve as our key field  
article_df['id'] = article_df['id'].astype(str)  
article_df['vector_id'] = article_df['vector_id'].astype(str)  
  
# Convert the DataFrame to a list of dictionaries  
documents = article_df.to_dict(orient='records')  
  
# Use SearchIndexingBufferedSender to upload the documents in batches optimized for indexing 
with SearchIndexingBufferedSender(search_service_endpoint, index_name, AzureKeyCredential(search_service_api_key)) as batch_client:  
    # Add upload actions for all documents  
    batch_client.upload_documents(documents=documents)  
  
print(f"Uploaded {len(documents)} documents in total")  

# Example function to generate document embedding  
def generate_document_embeddings(text):  
    response = openai.Embedding.create(  
        input=text, engine=model)  
    embeddings = response['data'][0]['embedding']  
    return embeddings  
  
# Sampling the first document content as an example 
first_document_content = documents[0]['text']  
print(f"Content: {first_document_content[:100]}")    
    
# Generate the content vector using the `generate_document_embeddings` function    
content_vector = generate_document_embeddings(first_document_content)    
print(f"Content vector generated")    

# Function to generate query embedding
def generate_embeddings(text):
    response = openai.Embedding.create(
        input=text, engine=model)
    embeddings = response['data'][0]['embedding']
    return embeddings

# Pure Vector Search
query = "modern art in Europe"
  
search_client = SearchClient(search_service_endpoint, index_name, AzureKeyCredential(search_service_api_key))  
vector = Vector(value=generate_embeddings(query), k=3, fields="content_vector")  
  
results = search_client.search(  
    search_text=None,  
    vectors=[vector],  
    select=["title", "text", "url"] 
)
  
for result in results:  
    print(f"Title: {result['title']}")  
    print(f"Score: {result['@search.score']}")  
    print(f"URL: {result['url']}\n")  

# Hybrid Search
query = "Famous battles in Scottish history"  
  
search_client = SearchClient(search_service_endpoint, index_name, AzureKeyCredential(search_service_api_key))  
vector = Vector(value=generate_embeddings(query), k=3, fields="content_vector")  
  
results = search_client.search(  
    search_text=query,  
    vectors=[vector],
    select=["title", "text", "url"],
    top=3
)  
  
for result in results:  
    print(f"Title: {result['title']}")  
    print(f"Score: {result['@search.score']}")  
    print(f"URL: {result['url']}\n")  

# Semantic Hybrid Search
query = "Famous battles in Scottish history" 

search_client = SearchClient(search_service_endpoint, index_name, AzureKeyCredential(search_service_api_key))  
vector = Vector(value=generate_embeddings(query), k=3, fields="content_vector")  

results = search_client.search(  
    search_text=query,  
    vectors=[vector], 
    select=["title", "text", "url"],
    query_type="semantic", query_language="en-us", semantic_configuration_name='my-semantic-config', query_caption="extractive", query_answer="extractive",
    top=3
)

semantic_answers = results.get_answers()
for answer in semantic_answers:
    if answer.highlights:
        print(f"Semantic Answer: {answer.highlights}")
    else:
        print(f"Semantic Answer: {answer.text}")
    print(f"Semantic Answer Score: {answer.score}\n")

for result in results:
    print(f"Title: {result['title']}")
    print(f"URL: {result['url']}")
    captions = result["@search.captions"]
    if captions:
        caption = captions[0]
        if caption.highlights:
            print(f"Caption: {caption.highlights}\n")
        else:
            print(f"Caption: {caption.text}\n")