691 lines
23 KiB
Python
691 lines
23 KiB
Python
"""
|
|
Tests for Azure Foundry Agent Service integration.
|
|
|
|
These tests require an Azure Foundry Agent Service endpoint and a pre-configured agent.
|
|
|
|
The Azure Foundry Agent Service uses the Assistants API pattern:
|
|
1. Create a thread
|
|
2. Add messages to the thread
|
|
3. Create and poll a run
|
|
4. Get the agent's response messages
|
|
|
|
Model format: azure_ai/agents/<agent_id>
|
|
|
|
API Base format: https://<AIFoundryResourceName>.services.ai.azure.com/api/projects/<ProjectName>
|
|
|
|
Authentication: Uses Azure AD Bearer tokens (not API keys)
|
|
Get token via: az account get-access-token --resource 'https://ai.azure.com'
|
|
|
|
Example environment variables:
|
|
AZURE_AGENTS_API_BASE=https://litellm-ci-cd-prod.services.ai.azure.com/api/projects/litellm-ci-cd
|
|
AZURE_AGENTS_API_KEY=<Azure AD Bearer token>
|
|
|
|
See: https://learn.microsoft.com/en-us/azure/ai-foundry/agents/quickstart
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import sys
|
|
|
|
sys.path.insert(0, os.path.abspath("../.."))
|
|
|
|
import pytest
|
|
from unittest.mock import MagicMock
|
|
|
|
import litellm
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_azure_ai_agents_acompletion_non_streaming():
|
|
"""
|
|
Test non-streaming acompletion call to Azure Foundry Agent Service.
|
|
Uses the multi-step flow: create thread -> add messages -> create/poll run -> get messages
|
|
"""
|
|
api_base = os.environ.get("AZURE_AGENTS_API_BASE")
|
|
api_key = os.environ.get("AZURE_AGENTS_API_KEY")
|
|
agent_id = os.environ.get("AZURE_AGENTS_AGENT_ID", "asst_hbnoK9BOCcHhC3lC4MDroVGG")
|
|
|
|
if not api_base or not api_key:
|
|
pytest.skip(
|
|
"AZURE_AGENTS_API_BASE and AZURE_AGENTS_API_KEY environment variables required"
|
|
)
|
|
|
|
response = await litellm.acompletion(
|
|
model=f"azure_ai/agents/{agent_id}",
|
|
messages=[{"role": "user", "content": "Hi Agent, what is 25 * 4?"}],
|
|
api_base=api_base,
|
|
api_key=api_key,
|
|
stream=False,
|
|
)
|
|
|
|
assert response is not None
|
|
assert response.choices is not None
|
|
assert len(response.choices) > 0
|
|
assert response.choices[0].message is not None
|
|
assert response.choices[0].message.content is not None
|
|
assert len(response.choices[0].message.content) > 0
|
|
|
|
# Verify thread_id is returned for conversation continuity
|
|
if hasattr(response, "_hidden_params") and response._hidden_params:
|
|
assert "thread_id" in response._hidden_params
|
|
|
|
print(f"Response: {response.choices[0].message.content}")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_azure_ai_agents_acompletion_streaming():
|
|
"""
|
|
Test native streaming acompletion call to Azure Foundry Agent Service.
|
|
Uses the create-thread-and-run endpoint with stream=True for SSE streaming.
|
|
"""
|
|
api_base = os.environ.get("AZURE_AGENTS_API_BASE")
|
|
api_key = os.environ.get("AZURE_AGENTS_API_KEY")
|
|
agent_id = os.environ.get("AZURE_AGENTS_AGENT_ID", "asst_hbnoK9BOCcHhC3lC4MDroVGG")
|
|
|
|
if not api_base or not api_key:
|
|
pytest.skip(
|
|
"AZURE_AGENTS_API_BASE and AZURE_AGENTS_API_KEY environment variables required"
|
|
)
|
|
|
|
response = await litellm.acompletion(
|
|
model=f"azure_ai/agents/{agent_id}",
|
|
messages=[{"role": "user", "content": "Hi Agent, what is 10 + 5?"}],
|
|
api_base=api_base,
|
|
api_key=api_key,
|
|
stream=True,
|
|
)
|
|
|
|
# Native streaming - collect chunks from the async iterator
|
|
chunks = []
|
|
full_content = ""
|
|
async for chunk in response:
|
|
print("Streaming chunk: ", chunk)
|
|
chunks.append(chunk)
|
|
if hasattr(chunk, "choices") and chunk.choices:
|
|
delta = chunk.choices[0].delta
|
|
if hasattr(delta, "content") and delta.content:
|
|
full_content += delta.content
|
|
|
|
assert len(chunks) > 0, "Expected at least one streaming chunk"
|
|
assert len(full_content) > 0, "Expected content from streaming response"
|
|
print(f"Streamed response ({len(chunks)} chunks): {full_content}")
|
|
|
|
|
|
def test_azure_ai_agents_is_agents_route():
|
|
"""
|
|
Test the is_azure_ai_agents_route detection method.
|
|
"""
|
|
from litellm.llms.azure_ai.agents.transformation import AzureAIAgentsConfig
|
|
|
|
# Should be recognized as agents route
|
|
assert (
|
|
AzureAIAgentsConfig.is_azure_ai_agents_route("azure_ai/agents/asst_123") is True
|
|
)
|
|
assert AzureAIAgentsConfig.is_azure_ai_agents_route("agents/asst_123") is True
|
|
|
|
# Should NOT be recognized as agents route
|
|
assert AzureAIAgentsConfig.is_azure_ai_agents_route("azure_ai/gpt-4") is False
|
|
assert AzureAIAgentsConfig.is_azure_ai_agents_route("gpt-4") is False
|
|
|
|
|
|
def test_azure_ai_get_azure_ai_route():
|
|
"""
|
|
Test the get_azure_ai_route dispatch method.
|
|
"""
|
|
from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
|
|
|
|
# Should return "agents" for agents routes
|
|
assert AzureFoundryModelInfo.get_azure_ai_route("agents/asst_123") == "agents"
|
|
assert (
|
|
AzureFoundryModelInfo.get_azure_ai_route("azure_ai/agents/asst_abc") == "agents"
|
|
)
|
|
|
|
# Should return "default" for non-agents routes
|
|
assert AzureFoundryModelInfo.get_azure_ai_route("gpt-4") == "default"
|
|
assert AzureFoundryModelInfo.get_azure_ai_route("claude-3-sonnet") == "default"
|
|
assert AzureFoundryModelInfo.get_azure_ai_route("azure_ai/gpt-4o") == "default"
|
|
|
|
|
|
def test_azure_ai_agents_get_agent_id_from_model():
|
|
"""
|
|
Test agent ID extraction from model name.
|
|
"""
|
|
from litellm.llms.azure_ai.agents.transformation import AzureAIAgentsConfig
|
|
|
|
# Test with full model name
|
|
agent_id = AzureAIAgentsConfig.get_agent_id_from_model(
|
|
"azure_ai/agents/asst_abc123"
|
|
)
|
|
assert agent_id == "asst_abc123"
|
|
|
|
# Test with just agents/id
|
|
agent_id = AzureAIAgentsConfig.get_agent_id_from_model("agents/asst_xyz789")
|
|
assert agent_id == "asst_xyz789"
|
|
|
|
# Test with just agent ID (fallback)
|
|
agent_id = AzureAIAgentsConfig.get_agent_id_from_model("asst_plain")
|
|
assert agent_id == "asst_plain"
|
|
|
|
|
|
def test_azure_ai_agents_config_get_agent_id():
|
|
"""
|
|
Test agent ID extraction via config method.
|
|
"""
|
|
from litellm.llms.azure_ai.agents.transformation import AzureAIAgentsConfig
|
|
|
|
config = AzureAIAgentsConfig()
|
|
|
|
# Test with full model name
|
|
agent_id = config._get_agent_id("azure_ai/agents/asst_abc123", {})
|
|
assert agent_id == "asst_abc123"
|
|
|
|
# Test with optional_params override
|
|
agent_id = config._get_agent_id(
|
|
"azure_ai/agents/asst_abc123", {"agent_id": "asst_override"}
|
|
)
|
|
assert agent_id == "asst_override"
|
|
|
|
# Test with assistant_id in optional_params
|
|
agent_id = config._get_agent_id(
|
|
"azure_ai/agents/asst_abc123", {"assistant_id": "asst_assistant"}
|
|
)
|
|
assert agent_id == "asst_assistant"
|
|
|
|
|
|
def test_azure_ai_agents_config_get_complete_url():
|
|
"""
|
|
Test that AzureAIAgentsConfig correctly generates base URLs.
|
|
"""
|
|
from litellm.llms.azure_ai.agents.transformation import AzureAIAgentsConfig
|
|
|
|
config = AzureAIAgentsConfig()
|
|
|
|
# Test URL generation
|
|
url = config.get_complete_url(
|
|
api_base="https://test-project.services.ai.azure.com",
|
|
api_key=None,
|
|
model="agents/asst_123",
|
|
optional_params={},
|
|
litellm_params={},
|
|
stream=False,
|
|
)
|
|
assert url == "https://test-project.services.ai.azure.com"
|
|
|
|
# Test URL with trailing slash
|
|
url_with_slash = config.get_complete_url(
|
|
api_base="https://test-project.services.ai.azure.com/",
|
|
api_key=None,
|
|
model="agents/asst_123",
|
|
optional_params={},
|
|
litellm_params={},
|
|
stream=False,
|
|
)
|
|
assert url_with_slash == "https://test-project.services.ai.azure.com"
|
|
|
|
|
|
def test_azure_ai_agents_config_transform_request():
|
|
"""
|
|
Test that AzureAIAgentsConfig correctly transforms requests.
|
|
"""
|
|
from litellm.llms.azure_ai.agents.transformation import AzureAIAgentsConfig
|
|
|
|
config = AzureAIAgentsConfig()
|
|
|
|
messages = [
|
|
{"role": "system", "content": "You are a helpful assistant."},
|
|
{"role": "user", "content": "What is 2 + 2?"},
|
|
]
|
|
|
|
request = config.transform_request(
|
|
model="azure_ai/agents/asst_123",
|
|
messages=messages,
|
|
optional_params={},
|
|
litellm_params={"stream": False},
|
|
headers={},
|
|
)
|
|
|
|
assert request["agent_id"] == "asst_123"
|
|
assert "messages" in request
|
|
assert len(request["messages"]) == 2
|
|
assert request["messages"][0]["role"] == "system"
|
|
assert request["messages"][1]["role"] == "user"
|
|
assert "api_version" in request
|
|
assert request["api_version"] == "2025-05-01"
|
|
|
|
|
|
def test_azure_ai_agents_provider_detection():
|
|
"""
|
|
Test that the azure_ai provider is correctly detected from model name.
|
|
"""
|
|
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
|
|
|
model, provider, api_key, api_base = get_llm_provider(
|
|
model="azure_ai/agents/asst_abc123",
|
|
api_base="https://test.services.ai.azure.com",
|
|
)
|
|
|
|
assert provider == "azure_ai"
|
|
assert model == "agents/asst_abc123"
|
|
|
|
|
|
def test_azure_ai_agents_validate_environment():
|
|
"""
|
|
Test that headers are correctly set up with Bearer token authentication.
|
|
|
|
Azure Foundry Agents uses Bearer token authentication (Azure AD tokens).
|
|
"""
|
|
from litellm.llms.azure_ai.agents.transformation import AzureAIAgentsConfig
|
|
|
|
config = AzureAIAgentsConfig()
|
|
|
|
headers = config.validate_environment(
|
|
headers={},
|
|
model="agents/asst_123",
|
|
messages=[],
|
|
optional_params={},
|
|
litellm_params={},
|
|
api_key="test-azure-ad-token",
|
|
api_base="https://test.services.ai.azure.com/api/projects/test-project",
|
|
)
|
|
|
|
assert headers["Content-Type"] == "application/json"
|
|
assert headers["Authorization"] == "Bearer test-azure-ad-token"
|
|
|
|
|
|
def test_azure_ai_agents_handler_url_builders():
|
|
"""
|
|
Test the URL building methods in the handler.
|
|
|
|
Azure Foundry Agents API uses direct paths without /openai/ prefix.
|
|
See: https://learn.microsoft.com/en-us/azure/ai-foundry/agents/quickstart
|
|
"""
|
|
from litellm.llms.azure_ai.agents.handler import AzureAIAgentsHandler
|
|
|
|
handler = AzureAIAgentsHandler()
|
|
api_base = "https://test.services.ai.azure.com/api/projects/test-project"
|
|
api_version = "2025-05-01"
|
|
thread_id = "thread_abc123"
|
|
run_id = "run_xyz789"
|
|
|
|
# Test thread URL - direct path without /openai/ prefix
|
|
thread_url = handler._build_thread_url(api_base, api_version)
|
|
assert thread_url == f"{api_base}/threads?api-version={api_version}"
|
|
|
|
# Test messages URL
|
|
messages_url = handler._build_messages_url(api_base, thread_id, api_version)
|
|
assert (
|
|
messages_url
|
|
== f"{api_base}/threads/{thread_id}/messages?api-version={api_version}"
|
|
)
|
|
|
|
# Test runs URL
|
|
runs_url = handler._build_runs_url(api_base, thread_id, api_version)
|
|
assert runs_url == f"{api_base}/threads/{thread_id}/runs?api-version={api_version}"
|
|
|
|
# Test run status URL
|
|
status_url = handler._build_run_status_url(api_base, thread_id, run_id, api_version)
|
|
assert (
|
|
status_url
|
|
== f"{api_base}/threads/{thread_id}/runs/{run_id}?api-version={api_version}"
|
|
)
|
|
|
|
|
|
def test_azure_ai_agents_extract_content_from_messages():
|
|
"""
|
|
Test content extraction from Azure Agents message response.
|
|
"""
|
|
from litellm.llms.azure_ai.agents.handler import AzureAIAgentsHandler
|
|
|
|
handler = AzureAIAgentsHandler()
|
|
|
|
# Test typical message response
|
|
messages_data = {
|
|
"data": [
|
|
{
|
|
"id": "msg_123",
|
|
"role": "assistant",
|
|
"content": [{"type": "text", "text": {"value": "The answer is 100."}}],
|
|
},
|
|
{
|
|
"id": "msg_122",
|
|
"role": "user",
|
|
"content": [{"type": "text", "text": {"value": "What is 25 * 4?"}}],
|
|
},
|
|
]
|
|
}
|
|
|
|
content, annotations = handler._extract_content_from_messages(messages_data)
|
|
assert content == "The answer is 100."
|
|
assert annotations is None
|
|
|
|
# Test empty response
|
|
empty_data = {"data": []}
|
|
content, annotations = handler._extract_content_from_messages(empty_data)
|
|
assert content == ""
|
|
assert annotations is None
|
|
|
|
|
|
def test_azure_ai_agents_extract_content_with_annotations():
|
|
"""
|
|
Test that annotations (e.g., Bing Search citations) are extracted from
|
|
Azure Agents message responses and transformed to OpenAI-compatible format.
|
|
|
|
Ref: https://github.com/BerriAI/litellm/issues/19126
|
|
"""
|
|
from litellm.llms.azure_ai.agents.handler import AzureAIAgentsHandler
|
|
|
|
handler = AzureAIAgentsHandler()
|
|
|
|
messages_data = {
|
|
"data": [
|
|
{
|
|
"id": "msg_abc",
|
|
"role": "assistant",
|
|
"content": [
|
|
{
|
|
"type": "text",
|
|
"text": {
|
|
"value": "According to sources [1], the answer is yes.",
|
|
"annotations": [
|
|
{
|
|
"type": "url_citation",
|
|
"text": "[1]",
|
|
"start_index": 22,
|
|
"end_index": 25,
|
|
"url_citation": {
|
|
"url": "https://example.com/source",
|
|
"title": "Example Source",
|
|
},
|
|
}
|
|
],
|
|
},
|
|
}
|
|
],
|
|
}
|
|
]
|
|
}
|
|
|
|
content, annotations = handler._extract_content_from_messages(messages_data)
|
|
assert content == "According to sources [1], the answer is yes."
|
|
assert annotations is not None
|
|
assert len(annotations) == 1
|
|
assert annotations[0]["type"] == "url_citation"
|
|
assert annotations[0]["url_citation"]["url"] == "https://example.com/source"
|
|
assert annotations[0]["url_citation"]["title"] == "Example Source"
|
|
# start/end_index should be moved into url_citation for OpenAI compatibility
|
|
assert annotations[0]["url_citation"]["start_index"] == 22
|
|
assert annotations[0]["url_citation"]["end_index"] == 25
|
|
|
|
|
|
def test_azure_ai_agents_build_model_response_with_annotations():
|
|
"""
|
|
Test that _build_model_response includes annotations in the Message object.
|
|
"""
|
|
from litellm.llms.azure_ai.agents.handler import AzureAIAgentsHandler
|
|
from litellm.types.utils import ModelResponse
|
|
|
|
handler = AzureAIAgentsHandler()
|
|
model_response = ModelResponse()
|
|
|
|
annotations = [
|
|
{
|
|
"type": "url_citation",
|
|
"url_citation": {
|
|
"url": "https://example.com",
|
|
"title": "Example",
|
|
"start_index": 0,
|
|
"end_index": 5,
|
|
},
|
|
}
|
|
]
|
|
|
|
result = handler._build_model_response(
|
|
model="azure_ai/agents/asst_123",
|
|
content="Hello [1]",
|
|
model_response=model_response,
|
|
thread_id="thread_abc",
|
|
messages=[{"role": "user", "content": "test"}],
|
|
annotations=annotations,
|
|
)
|
|
|
|
assert result.choices[0].message.content == "Hello [1]"
|
|
assert result.choices[0].message.annotations is not None
|
|
assert len(result.choices[0].message.annotations) == 1
|
|
assert result.choices[0].message.annotations[0]["type"] == "url_citation"
|
|
|
|
|
|
def test_azure_ai_agents_build_model_response_without_annotations():
|
|
"""
|
|
Test that _build_model_response works correctly without annotations.
|
|
"""
|
|
from litellm.llms.azure_ai.agents.handler import AzureAIAgentsHandler
|
|
from litellm.types.utils import ModelResponse
|
|
|
|
handler = AzureAIAgentsHandler()
|
|
model_response = ModelResponse()
|
|
|
|
result = handler._build_model_response(
|
|
model="azure_ai/agents/asst_123",
|
|
content="Hello",
|
|
model_response=model_response,
|
|
thread_id="thread_abc",
|
|
messages=[{"role": "user", "content": "test"}],
|
|
)
|
|
|
|
assert result.choices[0].message.content == "Hello"
|
|
assert getattr(result.choices[0].message, "annotations", None) is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_azure_ai_agents_streaming_annotations_from_completed_message():
|
|
"""
|
|
Test that annotations from thread.message.completed SSE events are collected
|
|
and attached to the final chunk's delta.
|
|
|
|
Ref: https://github.com/BerriAI/litellm/issues/19126
|
|
"""
|
|
from litellm.llms.azure_ai.agents.handler import AzureAIAgentsHandler
|
|
|
|
handler = AzureAIAgentsHandler()
|
|
|
|
# SSE lines simulating a stream with annotations in thread.message.completed
|
|
completed_data = {
|
|
"content": [
|
|
{
|
|
"type": "text",
|
|
"text": {
|
|
"value": "According to [1], the answer is 42.",
|
|
"annotations": [
|
|
{
|
|
"type": "url_citation",
|
|
"text": "[1]",
|
|
"start_index": 12,
|
|
"end_index": 15,
|
|
"url_citation": {
|
|
"url": "https://example.com/citation",
|
|
"title": "Citation Source",
|
|
},
|
|
}
|
|
],
|
|
},
|
|
}
|
|
]
|
|
}
|
|
|
|
sse_lines = [
|
|
"event: thread.created",
|
|
"",
|
|
'data: {"id": "thread_stream_123"}',
|
|
"",
|
|
"event: thread.message.delta",
|
|
"",
|
|
'data: {"delta": {"content": [{"type": "text", "text": {"value": "According to [1], the answer is 42."}}]}}',
|
|
"",
|
|
"event: thread.message.completed",
|
|
"",
|
|
f"data: {json.dumps(completed_data)}",
|
|
"",
|
|
"data: [DONE]",
|
|
]
|
|
|
|
async def mock_aiter_lines():
|
|
for line in sse_lines:
|
|
yield line
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.aiter_lines = MagicMock(return_value=mock_aiter_lines())
|
|
|
|
chunks = []
|
|
async for chunk in handler._process_sse_stream(
|
|
mock_response, "azure_ai/agents/asst_123"
|
|
):
|
|
chunks.append(chunk)
|
|
|
|
# Should have content chunks + final [DONE] chunk
|
|
assert len(chunks) >= 1
|
|
final_chunk = chunks[-1]
|
|
assert final_chunk.choices[0].finish_reason == "stop"
|
|
assert final_chunk.choices[0].delta.annotations is not None
|
|
assert len(final_chunk.choices[0].delta.annotations) == 1
|
|
ann = final_chunk.choices[0].delta.annotations[0]
|
|
assert ann["type"] == "url_citation"
|
|
assert ann["url_citation"]["url"] == "https://example.com/citation"
|
|
assert ann["url_citation"]["title"] == "Citation Source"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_azure_ai_agents_streaming_accumulates_annotations_from_multiple_text_items():
|
|
"""
|
|
Test that annotations from multiple text content items in thread.message.completed
|
|
are accumulated (not overwritten).
|
|
|
|
Ref: Greptile review on PR #23849
|
|
"""
|
|
from litellm.llms.azure_ai.agents.handler import AzureAIAgentsHandler
|
|
|
|
handler = AzureAIAgentsHandler()
|
|
|
|
# Two text blocks, each with distinct citations
|
|
completed_data = {
|
|
"content": [
|
|
{
|
|
"type": "text",
|
|
"text": {
|
|
"value": "First source [1].",
|
|
"annotations": [
|
|
{
|
|
"type": "url_citation",
|
|
"text": "[1]",
|
|
"start_index": 12,
|
|
"end_index": 15,
|
|
"url_citation": {
|
|
"url": "https://example.com/first",
|
|
"title": "First",
|
|
},
|
|
}
|
|
],
|
|
},
|
|
},
|
|
{
|
|
"type": "text",
|
|
"text": {
|
|
"value": "Second source [2].",
|
|
"annotations": [
|
|
{
|
|
"type": "url_citation",
|
|
"text": "[2]",
|
|
"start_index": 13,
|
|
"end_index": 16,
|
|
"url_citation": {
|
|
"url": "https://example.com/second",
|
|
"title": "Second",
|
|
},
|
|
}
|
|
],
|
|
},
|
|
},
|
|
]
|
|
}
|
|
|
|
sse_lines = [
|
|
"event: thread.created",
|
|
"",
|
|
'data: {"id": "thread_multi"}',
|
|
"",
|
|
"event: thread.message.completed",
|
|
"",
|
|
f"data: {json.dumps(completed_data)}",
|
|
"",
|
|
"data: [DONE]",
|
|
]
|
|
|
|
async def mock_aiter_lines():
|
|
for line in sse_lines:
|
|
yield line
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.aiter_lines = MagicMock(return_value=mock_aiter_lines())
|
|
|
|
chunks = []
|
|
async for chunk in handler._process_sse_stream(
|
|
mock_response, "azure_ai/agents/asst_123"
|
|
):
|
|
chunks.append(chunk)
|
|
|
|
final_chunk = chunks[-1]
|
|
assert final_chunk.choices[0].delta.annotations is not None
|
|
assert len(final_chunk.choices[0].delta.annotations) == 2
|
|
urls = [a["url_citation"]["url"] for a in final_chunk.choices[0].delta.annotations]
|
|
assert "https://example.com/first" in urls
|
|
assert "https://example.com/second" in urls
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_azure_ai_agents_conversation_continuity():
|
|
"""
|
|
Test that thread_id can be used for conversation continuity.
|
|
"""
|
|
api_base = os.environ.get("AZURE_AGENTS_API_BASE")
|
|
api_key = os.environ.get("AZURE_AGENTS_API_KEY")
|
|
agent_id = os.environ.get("AZURE_AGENTS_AGENT_ID", "asst_hbnoK9BOCcHhC3lC4MDroVGG")
|
|
|
|
if not api_base or not api_key:
|
|
pytest.skip(
|
|
"AZURE_AGENTS_API_BASE and AZURE_AGENTS_API_KEY environment variables required"
|
|
)
|
|
|
|
try:
|
|
# First message
|
|
response1 = await litellm.acompletion(
|
|
model=f"azure_ai/agents/{agent_id}",
|
|
messages=[{"role": "user", "content": "My name is Alice. Remember this."}],
|
|
api_base=api_base,
|
|
api_key=api_key,
|
|
stream=False,
|
|
)
|
|
|
|
assert response1 is not None
|
|
|
|
# Get thread_id for continuity
|
|
thread_id = None
|
|
if hasattr(response1, "_hidden_params") and response1._hidden_params:
|
|
thread_id = response1._hidden_params.get("thread_id")
|
|
|
|
if thread_id:
|
|
# Second message using the same thread
|
|
response2 = await litellm.acompletion(
|
|
model=f"azure_ai/agents/{agent_id}",
|
|
messages=[{"role": "user", "content": "What is my name?"}],
|
|
api_base=api_base,
|
|
api_key=api_key,
|
|
thread_id=thread_id, # Continue the conversation
|
|
stream=False,
|
|
)
|
|
|
|
assert response2 is not None
|
|
# The agent should remember the name from the previous message
|
|
print(f"Response to name question: {response2.choices[0].message.content}")
|
|
|
|
except Exception as e:
|
|
pytest.skip(f"Azure Agent Service not available: {e}")
|