mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
6a0d03914c
* test: drop the cwd-relative sys.path.insert calls from the test suite
TQ003 stands at 1,077 across 1,058 files, and 1,015 of them are the same shape:
sys.path.insert(0, os.path.abspath("../..")) and its deeper siblings. The
argument resolves against the working directory rather than the file, so from
the repo root, where every job runs pytest, it inserts the directory two levels
above the checkout. It has never pointed at litellm. The package is installed
into the environment anyway, which is what actually makes the import work, and
what the rule's message has said all along.
Removing them leaves 1,634 imports of sys and os with no remaining reference,
and those go too, except where another test module imports the name back out of
the file. The rest of TQ003 is 62 call sites that resolve against __file__ or a
variable, which are a different question and are left alone.
Collection is identical either way: 45,871 tests and the same 51 pre-existing
collection errors before and after, and ruff reports no new undefined name.
* test: drop the duplicate imports the sys.path sweep exposed to F811
* test(pre-call-utils): restore the os import the new bedrock tests need
688 lines
23 KiB
Python
688 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 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}")
|