Skip to content

Conversation

divya-garak
Copy link

Summary

This PR adds a missing Azure OpenAI embedding provider to fix the KeyError: 'azure' does not exist in the registry error reported in issue #1268.

Changes Made

  • New Azure OpenAI Embedding Provider: Implemented AzureOpenAIEmbeddingModel class in nemoguardrails/embeddings/providers/azure.py

    • Full Azure OpenAI API integration using the openai package
    • Supports both known embedding models and custom Azure deployment names
    • Automatic embedding size detection for custom deployments
    • Proper error handling for missing OpenAI dependency
  • Provider Registration: Added azure provider to the embedding provider registry in __init__.py

  • Comprehensive Test Suite: Created tests/test_azure_embedding_provider.py with 6 test cases:

    • Provider registration verification
    • Model initialization with different parameters
    • Custom deployment name handling
    • Encoding functionality testing
    • Error handling for missing dependencies
    • Integration with provider registry

Technical Details

The Azure provider uses the same embedding models as OpenAI (text-embedding-ada-002, text-embedding-3-small, text-embedding-3-large) but connects to Azure OpenAI endpoints. For custom deployment names, it performs an initial encoding call to determine the embedding size automatically.

Testing

All 6 test cases pass, covering:

  • ✅ Provider registration and retrieval
  • ✅ Initialization with various parameters
  • ✅ Custom deployment support
  • ✅ Encoding functionality
  • ✅ Error handling
  • ✅ Registry integration

Fixes

Closes #1268

Type of Change

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • This change requires a documentation update

- Implement AzureOpenAIEmbeddingModel class with full Azure OpenAI API support
- Add automatic embedding size detection for custom deployment names
- Register azure embedding provider in the provider registry
- Add comprehensive test suite with 6 test cases covering:
  - Provider registration and initialization
  - Custom deployment name handling
  - Encoding functionality and API calls
  - Error handling for missing dependencies
- Fix KeyError: 'azure does not exist in the registry' when using Azure embeddings

Fixes NVIDIA-NeMo#1268
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant