在Azure Machine Learning的PromptFlow中调用CosmosDB时遇400错误

问题背景 #

基于Azure Contoso Chat仓库开发项目，使用Azure Machine Learning（AML）的PromptFlow组件调用CosmosDB时，在AI Studio测试流程遇到错误：
Request failed with status code 400 - NotSupported ConnectionType 0 for connection contoso-cosmos.
本地运行流程正常，但AML环境中报错，且已确认连接信息无误。

解决思路 #

1. 更换连接类型 #

在AML AI Studio中，删除原有的自定义连接（CustomConnection），重新创建Azure Cosmos DB专属连接类型。自定义连接在AML托管环境中易出现类型识别异常，导致ConnectionType 0错误。创建时填写endpoint、密钥、数据库ID、容器ID等必要参数即可。

2. 调整工具代码的连接类 #

将代码中的CustomConnection替换为AzureCosmosDBConnection，直接使用内置属性读取连接信息，避免手动解析配置的类型匹配问题。修改后的代码如下：

from typing import Dict
from promptflow import tool
from azure.cosmos import CosmosClient
from promptflow.connections import AzureCosmosDBConnection


@tool
def customer_lookup(customerId: str, conn: AzureCosmosDBConnection) -> str:
    client = CosmosClient(url=conn.endpoint, credential=conn.key)
    db = client.get_database_client(conn.database_id)
    container = db.get_container_client(conn.container_id)
    response = container.read_item(item=customerId, partition_key=customerId)
    response["orders"] = response["orders"][:2]
    return response

3. 验证托管环境权限 #

确保AML工作区的系统分配托管身份拥有CosmosDB的访问权限：

• 进入CosmosDB的IAM设置页面
• 添加AML工作区的系统托管身份，赋予「Cosmos DB Account Contributor」或「Cosmos DB Reader」角色

4. 重新绑定流程连接 #

在AI Studio的PromptFlow编辑页面，确认工具节点的conn参数已绑定到新建的CosmosDB专属连接，保存后重新测试流程。

内容的提问来源于stack exchange，提问作者user24253039