Tutorial: Migrate Cortex AI to Azure AI Services¶

# cortex_replacement.py
from openai import AzureOpenAI
import os

client = AzureOpenAI(
    api_key=os.environ["AZURE_OPENAI_API_KEY"],
    api_version="2024-12-01-preview",
    azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"]
)

def cortex_complete_replacement(prompt: str, model: str = "gpt-4o",
                                  max_tokens: int = 1000,
                                  temperature: float = 0.3) -> str:
    """
    Drop-in replacement for Snowflake CORTEX.COMPLETE.

    Cortex: SNOWFLAKE.CORTEX.COMPLETE('model', 'prompt')
    Azure:  cortex_complete_replacement(prompt, model='gpt-4o')
    """
    response = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": prompt}],
        max_tokens=max_tokens,
        temperature=temperature
    )
    return response.choices[0].message.content


def cortex_summarize_replacement(text: str, max_length: int = 500) -> str:
    """
    Drop-in replacement for Snowflake CORTEX.SUMMARIZE.

    Cortex: SNOWFLAKE.CORTEX.SUMMARIZE(text)
    Azure:  cortex_summarize_replacement(text)
    """
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "system",
                "content": "You are a federal analyst. Provide concise, accurate summaries."
            },
            {
                "role": "user",
                "content": f"Summarize the following text in a clear, concise paragraph:\n\n{text}"
            }
        ],
        max_tokens=max_length,
        temperature=0.2
    )
    return response.choices[0].message.content


def cortex_translate_replacement(text: str, source_lang: str = "en",
                                   target_lang: str = "es") -> str:
    """
    Drop-in replacement for Snowflake CORTEX.TRANSLATE.

    Cortex: SNOWFLAKE.CORTEX.TRANSLATE(text, 'en', 'es')
    Azure:  cortex_translate_replacement(text, 'en', 'es')
    """
    lang_map = {
        "en": "English", "es": "Spanish", "fr": "French",
        "de": "German", "ja": "Japanese", "zh": "Chinese",
        "ko": "Korean", "pt": "Portuguese", "ar": "Arabic"
    }
    source = lang_map.get(source_lang, source_lang)
    target = lang_map.get(target_lang, target_lang)

    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "system",
                "content": f"You are a professional translator. Translate from {source} to {target}. "
                           f"Preserve formatting, technical terminology, and tone."
            },
            {"role": "user", "content": text}
        ],
        max_tokens=len(text) * 2,  # Allow room for expansion
        temperature=0.1
    )
    return response.choices[0].message.content


def cortex_sentiment_replacement(text: str) -> float:
    """
    Drop-in replacement for Snowflake CORTEX.SENTIMENT.

    Cortex: SNOWFLAKE.CORTEX.SENTIMENT(text) -> float (-1 to 1)
    Azure:  cortex_sentiment_replacement(text) -> float (-1 to 1)
    """
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "system",
                "content": "Analyze the sentiment of the text. Return ONLY a number between -1.0 (very negative) "
                           "and 1.0 (very positive). 0.0 is neutral. Return only the number, nothing else."
            },
            {"role": "user", "content": text}
        ],
        max_tokens=10,
        temperature=0.0
    )
    try:
        return float(response.choices[0].message.content.strip())
    except ValueError:
        return 0.0

# Register functions as Databricks UDFs for SQL access
from pyspark.sql.functions import udf
from pyspark.sql.types import StringType, FloatType

# Register UDFs
spark.udf.register("ai_complete", cortex_complete_replacement, StringType())
spark.udf.register("ai_summarize", cortex_summarize_replacement, StringType())
spark.udf.register("ai_translate", cortex_translate_replacement, StringType())
spark.udf.register("ai_sentiment", cortex_sentiment_replacement, FloatType())

# create_search_index.py
from azure.search.documents.indexes import SearchIndexClient
from azure.search.documents.indexes.models import (
    SearchIndex, SearchField, SearchFieldDataType,
    VectorSearch, HnswAlgorithmConfiguration,
    VectorSearchProfile, SemanticConfiguration,
    SemanticSearch, SemanticPrioritizedFields, SemanticField
)
from azure.core.credentials import AzureKeyCredential
import os

index_client = SearchIndexClient(
    endpoint=os.environ["AZURE_SEARCH_ENDPOINT"],
    credential=AzureKeyCredential(os.environ["AZURE_SEARCH_KEY"])
)

fields = [
    SearchField(name="id", type=SearchFieldDataType.String, key=True, filterable=True),
    SearchField(name="title", type=SearchFieldDataType.String, searchable=True, filterable=True),
    SearchField(name="content", type=SearchFieldDataType.String, searchable=True),
    SearchField(name="category", type=SearchFieldDataType.String, filterable=True, facetable=True),
    SearchField(name="agency", type=SearchFieldDataType.String, filterable=True, facetable=True),
    SearchField(name="published_date", type=SearchFieldDataType.DateTimeOffset, filterable=True, sortable=True),
    SearchField(
        name="content_vector",
        type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
        searchable=True,
        vector_search_dimensions=3072,
        vector_search_profile_name="vector-profile"
    )
]

vector_search = VectorSearch(
    algorithms=[HnswAlgorithmConfiguration(name="hnsw-config", parameters={"m": 4, "efConstruction": 400})],
    profiles=[VectorSearchProfile(name="vector-profile", algorithm_configuration_name="hnsw-config")]
)

semantic_config = SemanticConfiguration(
    name="semantic-config",
    prioritized_fields=SemanticPrioritizedFields(
        title_field=SemanticField(field_name="title"),
        content_fields=[SemanticField(field_name="content")]
    )
)

index = SearchIndex(
    name="government-documents",
    fields=fields,
    vector_search=vector_search,
    semantic_search=SemanticSearch(configurations=[semantic_config])
)

index_client.create_or_update_index(index)
print(f"Index 'government-documents' created successfully")

# index_documents.py
from azure.search.documents import SearchClient
from azure.core.credentials import AzureKeyCredential
from openai import AzureOpenAI
import os

openai_client = AzureOpenAI(
    api_key=os.environ["AZURE_OPENAI_API_KEY"],
    api_version="2024-12-01-preview",
    azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"]
)

search_client = SearchClient(
    endpoint=os.environ["AZURE_SEARCH_ENDPOINT"],
    index_name="government-documents",
    credential=AzureKeyCredential(os.environ["AZURE_SEARCH_KEY"])
)

def generate_embedding(text: str) -> list:
    """Generate embedding using Azure OpenAI."""
    response = openai_client.embeddings.create(
        model="text-embedding-3-large",
        input=text
    )
    return response.data[0].embedding

# Read documents from Delta Lake
documents_df = spark.table("analytics_prod.raw.government_documents").toPandas()

# Batch index documents
batch_size = 100
for i in range(0, len(documents_df), batch_size):
    batch = documents_df.iloc[i:i+batch_size]

    docs_to_upload = []
    for _, row in batch.iterrows():
        embedding = generate_embedding(row["content"][:8000])  # Truncate for token limit

        docs_to_upload.append({
            "id": str(row["document_id"]),
            "title": row["title"],
            "content": row["content"],
            "category": row["category"],
            "agency": row["agency"],
            "published_date": row["published_date"].isoformat(),
            "content_vector": embedding
        })

    result = search_client.upload_documents(documents=docs_to_upload)
    print(f"Indexed batch {i//batch_size + 1}: {len(docs_to_upload)} documents")

print(f"Total documents indexed: {len(documents_df)}")

# rag_pipeline.py
from azure.search.documents import SearchClient
from azure.search.documents.models import VectorizedQuery
from azure.core.credentials import AzureKeyCredential
from openai import AzureOpenAI
import os

openai_client = AzureOpenAI(
    api_key=os.environ["AZURE_OPENAI_API_KEY"],
    api_version="2024-12-01-preview",
    azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"]
)

search_client = SearchClient(
    endpoint=os.environ["AZURE_SEARCH_ENDPOINT"],
    index_name="government-documents",
    credential=AzureKeyCredential(os.environ["AZURE_SEARCH_KEY"])
)

def ask_question(question: str, top_k: int = 5,
                 category_filter: str = None) -> dict:
    """
    RAG pipeline: replace Cortex Search + EXTRACT_ANSWER.

    1. Embed the question
    2. Search Azure AI Search (hybrid: vector + keyword + semantic)
    3. Generate answer with Azure OpenAI using retrieved context
    """
    # Step 1: Generate question embedding
    embedding = openai_client.embeddings.create(
        model="text-embedding-3-large",
        input=question
    ).data[0].embedding

    # Step 2: Hybrid search
    filter_expr = f"category eq '{category_filter}'" if category_filter else None

    results = search_client.search(
        search_text=question,
        vector_queries=[
            VectorizedQuery(
                vector=embedding,
                k_nearest_neighbors=top_k,
                fields="content_vector"
            )
        ],
        filter=filter_expr,
        query_type="semantic",
        semantic_configuration_name="semantic-config",
        top=top_k
    )

    # Step 3: Build context from search results
    sources = []
    context_parts = []
    for result in results:
        sources.append({
            "id": result["id"],
            "title": result["title"],
            "score": result["@search.score"]
        })
        context_parts.append(
            f"[Source: {result['title']}]\n{result['content'][:2000]}"
        )

    context = "\n\n---\n\n".join(context_parts)

    # Step 4: Generate answer
    response = openai_client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "system",
                "content": (
                    "You are a federal analyst answering questions based on "
                    "government documents. Use ONLY the provided context to "
                    "answer. Cite your sources by title. If the context does "
                    "not contain the answer, say 'Information not found in "
                    "available documents.'"
                )
            },
            {
                "role": "user",
                "content": f"Context:\n{context}\n\nQuestion: {question}"
            }
        ],
        max_tokens=1000,
        temperature=0.1
    )

    return {
        "answer": response.choices[0].message.content,
        "sources": sources,
        "tokens_used": response.usage.total_tokens
    }


# Usage
result = ask_question("What is the budget allocation for cybersecurity in FY2026?")
print(f"Answer: {result['answer']}")
print(f"Sources: {result['sources']}")

# content_safety.py
from azure.ai.contentsafety import ContentSafetyClient
from azure.ai.contentsafety.models import AnalyzeTextOptions
from azure.core.credentials import AzureKeyCredential
import os

safety_client = ContentSafetyClient(
    endpoint=os.environ["AZURE_CONTENT_SAFETY_ENDPOINT"],
    credential=AzureKeyCredential(os.environ["AZURE_CONTENT_SAFETY_KEY"])
)

def check_safety(text: str, max_severity: int = 2) -> dict:
    """
    Check text for content safety issues.
    Returns: {"safe": bool, "categories": {...}, "blocked_reason": str|None}
    """
    request = AnalyzeTextOptions(text=text[:10000])  # API limit
    response = safety_client.analyze_text(request)

    categories = {}
    blocked = False
    blocked_reason = None

    for category in response.categories_analysis:
        categories[category.category] = category.severity
        if category.severity > max_severity:
            blocked = True
            blocked_reason = f"{category.category} severity: {category.severity}"

    return {
        "safe": not blocked,
        "categories": categories,
        "blocked_reason": blocked_reason
    }


def safe_ai_call(prompt: str, **kwargs) -> str:
    """Wrapper that adds content safety guardrails to any AI call."""
    # Check input safety
    input_check = check_safety(prompt)
    if not input_check["safe"]:
        return f"[BLOCKED] Input blocked: {input_check['blocked_reason']}"

    # Make AI call
    result = cortex_complete_replacement(prompt, **kwargs)

    # Check output safety
    output_check = check_safety(result)
    if not output_check["safe"]:
        return f"[BLOCKED] Output blocked: {output_check['blocked_reason']}"

    return result

# validate_quality.py
import pandas as pd

# Prepare test cases from actual Cortex usage
test_cases = spark.sql("""
    SELECT DISTINCT
        document_id,
        document_text,
        -- Capture existing Cortex outputs for comparison
        existing_summary,
        existing_sentiment
    FROM analytics_prod.validation.cortex_test_cases
    LIMIT 100
""").toPandas()

results = []
for _, row in test_cases.iterrows():
    # Azure OpenAI summary
    azure_summary = cortex_summarize_replacement(row["document_text"])

    # Azure OpenAI sentiment
    azure_sentiment = cortex_sentiment_replacement(row["document_text"])

    results.append({
        "document_id": row["document_id"],
        "cortex_summary": row["existing_summary"],
        "azure_summary": azure_summary,
        "cortex_sentiment": row["existing_sentiment"],
        "azure_sentiment": azure_sentiment,
        "sentiment_diff": abs(float(row["existing_sentiment"]) - azure_sentiment)
    })

results_df = pd.DataFrame(results)

# Summary statistics
print(f"Average sentiment difference: {results_df['sentiment_diff'].mean():.3f}")
print(f"Max sentiment difference: {results_df['sentiment_diff'].max():.3f}")
print(f"Sentiment correlation: {results_df['cortex_sentiment'].astype(float).corr(results_df['azure_sentiment']):.3f}")