Python: Update to AOAI OYD 2024-02-15-preview API version (#5684)

### Motivation and Context  ### Description * Updated models to support the new preview API version 2024-02-15-preview, with some backward compatibility with previous API shape through the use of aliases * Updated samples, unit tests and integration tests * Added experimental settings management using pydantic-settings -- not strictly required here, but I thought it was useful for testing to be able to set all values through the .env file, and validation comes for free using Pydantic. Looking forward to feedback on this. :) ### Contribution Checklist  - [ ] The code builds clean without any errors or warnings - [ ] The PR follows the [SK Contribution Guidelines](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md) and the [pre-submission formatting script](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md#development-scripts) raises no violations - [ ] All unit tests pass, and I have added new tests where possible - [ ] I didn't break anyone 😄 --------- Co-authored-by: Abby Hartman <[email protected]> Co-authored-by: Evan Mattson <[email protected]>
2026-10-12 02:19:06 +00:00 · 2024-04-03 12:52:25 +00:00
parent 2ddb5efd5a
commit ca2ad19130
11 changed files with 191 additions and 145 deletions
@@ -1,10 +1,33 @@
 OPENAI_API_KEY=""
 OPENAI_ORG_ID=""
+AZURE_OPENAI_SYSTEM_MESSAGE="You are an AI assistant that helps people find information"
+AZURE_OPENAI_API_VERSION="2024-02-15-preview"
 AZURE_OPENAI_DEPLOYMENT_NAME=""
 AZURE_OPENAI_ENDPOINT=""
 AZURE_OPENAI_API_KEY=""
-AZURE_COGNITIVE_SEARCH_ENDPOINT=""
-AZURE_COGNITIVE_SEARCH_ADMIN_KEY=""
+AZURE_OPENAI_TEMPERATURE=0
+AZURE_OPENAI_MAX_TOKENS=1000
+AZURE_OPENAI_TOP_P=1.0
+AZURE_OPENAI_STREAM=true
+AZURE_AISEARCH_URL=""
+AZURE_AISEARCH_SERVICE=""
+AZURE_AISEARCH_API_KEY=""
+AZURE_AISEARCH_INDEX_NAME=""
+AZURE_AISEARCH_EMBEDDING_DEPLOYMENT_NAME=""
+AZURE_AISEARCH_USE_SEMANTIC_SEARCH=false
+AZURE_AISEARCH_SEMANTIC_SEARCH_CONFIG=default
+AZURE_AISEARCH_INDEX_IS_PRECHUNKED=false
+AZURE_AISEARCH_TOP_K=5
+AZURE_AISEARCH_ENABLE_IN_DOMAIN=true
+AZURE_AISEARCH_CONTENT_COLUMNS=content
+AZURE_AISEARCH_FILEPATH_COLUMN=filepath
+AZURE_AISEARCH_TITLE_COLUMN=title
+AZURE_AISEARCH_URL_COLUMN=url
+AZURE_AISEARCH_VECTOR_COLUMNS=contentVector
+AZURE_AISEARCH_QUERY_TYPE=simple
+AZURE_AISEARCH_PERMITTED_GROUPS_COLUMN=
+AZURE_AISEARCH_STRICTNESS=3
+AZURE_AISEARCH_FILTER=""
 MONGODB_ATLAS_CONNECTION_STRING=""
 PINECONE_API_KEY=""
 PINECONE_ENVIRONMENT=""
@@ -1389,12 +1389,12 @@ files = [
 google-auth = ">=2.14.1,<3.0.dev0"
 googleapis-common-protos = ">=1.56.2,<2.0.dev0"
 grpcio = [
-    {version = ">=1.49.1,<2.0dev", optional = true, markers = "python_version >= \"3.11\" and extra == \"grpc\""},
    {version = ">=1.33.2,<2.0dev", optional = true, markers = "python_version < \"3.11\" and extra == \"grpc\""},
+    {version = ">=1.49.1,<2.0dev", optional = true, markers = "python_version >= \"3.11\" and extra == \"grpc\""},
 ]
 grpcio-status = [
-    {version = ">=1.49.1,<2.0.dev0", optional = true, markers = "python_version >= \"3.11\" and extra == \"grpc\""},
    {version = ">=1.33.2,<2.0.dev0", optional = true, markers = "python_version < \"3.11\" and extra == \"grpc\""},
+    {version = ">=1.49.1,<2.0.dev0", optional = true, markers = "python_version >= \"3.11\" and extra == \"grpc\""},
 ]
 proto-plus = ">=1.22.3,<2.0.0dev"
 protobuf = ">=3.19.5,<3.20.0 || >3.20.0,<3.20.1 || >3.20.1,<4.21.0 || >4.21.0,<4.21.1 || >4.21.1,<4.21.2 || >4.21.2,<4.21.3 || >4.21.3,<4.21.4 || >4.21.4,<4.21.5 || >4.21.5,<5.0.0.dev0"
@@ -3619,9 +3619,9 @@ files = [

 [package.dependencies]
 numpy = [
-    {version = ">=1.26.0,<2", markers = "python_version >= \"3.12\""},
    {version = ">=1.22.4,<2", markers = "python_version < \"3.11\""},
    {version = ">=1.23.2,<2", markers = "python_version == \"3.11\""},
+    {version = ">=1.26.0,<2", markers = "python_version >= \"3.12\""},
 ]
 python-dateutil = ">=2.8.2"
 pytz = ">=2020.1"
@@ -5003,8 +5003,8 @@ grpcio = ">=1.41.0"
 grpcio-tools = ">=1.41.0"
 httpx = {version = ">=0.20.0", extras = ["http2"]}
 numpy = [
-    {version = ">=1.26", markers = "python_version >= \"3.12\""},
    {version = ">=1.21", markers = "python_version >= \"3.8\" and python_version < \"3.12\""},
+    {version = ">=1.26", markers = "python_version >= \"3.12\""},
 ]
 portalocker = ">=2.7.0,<3.0.0"
 pydantic = ">=1.10.8"
@@ -1,6 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.

 import asyncio
+import logging

 import semantic_kernel as sk
 import semantic_kernel.connectors.ai.open_ai as sk_oai
@@ -8,9 +9,8 @@ from semantic_kernel.connectors.ai.open_ai.contents.azure_chat_message_content i
 from semantic_kernel.connectors.ai.open_ai.contents.function_call import FunctionCall
 from semantic_kernel.connectors.ai.open_ai.contents.tool_calls import ToolCall
 from semantic_kernel.connectors.ai.open_ai.prompt_execution_settings.azure_chat_prompt_execution_settings import (
-    AzureAISearchDataSources,
+    AzureAISearchDataSource,
    AzureChatPromptExecutionSettings,
-    AzureDataSources,
    ExtraBody,
 )
 from semantic_kernel.contents.chat_history import ChatHistory
@@ -24,9 +24,10 @@ from semantic_kernel.utils.settings import (
 )

 kernel = sk.Kernel()
+logging.basicConfig(level=logging.DEBUG)

 # Load Azure OpenAI Settings
-aoai_settings = azure_openai_settings_from_dot_env_as_dict()
+aoai_settings = azure_openai_settings_from_dot_env_as_dict(include_api_version=True)

 # For example, AI Search index may contain the following document:

@@ -46,15 +47,14 @@ azure_ai_search_settings["fieldsMapping"] = {
 }

 # Create the data source settings
-az_source = AzureAISearchDataSources(**azure_ai_search_settings)
-az_data = AzureDataSources(type="AzureCognitiveSearch", parameters=az_source)
-extra = ExtraBody(dataSources=[az_data])
+
+az_source = AzureAISearchDataSource(parameters=azure_ai_search_settings)
+extra = ExtraBody(data_sources=[az_source])
 req_settings = AzureChatPromptExecutionSettings(service_id="default", extra_body=extra)

-# When using data, set use_extensions=True and use the 2023-12-01-preview API version.
+# When using data, use the 2024-02-15-preview API version.
 chat_service = sk_oai.AzureChatCompletion(
    service_id="chat-gpt",
-    use_extensions=True,
    **aoai_settings,
 )
 kernel.add_service(chat_service)
@@ -121,7 +121,7 @@ async def chat() -> bool:
            chat_history.add_tool_message(full_message.tool_message, {"tool_call_id": "chat_with_your_data"})
        if full_message.role is None:
            full_message.role = ChatRole.ASSISTANT
-        chat_history.add_message(full_message)
+        chat_history.add_assistant_message(full_message.content)
    return True


@@ -1,14 +1,14 @@
 # Copyright (c) Microsoft. All rights reserved.

 import asyncio
+import logging
 import os

 import semantic_kernel as sk
 import semantic_kernel.connectors.ai.open_ai as sk_oai
 from semantic_kernel.connectors.ai.open_ai.prompt_execution_settings.azure_chat_prompt_execution_settings import (
-    AzureAISearchDataSources,
+    AzureAISearchDataSource,
    AzureChatPromptExecutionSettings,
-    AzureDataSources,
    ExtraBody,
 )
 from semantic_kernel.connectors.ai.open_ai.utils import (
@@ -20,21 +20,21 @@ from semantic_kernel.functions.kernel_arguments import KernelArguments
 from semantic_kernel.prompt_template.input_variable import InputVariable
 from semantic_kernel.prompt_template.prompt_template_config import PromptTemplateConfig

+logging.basicConfig(level=logging.DEBUG)
 # NOTE:
 # AzureOpenAI function calling requires the following models: gpt-35-turbo (1106) or gpt-4 (1106-preview)
-# along with the API version: 2023-12-01-preview
+# along with the API version: 2024-02-15-preview
 # https://learn.microsoft.com/en-us/azure/ai-services/openai/how-to/function-calling?tabs=python

 kernel = sk.Kernel()

 # Load Azure OpenAI Settings
-deployment, api_key, endpoint = sk.azure_openai_settings_from_dot_env()
+deployment, api_key, endpoint = sk.azure_openai_settings_from_dot_env(include_deployment=True)

 # Create the data source settings
 azure_ai_search_settings = sk.azure_aisearch_settings_from_dot_env_as_dict()
-az_source = AzureAISearchDataSources(**azure_ai_search_settings)
-az_data = AzureDataSources(type="AzureCognitiveSearch", parameters=az_source)
-extra = ExtraBody(dataSources=[az_data])
+az_source = AzureAISearchDataSource(parameters=azure_ai_search_settings)
+extra = ExtraBody(data_sources=[az_source])
 req_settings = AzureChatPromptExecutionSettings(service_id="chat-gpt", extra_body=extra, tool_choice="auto")

 # For example, AI Search index may contain the following document:
@@ -45,11 +45,10 @@ req_settings = AzureChatPromptExecutionSettings(service_id="chat-gpt", extra_bod

 chat_service = sk_oai.AzureChatCompletion(
    service_id="chat-gpt",
-    deployment_name="gpt-35-turbo-16k",
+    deployment_name=deployment,
    api_key=api_key,
    endpoint=endpoint,
-    api_version="2023-12-01-preview",
-    use_extensions=True,
+    api_version="2024-02-15-preview",
 )
 kernel.add_service(
    chat_service,
@@ -1,18 +1,18 @@
 # Copyright (c) Microsoft. All rights reserved.

 import asyncio
+import logging

 import semantic_kernel as sk
 import semantic_kernel.connectors.ai.open_ai as sk_oai
-from semantic_kernel.connectors.ai.open_ai.contents.azure_streaming_chat_message_content import (
-    AzureStreamingChatMessageContent,
+from semantic_kernel.connectors.ai.open_ai.contents.azure_chat_message_content import (
+    AzureChatMessageContent,
 )
-from semantic_kernel.connectors.ai.open_ai.contents.open_ai_chat_message_content import OpenAIChatMessageContent
+from semantic_kernel.connectors.ai.open_ai.contents.function_call import FunctionCall
 from semantic_kernel.connectors.ai.open_ai.contents.tool_calls import ToolCall
 from semantic_kernel.connectors.ai.open_ai.prompt_execution_settings.azure_chat_prompt_execution_settings import (
-    AzureAISearchDataSources,
+    AzureAISearchDataSource,
    AzureChatPromptExecutionSettings,
-    AzureDataSources,
    ExtraBody,
 )
 from semantic_kernel.contents.chat_history import ChatHistory
@@ -20,11 +20,16 @@ from semantic_kernel.contents.chat_role import ChatRole
 from semantic_kernel.functions.kernel_arguments import KernelArguments
 from semantic_kernel.prompt_template.input_variable import InputVariable
 from semantic_kernel.prompt_template.prompt_template_config import PromptTemplateConfig
+from semantic_kernel.utils.settings import (
+    azure_aisearch_settings_from_dot_env_as_dict,
+    azure_openai_settings_from_dot_env_as_dict,
+)

 kernel = sk.Kernel()
+logging.basicConfig(level=logging.DEBUG)

 # Load Azure OpenAI Settings
-deployment, api_key, endpoint = sk.azure_openai_settings_from_dot_env()
+aoai_settings = azure_openai_settings_from_dot_env_as_dict(include_api_version=True)

 # For example, AI Search index may contain the following document:

@@ -32,7 +37,7 @@ deployment, api_key, endpoint = sk.azure_openai_settings_from_dot_env()
 # Bonded by their love for the natural world and shared curiosity, they uncovered a
 # groundbreaking phenomenon in glaciology that could potentially reshape our understanding of climate change.

-azure_ai_search_settings = sk.azure_aisearch_settings_from_dot_env_as_dict()
+azure_ai_search_settings = azure_aisearch_settings_from_dot_env_as_dict()

 # This example index has fields "title", "chunk", and "vector".
 # Add fields mapping to the settings.
@@ -41,47 +46,41 @@ azure_ai_search_settings["fieldsMapping"] = {
    "contentFields": ["chunk"],
    "vectorFields": ["vector"],
 }
+
 # Add Ada embedding deployment name to the settings and use vector search.
 azure_ai_search_settings["embeddingDependency"] = {
    "type": "DeploymentName",
    "deploymentName": "ada-002",
 }
-azure_ai_search_settings["queryType"] = "vector"
+azure_ai_search_settings["query_type"] = "vector"

 # Create the data source settings
-az_source = AzureAISearchDataSources(**azure_ai_search_settings)
-az_data = AzureDataSources(type="AzureCognitiveSearch", parameters=az_source)
-extra = ExtraBody(dataSources=[az_data])
+az_source = AzureAISearchDataSource(parameters=azure_ai_search_settings)
+extra = ExtraBody(data_sources=[az_source])
 req_settings = AzureChatPromptExecutionSettings(service_id="default", extra_body=extra)

-# When using data, set use_extensions=True and use the 2023-12-01-preview API version.
+# When using data, use the 2024-02-15-preview API version.
 chat_service = sk_oai.AzureChatCompletion(
    service_id="chat-gpt",
-    deployment_name=deployment,
-    api_key=api_key,
-    endpoint=endpoint,
-    api_version="2023-12-01-preview",
-    use_extensions=True,
+    **aoai_settings,
 )
 kernel.add_service(chat_service)

 prompt_template_config = PromptTemplateConfig(
-    template="{{$user_input}}",
+    template="{{$chat_history}}{{$user_input}}",
    name="chat",
    template_format="semantic-kernel",
    input_variables=[
-        InputVariable(name="chat", description="The history of the conversation", is_required=True, default=""),
+        InputVariable(name="chat_history", description="The history of the conversation", is_required=True, default=""),
        InputVariable(name="request", description="The user input", is_required=True),
    ],
    execution_settings={"default": req_settings},
 )

-chat = ChatHistory()
+chat_history = ChatHistory()

-chat.add_user_message("Hi there, who are you?")
-chat.add_assistant_message("I am an AI assistant here to answer your questions.")
-
-arguments = KernelArguments()
+chat_history.add_user_message("Hi there, who are you?")
+chat_history.add_assistant_message("I am an AI assistant here to answer your questions.")

 chat_function = kernel.create_function_from_prompt(
    plugin_name="ChatBot", function_name="Chat", prompt_template_config=prompt_template_config
@@ -102,28 +101,38 @@ async def chat() -> bool:
        print("\n\nExiting chat...")
        return False

-    arguments = KernelArguments(user_input=user_input, execution_settings=req_settings)
    # Non streaming
    # answer = await kernel.invoke(chat_function, input_vars=context_vars)
    # print(f"Assistant:> {answer}")
-    arguments = KernelArguments(user_input=user_input, execution_settings=req_settings)
+    arguments = KernelArguments(chat_history=chat_history, user_input=user_input, execution_settings=req_settings)

    full_message = None
    print("Assistant:> ", end="")
    async for message in kernel.invoke_stream(chat_function, arguments=arguments):
        print(str(message[0]), end="")
        full_message = message[0] if not full_message else full_message + message[0]
-    chat.add_assistant_message(str(full_message))
+    chat_history.add_assistant_message(str(full_message))
    print("\n")
+
    # The tool message containing cited sources is available in the context
-    if isinstance(full_message, AzureStreamingChatMessageContent):
-        tool_call = ToolCall(full_message.tool_message)
-        chat.add_message(
-            role=ChatRole.TOOL,
-            content=full_message,
-            metadata={OpenAIChatMessageContent.ToolIdProperty: tool_call.function.name},
-        )
-        print(f"Tool:> {full_message.tool_message}")
+    if full_message:
+        chat_history.add_user_message(user_input)
+        if hasattr(full_message, "tool_message"):
+            chat_history.add_message(
+                AzureChatMessageContent(
+                    role="assistant",
+                    tool_calls=[
+                        ToolCall(
+                            id="chat_with_your_data",
+                            function=FunctionCall(name="chat_with_your_data", arguments=""),
+                        )
+                    ],
+                )
+            )
+            chat_history.add_tool_message(full_message.tool_message, {"tool_call_id": "chat_with_your_data"})
+        if full_message.role is None:
+            full_message.role = ChatRole.ASSISTANT
+        chat_history.add_assistant_message(full_message.content)
    return True


@@ -1,8 +1,10 @@
 import logging
 from typing import Any, Dict, List, Literal, Optional, Union

-from pydantic import Field, SerializeAsAny
-from pydantic.dataclasses import dataclass
+from pydantic import AliasGenerator, ConfigDict, Field
+from pydantic.alias_generators import to_camel, to_snake
+from pydantic.functional_validators import AfterValidator
+from typing_extensions import Annotated

 from semantic_kernel.connectors.ai.open_ai.prompt_execution_settings.open_ai_prompt_execution_settings import (
    OpenAIChatPromptExecutionSettings,
@@ -12,68 +14,81 @@ from semantic_kernel.kernel_pydantic import KernelBaseModel
 logger = logging.getLogger(__name__)


-@dataclass
-class ConnectionStringAuthentication:
-    type: Literal["ConnectionString"] = "ConnectionString"
-    connectionString: Optional[str] = None
+class AzureChatRequestBase(KernelBaseModel):
+    model_config = ConfigDict(
+        alias_generator=AliasGenerator(validation_alias=to_camel, serialization_alias=to_snake),
+        use_enum_values=True,
+        extra="allow",
+    )


-@dataclass
-class ApiKeyAuthentication:
-    type: Literal["APIKey"] = "APIKey"
+class ConnectionStringAuthentication(AzureChatRequestBase):
+    type: Annotated[Literal["ConnectionString", "connection_string"], AfterValidator(to_snake)] = "connection_string"
+    connection_string: Optional[str] = None
+
+
+class ApiKeyAuthentication(AzureChatRequestBase):
+    type: Annotated[Literal["APIKey", "api_key"], AfterValidator(to_snake)] = "api_key"
    key: Optional[str] = None


-@dataclass
-class AzureEmbeddingDependency:
-    type: Literal["DeploymentName"] = "DeploymentName"
-    deploymentName: Optional[str] = None
+class AzureEmbeddingDependency(AzureChatRequestBase):
+    type: Annotated[Literal["DeploymentName", "deployment_name"], AfterValidator(to_snake)] = "deployment_name"
+    deployment_name: Optional[str] = None


-@dataclass
-class AzureDataSourceParameters:
-    indexName: str
-    indexLanguage: Optional[str] = None
-    fieldsMapping: Dict[str, Any] = Field(default_factory=dict)
-    inScope: Optional[bool] = True
-    topNDocuments: Optional[int] = 5
-    semanticConfiguration: Optional[str] = None
-    roleInformation: Optional[str] = None
+class DataSourceFieldsMapping(AzureChatRequestBase):
+    title_field: Optional[str] = None
+    url_field: Optional[str] = None
+    filepath_field: Optional[str] = None
+    content_fields: Optional[List[str]] = None
+    vector_fields: Optional[List[str]] = None
+    content_fields_separator: Optional[str] = "\n"
+
+
+class AzureDataSourceParameters(AzureChatRequestBase):
+    index_name: str
+    index_language: Optional[str] = None
+    fields_mapping: Optional[DataSourceFieldsMapping] = None
+    in_scope: Optional[bool] = True
+    top_n_documents: Optional[int] = 5
+    semantic_configuration: Optional[str] = None
+    role_information: Optional[str] = None
    filter: Optional[str] = None
-    embeddingKey: Optional[str] = None
-    embeddingEndpoint: Optional[str] = None
-    embeddingDeploymentName: Optional[str] = None
    strictness: int = 3
-    embeddingDependency: Optional[AzureEmbeddingDependency] = None
+    embedding_dependency: Optional[AzureEmbeddingDependency] = None


-@dataclass
-class AzureCosmosDBDataSource(AzureDataSourceParameters):
+class AzureCosmosDBDataSourceParameters(AzureDataSourceParameters):
    authentication: Optional[ConnectionStringAuthentication] = None
-    databaseName: Optional[str] = None
-    containerName: Optional[str] = None
-    embeddingDependencyType: Optional[AzureEmbeddingDependency] = None
+    database_name: Optional[str] = None
+    container_name: Optional[str] = None
+    embedding_dependency_type: Optional[AzureEmbeddingDependency] = None


-@dataclass
-class AzureAISearchDataSources(AzureDataSourceParameters):
+class AzureCosmosDBDataSource(AzureChatRequestBase):
+    type: Literal["azure_cosmos_db"] = "azure_cosmos_db"
+    parameters: AzureCosmosDBDataSourceParameters
+
+
+class AzureAISearchDataSourceParameters(AzureDataSourceParameters):
    endpoint: Optional[str] = None
-    key: Optional[str] = None
-    queryType: Literal["simple", "semantic", "vector", "vectorSimpleHybrid", "vectorSemanticHybrid"] = "simple"
+    query_type: Annotated[
+        Literal["simple", "semantic", "vector", "vectorSimpleHybrid", "vectorSemanticHybrid"], AfterValidator(to_snake)
+    ] = "simple"
    authentication: Optional[ApiKeyAuthentication] = None


-@dataclass
-class AzureDataSources:
-    """Class to hold Azure AI data source parameters."""
+class AzureAISearchDataSource(AzureChatRequestBase):
+    type: Literal["azure_search"] = "azure_search"
+    parameters: Annotated[dict, AzureAISearchDataSourceParameters]

-    type: Literal["AzureCognitiveSearch", "AzureCosmosDB"] = "AzureCognitiveSearch"
-    parameters: Optional[SerializeAsAny[AzureDataSourceParameters]] = None
+
+DataSource = Annotated[Union[AzureAISearchDataSource, AzureCosmosDBDataSource], Field(discriminator="type")]


-# @dataclass
 class ExtraBody(KernelBaseModel):
-    data_sources: Optional[List[AzureDataSources]] = Field(None, alias="dataSources")
+    data_sources: Optional[List[DataSource]] = None
    input_language: Optional[str] = Field(None, serialization_alias="inputLanguage")
    output_language: Optional[str] = Field(None, serialization_alias="outputLanguage")

@@ -1,4 +1,5 @@
 # Copyright (c) Microsoft. All rights reserved.
+import json
 import logging
 from typing import Any, Dict, Mapping, Optional, Union, overload

@@ -130,7 +131,6 @@ class AzureChatCompletion(AzureOpenAIConfigBase, OpenAIChatCompletionBase, OpenA
        ad_token: Optional[str] = None,
        ad_token_provider: Optional[AsyncAzureADTokenProvider] = None,
        default_headers: Optional[Mapping[str, str]] = None,
-        use_extensions: bool = False,
    ) -> None:
        """
        Initialize an AzureChatCompletion service.
@@ -154,9 +154,6 @@ class AzureChatCompletion(AzureOpenAIConfigBase, OpenAIChatCompletionBase, OpenA
        default_headers: The default headers mapping of string keys to
            string values for HTTP requests. (Optional)
        log: The logger instance to use. (Optional)
-        use_extensions: Whether to use extensions, for example when chatting with data. (Optional)
-            When True, base_url is overwritten to '{endpoint}/openai/deployments/{deployment_name}/extensions'.
-            The default value is False.
        """

    def __init__(
@@ -171,7 +168,6 @@ class AzureChatCompletion(AzureOpenAIConfigBase, OpenAIChatCompletionBase, OpenA
        ad_token_provider: Optional[AsyncAzureADTokenProvider] = None,
        default_headers: Optional[Mapping[str, str]] = None,
        async_client: Optional[AsyncAzureOpenAI] = None,
-        use_extensions: bool = False,
    ) -> None:
        """
        Initialize an AzureChatCompletion service.
@@ -201,14 +197,11 @@ class AzureChatCompletion(AzureOpenAIConfigBase, OpenAIChatCompletionBase, OpenA
            default_headers: The default headers mapping of string keys to
                string values for HTTP requests. (Optional)
            async_client {Optional[AsyncAzureOpenAI]} -- An existing client to use. (Optional)
-            use_extensions: Whether to use extensions, for example when chatting with data. (Optional)
-                When True, base_url is overwritten to '{endpoint}/openai/deployments/{deployment_name}/extensions'.
-                The default value is False.
        """
        if base_url and isinstance(base_url, str):
            base_url = HttpsUrl(base_url)
-        if use_extensions and endpoint and deployment_name:
-            base_url = HttpsUrl(f"{str(endpoint).rstrip('/')}/openai/deployments/{deployment_name}/extensions")
+        if endpoint and deployment_name:
+            base_url = HttpsUrl(f"{str(endpoint).rstrip('/')}/openai/deployments/{deployment_name}")
        super().__init__(
            deployment_name=deployment_name,
            endpoint=endpoint if not isinstance(endpoint, str) else HttpsUrl(endpoint),
@@ -295,10 +288,8 @@ class AzureChatCompletion(AzureOpenAIConfigBase, OpenAIChatCompletionBase, OpenA
        else:
            content = choice.delta
        if content.model_extra is not None and "context" in content.model_extra:
-            if "messages" in content.model_extra["context"]:
-                for message in content.model_extra["context"]["messages"]:
-                    if "tool" in message["role"]:
-                        return message["content"]
+            return json.dumps(content.model_extra["context"])
+
        return None

    def get_chat_message_content_type(self) -> str:
@@ -295,7 +295,7 @@ def azure_aisearch_settings_from_dot_env_as_dict() -> Dict[str, str]:
        Dict[str, str]: the Azure AI search environment variables
    """
    api_key, url, index_name = azure_aisearch_settings_from_dot_env(include_index_name=True)
-    return {"key": api_key, "endpoint": url, "indexName": index_name}
+    return {"authentication": {"type": "api_key", "key": api_key}, "endpoint": url, "index_name": index_name}


 def azure_key_vault_settings_from_dot_env(
@@ -9,8 +9,8 @@ import pytest

 import semantic_kernel.connectors.ai.open_ai as sk_oai
 from semantic_kernel.connectors.ai.open_ai.prompt_execution_settings.azure_chat_prompt_execution_settings import (
-    AzureAISearchDataSources,
-    AzureDataSources,
+    AzureAISearchDataSource,
+    AzureAISearchDataSourceParameters,
    ExtraBody,
 )
 from semantic_kernel.connectors.ai.prompt_execution_settings import PromptExecutionSettings
@@ -93,30 +93,28 @@ async def create_with_data_chat_function(get_aoai_config, kernel: Kernel, create

        extra = ExtraBody(
            data_sources=[
-                AzureDataSources(
-                    type="AzureCognitiveSearch",
-                    parameters=AzureAISearchDataSources(
-                        indexName=collection,
+                AzureAISearchDataSource(
+                    parameters=AzureAISearchDataSourceParameters(
+                        index_name=collection,
                        endpoint=search_endpoint,
-                        key=search_api_key,
-                        queryType="simple",
-                        fieldsMapping={
+                        authentication={"type": "api_key", "api_key": search_api_key},
+                        query_type="simple",
+                        fields_mapping={
                            "titleField": "Description",
                            "contentFields": ["Text"],
                        },
-                        topNDocuments=1,
+                        top_n_documents=1,
                    ),
                )
            ]
        )
-
+        print(f"deployment: {deployment_name}, endpoint: {endpoint}")
        chat_service = sk_oai.AzureChatCompletion(
            service_id="chat-gpt-extensions",
            deployment_name=deployment_name,
            api_key=api_key,
            endpoint=endpoint,
-            api_version="2023-12-01-preview",
-            use_extensions=True,
+            api_version="2024-02-15-preview",
        )
        kernel.add_service(chat_service)

@@ -143,7 +141,6 @@ async def create_with_data_chat_function(get_aoai_config, kernel: Kernel, create


@pytest.mark.asyncio
-@pytest.mark.xfail(reason="The test is failing a 400 saying the request body is invalid. Will investigate.")
@pytestmark
 async def test_azure_e2e_chat_completion_with_extensions(
    create_with_data_chat_function,
@@ -17,9 +17,8 @@ from semantic_kernel.connectors.ai.open_ai.exceptions.content_filter_ai_exceptio
    ContentFilterResultSeverity,
 )
 from semantic_kernel.connectors.ai.open_ai.prompt_execution_settings.azure_chat_prompt_execution_settings import (
-    AzureAISearchDataSources,
+    AzureAISearchDataSource,
    AzureChatPromptExecutionSettings,
-    AzureDataSources,
    ExtraBody,
 )
 from semantic_kernel.contents.chat_history import ChatHistory
@@ -334,7 +333,6 @@ async def test_azure_chat_completion_with_data_call_with_parameters(
        endpoint=endpoint,
        api_version=api_version,
        api_key=api_key,
-        use_extensions=True,
    )

    await azure_chat_completion.complete_chat(
@@ -368,15 +366,20 @@ async def test_azure_chat_completion_call_with_data_parameters_and_function_call
    prompt = "hello world"
    chat_history.add_user_message(prompt)

-    ai_source = AzureAISearchDataSources(indexName="test-index", endpoint="test-endpoint", key="test-key")
-    extra = ExtraBody(data_sources=[AzureDataSources(type="AzureCognitiveSearch", parameters=ai_source)])
+    ai_source = AzureAISearchDataSource(
+        parameters={
+            "indexName": "test-index",
+            "endpoint": "test-endpoint",
+            "authentication": {"type": "api_key", "api_key": "test-key"},
+        }
+    )
+    extra = ExtraBody(data_sources=[ai_source])

    azure_chat_completion = AzureChatCompletion(
        deployment_name=deployment_name,
        endpoint=endpoint,
        api_key=api_key,
        api_version=api_version,
-        use_extensions=True,
    )

    functions = [{"name": "test-function", "description": "test-description"}]
@@ -426,8 +429,14 @@ async def test_azure_chat_completion_call_with_data_with_parameters_and_Stop_Def
    stop = ["!"]
    complete_prompt_execution_settings.stop = stop

-    ai_source = AzureAISearchDataSources(indexName="test-index", endpoint="test-endpoint", key="test-key")
-    extra = ExtraBody(data_sources=[AzureDataSources(type="AzureCognitiveSearch", parameters=ai_source)])
+    ai_source = AzureAISearchDataSource(
+        parameters={
+            "indexName": "test-index",
+            "endpoint": "test-endpoint",
+            "authentication": {"type": "api_key", "api_key": "test-key"},
+        }
+    )
+    extra = ExtraBody(data_sources=[ai_source])

    complete_prompt_execution_settings.extra_body = extra

@@ -436,7 +445,6 @@ async def test_azure_chat_completion_call_with_data_with_parameters_and_Stop_Def
        endpoint=endpoint,
        api_key=api_key,
        api_version=api_version,
-        use_extensions=True,
    )

    await azure_chat_completion.complete_chat(chat_history, complete_prompt_execution_settings, kernel=kernel)
@@ -3,9 +3,8 @@
 import pytest

 from semantic_kernel.connectors.ai.open_ai.prompt_execution_settings.azure_chat_prompt_execution_settings import (
-    AzureAISearchDataSources,
+    AzureAISearchDataSource,
    AzureChatPromptExecutionSettings,
-    AzureDataSources,
    ExtraBody,
 )
 from semantic_kernel.connectors.ai.open_ai.prompt_execution_settings.open_ai_prompt_execution_settings import (
@@ -195,9 +194,14 @@ def test_create_options():


 def test_create_options_azure_data():
-    az_source = AzureAISearchDataSources(indexName="test-index", endpoint="test-endpoint", key="test-key")
-    az_data = AzureDataSources(type="AzureCognitiveSearch", parameters=az_source)
-    extra = ExtraBody(dataSources=[az_data])
+    az_source = AzureAISearchDataSource(
+        parameters={
+            "indexName": "test-index",
+            "endpoint": "test-endpoint",
+            "authentication": {"type": "api_key", "api_key": "test-key"},
+        }
+    )
+    extra = ExtraBody(dataSources=[az_source])
    settings = AzureChatPromptExecutionSettings(extra_body=extra)
    options = settings.prepare_settings_dict()
    assert options["extra_body"] == extra.model_dump(exclude_none=True, by_alias=True)