db620d33df
CodeQL / Analyze (csharp) (push) Waiting to run
CodeQL / Analyze (python) (push) Waiting to run
dotnet-build-and-test / dotnet-test-functions (push) Has been cancelled
dotnet-build-and-test / paths-filter (push) Has been cancelled
dotnet-build-and-test / dotnet-build (Debug, windows-latest, net9.0) (push) Has been cancelled
dotnet-build-and-test / dotnet-build (Release, ubuntu-latest, net10.0) (push) Has been cancelled
dotnet-build-and-test / dotnet-build (Release, ubuntu-latest, net8.0) (push) Has been cancelled
dotnet-build-and-test / dotnet-build (Release, windows-latest, net472) (push) Has been cancelled
dotnet-build-and-test / dotnet-test (Release, integration, true, ubuntu-latest, net10.0) (push) Has been cancelled
dotnet-build-and-test / dotnet-test (Release, integration, true, windows-latest, net472) (push) Has been cancelled
dotnet-build-and-test / dotnet-foundry-hosted-it (push) Has been cancelled
dotnet-build-and-test / dotnet-build-and-test-check (push) Has been cancelled
dotnet-build-and-test / Integration Test Report (push) Has been cancelled
1827 lines
72 KiB
Python
1827 lines
72 KiB
Python
# Copyright (c) Microsoft. All rights reserved.
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
import json
|
|
import re
|
|
from typing import Any, cast
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
from agent_framework import Content, Message, SessionContext
|
|
from agent_framework._sessions import AgentSession
|
|
from azure.ai.contentunderstanding.models import AnalysisResult
|
|
|
|
from agent_framework_azure_contentunderstanding import (
|
|
ContentUnderstandingContextProvider,
|
|
DocumentStatus,
|
|
)
|
|
from agent_framework_azure_contentunderstanding._detection import SUPPORTED_MEDIA_TYPES, derive_doc_key
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_SAMPLE_PDF_BYTES = b"%PDF-1.4 fake content for testing"
|
|
|
|
|
|
def _make_mock_poller(result: AnalysisResult) -> AsyncMock:
|
|
"""Create a mock poller that returns the given result immediately."""
|
|
poller = AsyncMock()
|
|
poller.result = AsyncMock(return_value=result)
|
|
return poller
|
|
|
|
|
|
def _make_slow_poller(result: AnalysisResult, delay: float = 10.0) -> MagicMock:
|
|
"""Create a mock poller that simulates a timeout then eventually returns."""
|
|
poller = MagicMock()
|
|
|
|
async def slow_result() -> AnalysisResult:
|
|
await asyncio.sleep(delay)
|
|
return result
|
|
|
|
poller.result = slow_result
|
|
poller.continuation_token = MagicMock(return_value="mock_slow_continuation_token")
|
|
return poller
|
|
|
|
|
|
def _make_failing_poller(error: Exception) -> AsyncMock:
|
|
"""Create a mock poller that raises an exception."""
|
|
poller = AsyncMock()
|
|
poller.result = AsyncMock(side_effect=error)
|
|
return poller
|
|
|
|
|
|
def _make_data_uri(data: bytes, media_type: str) -> str:
|
|
return f"data:{media_type};base64,{base64.b64encode(data).decode('ascii')}"
|
|
|
|
|
|
def _make_content_from_data(data: bytes, media_type: str, filename: str | None = None) -> Content:
|
|
props = {"filename": filename} if filename else None
|
|
return Content.from_data(data, media_type, additional_properties=props)
|
|
|
|
|
|
def _make_context(messages: list[Message]) -> SessionContext:
|
|
return SessionContext(input_messages=messages)
|
|
|
|
|
|
def _make_provider(
|
|
mock_client: AsyncMock | None = None,
|
|
**kwargs: Any,
|
|
) -> ContentUnderstandingContextProvider:
|
|
provider = ContentUnderstandingContextProvider(
|
|
endpoint="https://test.cognitiveservices.azure.com/",
|
|
credential=AsyncMock(),
|
|
**kwargs,
|
|
)
|
|
if mock_client:
|
|
provider._client = mock_client # type: ignore[assignment]
|
|
return provider
|
|
|
|
|
|
def _make_mock_agent() -> MagicMock:
|
|
return MagicMock()
|
|
|
|
|
|
# ===========================================================================
|
|
# Test Classes
|
|
# ===========================================================================
|
|
|
|
|
|
class TestInit:
|
|
def test_default_values(self) -> None:
|
|
provider = ContentUnderstandingContextProvider(
|
|
endpoint="https://test.cognitiveservices.azure.com/",
|
|
credential=AsyncMock(),
|
|
)
|
|
assert provider.analyzer_id is None
|
|
assert provider.max_wait == 5.0
|
|
assert provider.output_sections == ["markdown", "fields"]
|
|
assert provider.source_id == "azure_contentunderstanding"
|
|
|
|
def test_custom_values(self) -> None:
|
|
provider = ContentUnderstandingContextProvider(
|
|
endpoint="https://custom.cognitiveservices.azure.com/",
|
|
credential=AsyncMock(),
|
|
analyzer_id="prebuilt-invoice",
|
|
max_wait=10.0,
|
|
output_sections=["markdown"],
|
|
source_id="custom_cu",
|
|
)
|
|
assert provider.analyzer_id == "prebuilt-invoice"
|
|
assert provider.max_wait == 10.0
|
|
assert provider.output_sections == ["markdown"]
|
|
assert provider.source_id == "custom_cu"
|
|
|
|
def test_max_wait_none(self) -> None:
|
|
provider = ContentUnderstandingContextProvider(
|
|
endpoint="https://test.cognitiveservices.azure.com/",
|
|
credential=AsyncMock(),
|
|
max_wait=None,
|
|
)
|
|
assert provider.max_wait is None
|
|
|
|
def test_endpoint_from_env_var(self, monkeypatch: Any) -> None:
|
|
"""Endpoint can be loaded from AZURE_CONTENTUNDERSTANDING_ENDPOINT env var."""
|
|
monkeypatch.setenv(
|
|
"AZURE_CONTENTUNDERSTANDING_ENDPOINT",
|
|
"https://env-test.cognitiveservices.azure.com/",
|
|
)
|
|
provider = ContentUnderstandingContextProvider(credential=AsyncMock())
|
|
assert provider._endpoint == "https://env-test.cognitiveservices.azure.com/"
|
|
|
|
def test_explicit_endpoint_overrides_env_var(self, monkeypatch: Any) -> None:
|
|
"""Explicit endpoint kwarg takes priority over env var."""
|
|
monkeypatch.setenv(
|
|
"AZURE_CONTENTUNDERSTANDING_ENDPOINT",
|
|
"https://env-test.cognitiveservices.azure.com/",
|
|
)
|
|
provider = ContentUnderstandingContextProvider(
|
|
endpoint="https://explicit.cognitiveservices.azure.com/",
|
|
credential=AsyncMock(),
|
|
)
|
|
assert provider._endpoint == "https://explicit.cognitiveservices.azure.com/"
|
|
|
|
def test_missing_endpoint_raises(self) -> None:
|
|
"""Missing endpoint (no kwarg, no env var) raises an error."""
|
|
# Clear env var to ensure load_settings raises
|
|
import os
|
|
|
|
import pytest as _pytest
|
|
from agent_framework.exceptions import SettingNotFoundError
|
|
|
|
env_key = "AZURE_CONTENTUNDERSTANDING_ENDPOINT"
|
|
old_val = os.environ.pop(env_key, None)
|
|
try:
|
|
with _pytest.raises(SettingNotFoundError, match="endpoint"):
|
|
ContentUnderstandingContextProvider(credential=AsyncMock())
|
|
finally:
|
|
if old_val is not None:
|
|
os.environ[env_key] = old_val
|
|
|
|
def test_missing_credential_raises(self) -> None:
|
|
"""Missing credential raises ValueError."""
|
|
import pytest as _pytest
|
|
|
|
with _pytest.raises(ValueError, match="credential is required"):
|
|
ContentUnderstandingContextProvider(
|
|
endpoint="https://test.cognitiveservices.azure.com/",
|
|
)
|
|
|
|
|
|
class TestAsyncContextManager:
|
|
async def test_aenter_returns_self(self) -> None:
|
|
provider = ContentUnderstandingContextProvider(
|
|
endpoint="https://test.cognitiveservices.azure.com/",
|
|
credential=AsyncMock(),
|
|
)
|
|
result = await provider.__aenter__()
|
|
assert result is provider
|
|
await provider.__aexit__(None, None, None)
|
|
|
|
async def test_aexit_closes_client(self) -> None:
|
|
provider = ContentUnderstandingContextProvider(
|
|
endpoint="https://test.cognitiveservices.azure.com/",
|
|
credential=AsyncMock(),
|
|
)
|
|
mock_client = AsyncMock()
|
|
provider._client = mock_client # type: ignore[assignment]
|
|
await provider.__aexit__(None, None, None)
|
|
mock_client.close.assert_called_once()
|
|
|
|
|
|
class TestBeforeRunNewFile:
|
|
async def test_single_pdf_analyzed(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
pdf_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
mock_cu_client.begin_analyze_binary = AsyncMock(return_value=_make_mock_poller(pdf_analysis_result))
|
|
provider = _make_provider(mock_client=mock_cu_client)
|
|
|
|
msg = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("What's on this invoice?"),
|
|
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "invoice.pdf"),
|
|
],
|
|
)
|
|
context = _make_context([msg])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
|
|
|
|
# Document should be in state
|
|
assert "documents" in state
|
|
assert "invoice.pdf" in state["documents"]
|
|
assert state["documents"]["invoice.pdf"]["status"] == DocumentStatus.READY
|
|
|
|
# Binary should be stripped from input
|
|
for m in context.input_messages:
|
|
for c in m.contents:
|
|
assert c.media_type != "application/pdf"
|
|
|
|
# Context should have messages injected
|
|
assert len(context.context_messages) > 0
|
|
|
|
async def test_url_input_analyzed(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
pdf_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
mock_cu_client.begin_analyze = AsyncMock(return_value=_make_mock_poller(pdf_analysis_result))
|
|
provider = _make_provider(mock_client=mock_cu_client)
|
|
|
|
msg = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("Analyze this document"),
|
|
Content.from_uri("https://example.com/report.pdf", media_type="application/pdf"),
|
|
],
|
|
)
|
|
context = _make_context([msg])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
|
|
|
|
# URL input should use begin_analyze
|
|
mock_cu_client.begin_analyze.assert_called_once()
|
|
assert "report.pdf" in state["documents"]
|
|
assert state["documents"]["report.pdf"]["status"] == DocumentStatus.READY
|
|
|
|
async def test_text_only_skipped(self, mock_cu_client: AsyncMock) -> None:
|
|
provider = _make_provider(mock_client=mock_cu_client)
|
|
|
|
msg = Message(role="user", contents=[Content.from_text("What's the weather?")])
|
|
context = _make_context([msg])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
|
|
|
|
# No CU calls
|
|
mock_cu_client.begin_analyze.assert_not_called()
|
|
mock_cu_client.begin_analyze_binary.assert_not_called()
|
|
# No documents
|
|
assert state.get("documents", {}) == {}
|
|
|
|
|
|
class TestBeforeRunMultiFile:
|
|
async def test_two_files_both_analyzed(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
pdf_analysis_result: AnalysisResult,
|
|
image_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
mock_cu_client.begin_analyze_binary = AsyncMock(
|
|
side_effect=[
|
|
_make_mock_poller(pdf_analysis_result),
|
|
_make_mock_poller(image_analysis_result),
|
|
]
|
|
)
|
|
provider = _make_provider(mock_client=mock_cu_client)
|
|
|
|
msg = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("Compare these documents"),
|
|
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "doc1.pdf"),
|
|
_make_content_from_data(b"\x89PNG fake", "image/png", "chart.png"),
|
|
],
|
|
)
|
|
context = _make_context([msg])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
|
|
|
|
assert len(state["documents"]) == 2
|
|
assert state["documents"]["doc1.pdf"]["status"] == DocumentStatus.READY
|
|
assert state["documents"]["chart.png"]["status"] == DocumentStatus.READY
|
|
|
|
|
|
class TestBeforeRunTimeout:
|
|
async def test_exceeds_max_wait_defers_to_background(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
pdf_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
mock_cu_client.begin_analyze_binary = AsyncMock(return_value=_make_slow_poller(pdf_analysis_result, delay=10.0))
|
|
provider = _make_provider(mock_client=mock_cu_client, max_wait=0.1)
|
|
|
|
msg = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("Analyze this"),
|
|
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "big_doc.pdf"),
|
|
],
|
|
)
|
|
context = _make_context([msg])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
|
|
|
|
assert state["documents"]["big_doc.pdf"]["status"] == DocumentStatus.ANALYZING
|
|
assert "big_doc.pdf" in state.get("_pending_tokens", {})
|
|
token_info = state["_pending_tokens"]["big_doc.pdf"]
|
|
assert "continuation_token" in token_info
|
|
assert "analyzer_id" in token_info
|
|
|
|
# Context messages should mention analyzing
|
|
assert any("being analyzed" in m.text for msgs in context.context_messages.values() for m in msgs)
|
|
|
|
|
|
class TestBeforeRunPendingResolution:
|
|
async def test_pending_completes_on_next_turn(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
pdf_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
# Mock begin_analyze to return a completed poller when called with continuation_token
|
|
mock_poller = _make_mock_poller(pdf_analysis_result)
|
|
mock_poller.done = MagicMock(return_value=True)
|
|
mock_cu_client.begin_analyze = AsyncMock(return_value=mock_poller)
|
|
provider = _make_provider(mock_client=mock_cu_client)
|
|
|
|
state: dict[str, Any] = {
|
|
"_pending_tokens": {
|
|
"report.pdf": {"continuation_token": "tok_123", "analyzer_id": "prebuilt-documentSearch"}
|
|
},
|
|
"documents": {
|
|
"report.pdf": {
|
|
"status": DocumentStatus.ANALYZING,
|
|
"filename": "report.pdf",
|
|
"media_type": "application/pdf",
|
|
"analyzer_id": "prebuilt-documentSearch",
|
|
"analyzed_at": None,
|
|
"analysis_duration_s": None,
|
|
"upload_duration_s": None,
|
|
"result": None,
|
|
"error": None,
|
|
},
|
|
},
|
|
}
|
|
|
|
msg = Message(role="user", contents=[Content.from_text("Is the report ready?")])
|
|
context = _make_context([msg])
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
|
|
|
|
assert state["documents"]["report.pdf"]["status"] == DocumentStatus.READY
|
|
assert state["documents"]["report.pdf"]["result"] is not None
|
|
assert "report.pdf" not in state.get("_pending_tokens", {})
|
|
|
|
|
|
class TestBeforeRunPendingFailure:
|
|
async def test_pending_task_failure_updates_state(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
) -> None:
|
|
# Mock begin_analyze to raise when resuming from continuation token
|
|
mock_cu_client.begin_analyze = AsyncMock(side_effect=RuntimeError("CU service unavailable"))
|
|
provider = _make_provider(mock_client=mock_cu_client)
|
|
|
|
state: dict[str, Any] = {
|
|
"_pending_tokens": {
|
|
"bad_doc.pdf": {"continuation_token": "tok_fail", "analyzer_id": "prebuilt-documentSearch"}
|
|
},
|
|
"documents": {
|
|
"bad_doc.pdf": {
|
|
"status": DocumentStatus.ANALYZING,
|
|
"filename": "bad_doc.pdf",
|
|
"media_type": "application/pdf",
|
|
"analyzer_id": "prebuilt-documentSearch",
|
|
"analyzed_at": None,
|
|
"analysis_duration_s": None,
|
|
"upload_duration_s": None,
|
|
"result": None,
|
|
"error": None,
|
|
},
|
|
},
|
|
}
|
|
|
|
msg = Message(role="user", contents=[Content.from_text("Check status")])
|
|
context = _make_context([msg])
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
|
|
|
|
assert state["documents"]["bad_doc.pdf"]["status"] == DocumentStatus.FAILED
|
|
assert "CU service unavailable" in (state["documents"]["bad_doc.pdf"]["error"] or "")
|
|
|
|
|
|
class TestDocumentKeyDerivation:
|
|
def test_filename_from_additional_properties(self) -> None:
|
|
content = _make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "my_report.pdf")
|
|
key = derive_doc_key(content)
|
|
assert key == "my_report.pdf"
|
|
|
|
def test_url_basename(self) -> None:
|
|
content = Content.from_uri("https://example.com/docs/annual_report.pdf", media_type="application/pdf")
|
|
key = derive_doc_key(content)
|
|
assert key == "annual_report.pdf"
|
|
|
|
def test_content_hash_fallback(self) -> None:
|
|
content = Content.from_data(_SAMPLE_PDF_BYTES, "application/pdf")
|
|
key = derive_doc_key(content)
|
|
assert key.startswith("doc_")
|
|
assert len(key) == 12 # "doc_" + 8 hex chars
|
|
|
|
|
|
class TestSessionState:
|
|
async def test_documents_persist_across_turns(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
pdf_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
mock_cu_client.begin_analyze_binary = AsyncMock(return_value=_make_mock_poller(pdf_analysis_result))
|
|
provider = _make_provider(mock_client=mock_cu_client)
|
|
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
# Turn 1: upload
|
|
msg1 = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("Analyze this"),
|
|
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "doc.pdf"),
|
|
],
|
|
)
|
|
ctx1 = _make_context([msg1])
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=ctx1, state=state)
|
|
|
|
assert "doc.pdf" in state["documents"]
|
|
|
|
# Turn 2: follow-up (no file)
|
|
msg2 = Message(role="user", contents=[Content.from_text("What's the total?")])
|
|
ctx2 = _make_context([msg2])
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=ctx2, state=state)
|
|
|
|
# Document should still be there
|
|
assert "doc.pdf" in state["documents"]
|
|
assert state["documents"]["doc.pdf"]["status"] == DocumentStatus.READY
|
|
|
|
|
|
class TestListDocumentsTool:
|
|
async def test_returns_all_docs_with_status(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
pdf_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
mock_cu_client.begin_analyze_binary = AsyncMock(return_value=_make_mock_poller(pdf_analysis_result))
|
|
provider = _make_provider(mock_client=mock_cu_client)
|
|
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
msg = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("Analyze this"),
|
|
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "test.pdf"),
|
|
],
|
|
)
|
|
context = _make_context([msg])
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
|
|
|
|
# Find the list_documents tool
|
|
list_tool = None
|
|
for tool in context.tools:
|
|
if getattr(tool, "name", None) == "list_documents":
|
|
list_tool = tool
|
|
break
|
|
|
|
assert list_tool is not None
|
|
result = list_tool.func() # type: ignore[union-attr]
|
|
parsed = json.loads(result)
|
|
assert len(parsed) == 1
|
|
assert parsed[0]["name"] == "test.pdf"
|
|
assert parsed[0]["status"] == DocumentStatus.READY
|
|
|
|
|
|
class TestOutputFiltering:
|
|
"""Validate that output_sections controls what `_render_for_llm` emits.
|
|
|
|
Rendering is delegated to ``azure.ai.contentunderstanding.to_llm_input``:
|
|
- ``"markdown" in output_sections`` -> ``include_markdown=True``.
|
|
- ``"fields" in output_sections`` -> ``include_fields=True``.
|
|
- ``metadata={"source": <filename>}`` is always supplied.
|
|
|
|
Note: detailed field/JSON shape is owned by the SDK and exercised in the
|
|
SDK's own ``to_llm_input`` tests. We only assert MAF-level wiring here.
|
|
"""
|
|
|
|
def test_default_markdown_and_fields(self, pdf_analysis_result: AnalysisResult) -> None:
|
|
provider = _make_provider()
|
|
rendered = provider._render_for_llm(pdf_analysis_result, "report.pdf")
|
|
|
|
# YAML front matter with source key.
|
|
assert "source: report.pdf" in rendered
|
|
# PDF fixture contains "Contoso" in its markdown body.
|
|
assert "Contoso" in rendered
|
|
|
|
def test_markdown_only(self, pdf_analysis_result: AnalysisResult) -> None:
|
|
provider = _make_provider(output_sections=["markdown"])
|
|
rendered = provider._render_for_llm(pdf_analysis_result, "report.pdf")
|
|
|
|
# Markdown body still present; no ``fields:`` front-matter section.
|
|
assert "Contoso" in rendered
|
|
assert "\nfields:" not in rendered
|
|
assert not rendered.startswith("fields:")
|
|
|
|
def test_fields_only(self, invoice_analysis_result: AnalysisResult) -> None:
|
|
provider = _make_provider(output_sections=["fields"])
|
|
rendered = provider._render_for_llm(invoice_analysis_result, "invoice.pdf")
|
|
|
|
# ``fields:`` YAML key is emitted; vendor name appears under it.
|
|
assert "fields:" in rendered
|
|
assert "VendorName" in rendered
|
|
assert "TechServe Global Partners" in rendered
|
|
|
|
def test_field_values_extracted(self, invoice_analysis_result: AnalysisResult) -> None:
|
|
provider = _make_provider()
|
|
rendered = provider._render_for_llm(invoice_analysis_result, "invoice.pdf")
|
|
|
|
# Both sections present.
|
|
assert "fields:" in rendered
|
|
# Field values visible to the LLM (vendor + a known line-item description).
|
|
assert "TechServe Global Partners" in rendered
|
|
assert "Consulting Services" in rendered
|
|
|
|
def test_source_metadata_uses_filename(self, pdf_analysis_result: AnalysisResult) -> None:
|
|
"""Per-document ``source`` key carries the original filename."""
|
|
provider = _make_provider()
|
|
rendered = provider._render_for_llm(pdf_analysis_result, "custom_name.pdf")
|
|
assert "source: custom_name.pdf" in rendered
|
|
|
|
def test_page_markers_passed_through_to_llm_input(self, pdf_analysis_result: AnalysisResult) -> None:
|
|
"""Decision H: MAF must not strip page markers emitted by the SDK helper.
|
|
|
|
Today the SDK helper (``azure.ai.contentunderstanding.to_llm_input``)
|
|
injects ``<!-- page N -->`` markers per page. Per
|
|
``cognitive-services/ContentUnderstanding-Docs#249`` (Decision 4) it
|
|
will switch to ``<!-- InputPageNumber: N -->`` once the service ships
|
|
the marker natively. Either format must reach the LLM unchanged --
|
|
this test guards against MAF accidentally regex-stripping them.
|
|
"""
|
|
provider = _make_provider()
|
|
rendered = provider._render_for_llm(pdf_analysis_result, "report.pdf")
|
|
|
|
legacy = re.findall(r"<!--\s*page\s+\d+\s*-->", rendered)
|
|
future = re.findall(r"<!--\s*InputPageNumber:\s*\d+\s*-->", rendered)
|
|
# PDF fixture has 5 pages; expect 5 markers in whichever format is in use.
|
|
assert len(legacy) == 5 or len(future) == 5, (
|
|
"Expected SDK-injected page markers to be passed through to LLM input. "
|
|
f"Found legacy={len(legacy)}, future={len(future)}."
|
|
)
|
|
|
|
|
|
class TestDuplicateDocumentKey:
|
|
async def test_duplicate_filename_rejected(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
pdf_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
"""Uploading the same filename twice in the same session should reject the second."""
|
|
mock_cu_client.begin_analyze_binary = AsyncMock(return_value=_make_mock_poller(pdf_analysis_result))
|
|
provider = _make_provider(mock_client=mock_cu_client)
|
|
|
|
# Turn 1: upload invoice.pdf
|
|
msg1 = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("Analyze this"),
|
|
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "invoice.pdf"),
|
|
],
|
|
)
|
|
context1 = _make_context([msg1])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context1, state=state)
|
|
assert "invoice.pdf" in state["documents"]
|
|
assert state["documents"]["invoice.pdf"]["status"] == DocumentStatus.READY
|
|
|
|
# Turn 2: upload invoice.pdf again (different content but same filename)
|
|
msg2 = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("Analyze this too"),
|
|
_make_content_from_data(b"different-content", "application/pdf", "invoice.pdf"),
|
|
],
|
|
)
|
|
context2 = _make_context([msg2])
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context2, state=state)
|
|
|
|
# Should still have only one document, not re-analyzed
|
|
assert mock_cu_client.begin_analyze_binary.call_count == 1
|
|
# Context messages should mention duplicate
|
|
assert any("already uploaded" in m.text for msgs in context2.context_messages.values() for m in msgs)
|
|
|
|
async def test_duplicate_in_same_turn_rejected(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
pdf_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
"""Two files with the same filename in the same turn: first wins, second rejected."""
|
|
mock_cu_client.begin_analyze_binary = AsyncMock(return_value=_make_mock_poller(pdf_analysis_result))
|
|
provider = _make_provider(mock_client=mock_cu_client)
|
|
|
|
msg = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("Analyze both"),
|
|
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "report.pdf"),
|
|
_make_content_from_data(b"other-content", "application/pdf", "report.pdf"),
|
|
],
|
|
)
|
|
context = _make_context([msg])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
|
|
|
|
# Only analyzed once (first one wins)
|
|
assert mock_cu_client.begin_analyze_binary.call_count == 1
|
|
assert "report.pdf" in state["documents"]
|
|
assert any("already uploaded" in m.text for msgs in context.context_messages.values() for m in msgs)
|
|
|
|
|
|
class TestBinaryStripping:
|
|
async def test_supported_files_stripped(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
pdf_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
mock_cu_client.begin_analyze_binary = AsyncMock(return_value=_make_mock_poller(pdf_analysis_result))
|
|
provider = _make_provider(mock_client=mock_cu_client)
|
|
|
|
msg = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("What's in here?"),
|
|
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "doc.pdf"),
|
|
],
|
|
)
|
|
context = _make_context([msg])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
|
|
|
|
# PDF should be stripped; text should remain
|
|
for m in context.input_messages:
|
|
for c in m.contents:
|
|
assert c.media_type != "application/pdf"
|
|
assert any(c.text and "What's in here?" in c.text for c in m.contents)
|
|
|
|
async def test_unsupported_files_left_in_place(self, mock_cu_client: AsyncMock) -> None:
|
|
provider = _make_provider(mock_client=mock_cu_client)
|
|
|
|
msg = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("What's in this zip?"),
|
|
Content.from_data(
|
|
b"PK\x03\x04fake",
|
|
"application/zip",
|
|
additional_properties={"filename": "archive.zip"},
|
|
),
|
|
],
|
|
)
|
|
context = _make_context([msg])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
|
|
|
|
# Zip should NOT be stripped (unsupported)
|
|
found_zip = False
|
|
for m in context.input_messages:
|
|
for c in m.contents:
|
|
if c.media_type == "application/zip":
|
|
found_zip = True
|
|
assert found_zip
|
|
|
|
|
|
# Real magic-byte headers for binary sniffing tests
|
|
_MP4_MAGIC = b"\x00\x00\x00\x1cftypisom" + b"\x00" * 250
|
|
_WAV_MAGIC = b"RIFF\x00\x00\x00\x00WAVE" + b"\x00" * 250
|
|
_MP3_MAGIC = b"ID3\x04\x00\x00" + b"\x00" * 250
|
|
_FLAC_MAGIC = b"fLaC\x00\x00\x00\x00" + b"\x00" * 250
|
|
_OGG_MAGIC = b"OggS\x00\x02" + b"\x00" * 250
|
|
_AVI_MAGIC = b"RIFF\x00\x00\x00\x00AVI " + b"\x00" * 250
|
|
_MOV_MAGIC = b"\x00\x00\x00\x14ftypqt " + b"\x00" * 250
|
|
|
|
|
|
class TestMimeSniffing:
|
|
"""Tests for binary MIME sniffing via filetype when upstream MIME is unreliable."""
|
|
|
|
async def test_octet_stream_mp4_detected_and_stripped(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
pdf_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
"""MP4 uploaded as application/octet-stream should be sniffed, corrected, and stripped."""
|
|
mock_cu_client.begin_analyze_binary = AsyncMock(return_value=_make_mock_poller(pdf_analysis_result))
|
|
provider = _make_provider(mock_client=mock_cu_client)
|
|
|
|
msg = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("What's in this file?"),
|
|
_make_content_from_data(_MP4_MAGIC, "application/octet-stream", "video.mp4"),
|
|
],
|
|
)
|
|
context = _make_context([msg])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
|
|
|
|
# MP4 should be stripped from input
|
|
for m in context.input_messages:
|
|
for c in m.contents:
|
|
assert c.media_type != "application/octet-stream", "octet-stream content should be stripped"
|
|
|
|
# CU should have been called
|
|
assert mock_cu_client.begin_analyze_binary.called
|
|
|
|
async def test_octet_stream_wav_detected_via_sniff(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
pdf_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
"""WAV uploaded as application/octet-stream should be detected via filetype sniffing."""
|
|
mock_cu_client.begin_analyze_binary = AsyncMock(return_value=_make_mock_poller(pdf_analysis_result))
|
|
provider = _make_provider(mock_client=mock_cu_client)
|
|
|
|
msg = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("Transcribe"),
|
|
_make_content_from_data(_WAV_MAGIC, "application/octet-stream", "audio.wav"),
|
|
],
|
|
)
|
|
context = _make_context([msg])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
|
|
|
|
# Should be detected and analyzed
|
|
assert "audio.wav" in state["documents"]
|
|
# The media_type should be corrected to audio/wav (via _MIME_ALIASES)
|
|
assert state["documents"]["audio.wav"]["media_type"] == "audio/wav"
|
|
|
|
async def test_octet_stream_mp3_detected_via_sniff(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
pdf_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
"""MP3 uploaded as application/octet-stream should be detected as audio/mpeg."""
|
|
mock_cu_client.begin_analyze_binary = AsyncMock(return_value=_make_mock_poller(pdf_analysis_result))
|
|
provider = _make_provider(mock_client=mock_cu_client)
|
|
|
|
msg = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("Transcribe"),
|
|
_make_content_from_data(_MP3_MAGIC, "application/octet-stream", "song.mp3"),
|
|
],
|
|
)
|
|
context = _make_context([msg])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
|
|
|
|
assert "song.mp3" in state["documents"]
|
|
assert state["documents"]["song.mp3"]["media_type"] == "audio/mpeg"
|
|
|
|
async def test_octet_stream_flac_alias_normalized(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
pdf_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
"""FLAC sniffed as audio/x-flac should be normalized to audio/flac."""
|
|
mock_cu_client.begin_analyze_binary = AsyncMock(return_value=_make_mock_poller(pdf_analysis_result))
|
|
provider = _make_provider(mock_client=mock_cu_client)
|
|
|
|
msg = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("Transcribe"),
|
|
_make_content_from_data(_FLAC_MAGIC, "application/octet-stream", "music.flac"),
|
|
],
|
|
)
|
|
context = _make_context([msg])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
|
|
|
|
assert "music.flac" in state["documents"]
|
|
assert state["documents"]["music.flac"]["media_type"] == "audio/flac"
|
|
|
|
async def test_octet_stream_unknown_binary_not_stripped(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
) -> None:
|
|
"""Unknown binary with application/octet-stream should NOT be stripped."""
|
|
provider = _make_provider(mock_client=mock_cu_client)
|
|
|
|
unknown_bytes = b"\x00\x01\x02\x03random garbage" + b"\x00" * 250
|
|
msg = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("What is this?"),
|
|
_make_content_from_data(unknown_bytes, "application/octet-stream", "mystery.bin"),
|
|
],
|
|
)
|
|
context = _make_context([msg])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
|
|
|
|
# Unknown file should NOT be stripped
|
|
found_octet = False
|
|
for m in context.input_messages:
|
|
for c in m.contents:
|
|
if c.media_type == "application/octet-stream":
|
|
found_octet = True
|
|
assert found_octet
|
|
|
|
async def test_missing_mime_falls_back_to_filename(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
pdf_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
"""Content with empty MIME but a .mp4 filename should be detected via mimetypes fallback."""
|
|
mock_cu_client.begin_analyze_binary = AsyncMock(return_value=_make_mock_poller(pdf_analysis_result))
|
|
provider = _make_provider(mock_client=mock_cu_client)
|
|
|
|
# Use garbage binary (filetype won't detect) but filename has .mp4
|
|
garbage = b"\x00" * 300
|
|
content = Content.from_data(garbage, "", additional_properties={"filename": "recording.mp4"})
|
|
msg = Message(
|
|
role="user",
|
|
contents=[Content.from_text("Analyze"), content],
|
|
)
|
|
context = _make_context([msg])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
|
|
|
|
# Should be detected via filename and analyzed
|
|
assert "recording.mp4" in state["documents"]
|
|
|
|
async def test_correct_mime_not_sniffed(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
pdf_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
"""Files with correct MIME type should go through fast path without sniffing."""
|
|
mock_cu_client.begin_analyze_binary = AsyncMock(return_value=_make_mock_poller(pdf_analysis_result))
|
|
provider = _make_provider(mock_client=mock_cu_client)
|
|
|
|
msg = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("Analyze"),
|
|
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "doc.pdf"),
|
|
],
|
|
)
|
|
context = _make_context([msg])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
|
|
|
|
assert "doc.pdf" in state["documents"]
|
|
assert state["documents"]["doc.pdf"]["media_type"] == "application/pdf"
|
|
|
|
async def test_sniffed_video_uses_correct_analyzer(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
pdf_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
"""MP4 sniffed from octet-stream should use prebuilt-videoSearch analyzer."""
|
|
mock_cu_client.begin_analyze_binary = AsyncMock(return_value=_make_mock_poller(pdf_analysis_result))
|
|
provider = _make_provider(mock_client=mock_cu_client) # analyzer_id=None → auto-detect
|
|
|
|
msg = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("What's in this video?"),
|
|
_make_content_from_data(_MP4_MAGIC, "application/octet-stream", "demo.mp4"),
|
|
],
|
|
)
|
|
context = _make_context([msg])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
|
|
|
|
assert state["documents"]["demo.mp4"]["analyzer_id"] == "prebuilt-videoSearch"
|
|
|
|
|
|
class TestErrorHandling:
|
|
async def test_cu_service_error(self, mock_cu_client: AsyncMock) -> None:
|
|
mock_cu_client.begin_analyze_binary = AsyncMock(
|
|
return_value=_make_failing_poller(RuntimeError("Service unavailable"))
|
|
)
|
|
provider = _make_provider(mock_client=mock_cu_client)
|
|
|
|
msg = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("Analyze this"),
|
|
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "error.pdf"),
|
|
],
|
|
)
|
|
context = _make_context([msg])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
|
|
|
|
assert state["documents"]["error.pdf"]["status"] == DocumentStatus.FAILED
|
|
assert "Service unavailable" in (state["documents"]["error.pdf"]["error"] or "")
|
|
|
|
async def test_lazy_initialization_on_before_run(self) -> None:
|
|
"""before_run works with eagerly-initialized client."""
|
|
provider = ContentUnderstandingContextProvider(
|
|
endpoint="https://test.cognitiveservices.azure.com/",
|
|
credential=AsyncMock(),
|
|
)
|
|
assert provider._client is not None
|
|
|
|
mock_client = AsyncMock()
|
|
mock_client.begin_analyze_binary = AsyncMock(
|
|
side_effect=Exception("mock error"),
|
|
)
|
|
provider._client = mock_client # type: ignore[assignment]
|
|
|
|
msg = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("Analyze this"),
|
|
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "doc.pdf"),
|
|
],
|
|
)
|
|
context = _make_context([msg])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
|
|
# Client should still be set
|
|
assert provider._client is not None
|
|
|
|
|
|
class TestMultiModalFixtures:
|
|
"""Verify ``_render_for_llm`` produces sensible output for each modality.
|
|
|
|
Detailed shape of the YAML/Markdown payload is the SDK's responsibility and
|
|
is exercised by ``azure-ai-contentunderstanding`` tests. Here we only check
|
|
that the MAF wiring (filename surfaced as ``source``, key content visible)
|
|
works for each fixture kind.
|
|
"""
|
|
|
|
def test_pdf_fixture_loads(self, pdf_analysis_result: AnalysisResult) -> None:
|
|
provider = _make_provider()
|
|
rendered = provider._render_for_llm(pdf_analysis_result, "report.pdf")
|
|
assert "source: report.pdf" in rendered
|
|
assert "Contoso" in rendered
|
|
|
|
def test_audio_fixture_loads(self, audio_analysis_result: AnalysisResult) -> None:
|
|
provider = _make_provider()
|
|
rendered = provider._render_for_llm(audio_analysis_result, "call.mp3")
|
|
assert "source: call.mp3" in rendered
|
|
assert "Call Center" in rendered
|
|
|
|
def test_video_fixture_loads(self, video_analysis_result: AnalysisResult) -> None:
|
|
provider = _make_provider()
|
|
rendered = provider._render_for_llm(video_analysis_result, "demo.mp4")
|
|
assert "source: demo.mp4" in rendered
|
|
# All 3 segments should be visible in the rendered text.
|
|
assert "Contoso Product Demo" in rendered
|
|
assert "real-time monitoring" in rendered
|
|
assert "contoso.com/cloud-manager" in rendered
|
|
# Each segment must render its own YAML front matter with a timeRange entry.
|
|
# This guards against multi-segment results being collapsed into one block.
|
|
assert rendered.count("timeRange:") == 3
|
|
# Segments must be rendered in chronological order (1s, 15s, 36s starts).
|
|
assert (
|
|
rendered.index("Contoso Product Demo")
|
|
< rendered.index("real-time monitoring")
|
|
< rendered.index("contoso.com/cloud-manager")
|
|
)
|
|
|
|
def test_image_fixture_loads(self, image_analysis_result: AnalysisResult) -> None:
|
|
provider = _make_provider()
|
|
rendered = provider._render_for_llm(image_analysis_result, "image.png")
|
|
assert "source: image.png" in rendered
|
|
# Non-empty body (image markdown caption from CU).
|
|
assert len(rendered) > len("source: image.png")
|
|
|
|
def test_invoice_fixture_loads(self, invoice_analysis_result: AnalysisResult) -> None:
|
|
provider = _make_provider()
|
|
rendered = provider._render_for_llm(invoice_analysis_result, "invoice.pdf")
|
|
assert "source: invoice.pdf" in rendered
|
|
assert "fields:" in rendered
|
|
assert "VendorName" in rendered
|
|
|
|
|
|
# NOTE: ``TestFormatResult`` (4 tests) was deleted as part of the migration to
|
|
# ``azure.ai.contentunderstanding.to_llm_input``. The legacy ``format_result``
|
|
# helper no longer exists; rendering shape (YAML front matter + Markdown body,
|
|
# segment serialization, reserved-key handling) is owned and tested by the SDK.
|
|
|
|
|
|
class TestSupportedMediaTypes:
|
|
def test_pdf_supported(self) -> None:
|
|
assert "application/pdf" in SUPPORTED_MEDIA_TYPES
|
|
|
|
def test_audio_supported(self) -> None:
|
|
assert "audio/mp3" in SUPPORTED_MEDIA_TYPES
|
|
assert "audio/wav" in SUPPORTED_MEDIA_TYPES
|
|
|
|
def test_video_supported(self) -> None:
|
|
assert "video/mp4" in SUPPORTED_MEDIA_TYPES
|
|
|
|
def test_zip_not_supported(self) -> None:
|
|
assert "application/zip" not in SUPPORTED_MEDIA_TYPES
|
|
|
|
|
|
class TestAnalyzerAutoDetection:
|
|
"""Verify _resolve_analyzer_id auto-selects the right analyzer by media type."""
|
|
|
|
def test_explicit_analyzer_always_wins(self) -> None:
|
|
provider = _make_provider(analyzer_id="prebuilt-invoice")
|
|
assert provider._resolve_analyzer_id("audio/mp3") == "prebuilt-invoice"
|
|
assert provider._resolve_analyzer_id("video/mp4") == "prebuilt-invoice"
|
|
assert provider._resolve_analyzer_id("application/pdf") == "prebuilt-invoice"
|
|
|
|
def test_auto_detect_pdf(self) -> None:
|
|
provider = _make_provider() # analyzer_id=None
|
|
assert provider._resolve_analyzer_id("application/pdf") == "prebuilt-documentSearch"
|
|
|
|
def test_auto_detect_image(self) -> None:
|
|
provider = _make_provider()
|
|
assert provider._resolve_analyzer_id("image/jpeg") == "prebuilt-documentSearch"
|
|
assert provider._resolve_analyzer_id("image/png") == "prebuilt-documentSearch"
|
|
|
|
def test_auto_detect_audio(self) -> None:
|
|
provider = _make_provider()
|
|
assert provider._resolve_analyzer_id("audio/mp3") == "prebuilt-audioSearch"
|
|
assert provider._resolve_analyzer_id("audio/wav") == "prebuilt-audioSearch"
|
|
assert provider._resolve_analyzer_id("audio/mpeg") == "prebuilt-audioSearch"
|
|
|
|
def test_auto_detect_video(self) -> None:
|
|
provider = _make_provider()
|
|
assert provider._resolve_analyzer_id("video/mp4") == "prebuilt-videoSearch"
|
|
assert provider._resolve_analyzer_id("video/webm") == "prebuilt-videoSearch"
|
|
|
|
def test_auto_detect_unknown_falls_back_to_document(self) -> None:
|
|
provider = _make_provider()
|
|
assert provider._resolve_analyzer_id("application/octet-stream") == "prebuilt-documentSearch"
|
|
|
|
|
|
class TestFileSearchIntegration:
|
|
_FILE_SEARCH_TOOL = {"type": "file_search", "vector_store_ids": ["vs_test123"]}
|
|
|
|
def _make_mock_backend(self) -> AsyncMock:
|
|
"""Create a mock FileSearchBackend."""
|
|
backend = AsyncMock()
|
|
backend.upload_file = AsyncMock(return_value="file_test456")
|
|
backend.delete_file = AsyncMock()
|
|
return backend
|
|
|
|
def _make_file_search_config(self, backend: AsyncMock | None = None) -> Any:
|
|
from agent_framework_azure_contentunderstanding import FileSearchConfig
|
|
|
|
return FileSearchConfig(
|
|
backend=backend or self._make_mock_backend(),
|
|
vector_store_id="vs_test123",
|
|
file_search_tool=self._FILE_SEARCH_TOOL,
|
|
)
|
|
|
|
async def test_file_search_uploads_to_vector_store(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
pdf_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
mock_backend = self._make_mock_backend()
|
|
config = self._make_file_search_config(mock_backend)
|
|
mock_cu_client.begin_analyze_binary = AsyncMock(
|
|
return_value=_make_mock_poller(pdf_analysis_result),
|
|
)
|
|
provider = _make_provider(
|
|
mock_client=mock_cu_client,
|
|
file_search=config,
|
|
)
|
|
|
|
msg = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("Analyze this"),
|
|
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "doc.pdf"),
|
|
],
|
|
)
|
|
context = _make_context([msg])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(
|
|
agent=_make_mock_agent(),
|
|
session=session,
|
|
context=context,
|
|
state=state,
|
|
)
|
|
|
|
# File should be uploaded via backend
|
|
mock_backend.upload_file.assert_called_once()
|
|
call_args = mock_backend.upload_file.call_args
|
|
assert call_args[0][0] == "vs_test123" # vector_store_id
|
|
assert call_args[0][1] == "doc.pdf.md" # filename
|
|
# file_search tool should be registered on context
|
|
assert self._FILE_SEARCH_TOOL in context.tools
|
|
|
|
async def test_file_search_no_content_injection(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
pdf_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
"""When file_search is enabled, full content should NOT be injected into context."""
|
|
mock_cu_client.begin_analyze_binary = AsyncMock(
|
|
return_value=_make_mock_poller(pdf_analysis_result),
|
|
)
|
|
provider = _make_provider(
|
|
mock_client=mock_cu_client,
|
|
file_search=self._make_file_search_config(),
|
|
)
|
|
|
|
msg = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("Analyze this"),
|
|
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "doc.pdf"),
|
|
],
|
|
)
|
|
context = _make_context([msg])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(
|
|
agent=_make_mock_agent(),
|
|
session=session,
|
|
context=context,
|
|
state=state,
|
|
)
|
|
|
|
# Context messages should NOT contain full document content
|
|
# (file_search handles retrieval instead)
|
|
for msgs in context.context_messages.values():
|
|
for m in msgs:
|
|
assert "Document Content" not in m.text
|
|
|
|
async def test_cleanup_deletes_uploaded_files(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
pdf_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
mock_backend = self._make_mock_backend()
|
|
config = self._make_file_search_config(mock_backend)
|
|
mock_cu_client.begin_analyze_binary = AsyncMock(
|
|
return_value=_make_mock_poller(pdf_analysis_result),
|
|
)
|
|
provider = _make_provider(
|
|
mock_client=mock_cu_client,
|
|
file_search=config,
|
|
)
|
|
|
|
msg = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("Analyze this"),
|
|
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "doc.pdf"),
|
|
],
|
|
)
|
|
context = _make_context([msg])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(
|
|
agent=_make_mock_agent(),
|
|
session=session,
|
|
context=context,
|
|
state=state,
|
|
)
|
|
|
|
# Close should clean up uploaded files (not the vector store itself)
|
|
await provider.close()
|
|
mock_backend.delete_file.assert_called_once_with("file_test456")
|
|
|
|
async def test_no_file_search_injects_content(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
pdf_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
"""Without file_search, full content should be injected (default behavior)."""
|
|
mock_cu_client.begin_analyze_binary = AsyncMock(
|
|
return_value=_make_mock_poller(pdf_analysis_result),
|
|
)
|
|
provider = _make_provider(mock_client=mock_cu_client)
|
|
|
|
msg = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("Analyze this"),
|
|
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "doc.pdf"),
|
|
],
|
|
)
|
|
context = _make_context([msg])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(
|
|
agent=_make_mock_agent(),
|
|
session=session,
|
|
context=context,
|
|
state=state,
|
|
)
|
|
|
|
# Without file_search, content SHOULD be injected
|
|
found_content = False
|
|
for msgs in context.context_messages.values():
|
|
for m in msgs:
|
|
if "Document Content" in m.text or "Contoso" in m.text:
|
|
found_content = True
|
|
assert found_content
|
|
|
|
async def test_file_search_multiple_files(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
pdf_analysis_result: AnalysisResult,
|
|
audio_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
"""Multiple files should each be uploaded to the vector store."""
|
|
mock_backend = self._make_mock_backend()
|
|
# Return different file IDs for each upload
|
|
mock_backend.upload_file = AsyncMock(side_effect=["file_001", "file_002"])
|
|
config = self._make_file_search_config(mock_backend)
|
|
mock_cu_client.begin_analyze_binary = AsyncMock(
|
|
side_effect=[
|
|
_make_mock_poller(pdf_analysis_result),
|
|
_make_mock_poller(audio_analysis_result),
|
|
],
|
|
)
|
|
provider = _make_provider(
|
|
mock_client=mock_cu_client,
|
|
file_search=config,
|
|
)
|
|
|
|
msg = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("Compare these"),
|
|
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "doc.pdf"),
|
|
_make_content_from_data(b"\x00audio-fake", "audio/mp3", "call.mp3"),
|
|
],
|
|
)
|
|
context = _make_context([msg])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
|
|
|
|
# Two files uploaded via backend
|
|
assert mock_backend.upload_file.call_count == 2
|
|
|
|
async def test_file_search_skips_empty_markdown(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
) -> None:
|
|
"""Upload should be skipped when CU returns no markdown content."""
|
|
mock_backend = self._make_mock_backend()
|
|
config = self._make_file_search_config(mock_backend)
|
|
|
|
# Create a result with empty markdown
|
|
empty_result = AnalysisResult({"contents": [{"markdown": "", "fields": {}}]})
|
|
mock_cu_client.begin_analyze_binary = AsyncMock(
|
|
return_value=_make_mock_poller(empty_result),
|
|
)
|
|
provider = _make_provider(
|
|
mock_client=mock_cu_client,
|
|
file_search=config,
|
|
)
|
|
|
|
msg = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("Analyze this"),
|
|
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "empty.pdf"),
|
|
],
|
|
)
|
|
context = _make_context([msg])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
|
|
|
|
# No file should be uploaded (empty markdown)
|
|
mock_backend.upload_file.assert_not_called()
|
|
|
|
async def test_pending_resolution_uploads_to_vector_store(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
pdf_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
"""When a background task completes in file_search mode, content should be
|
|
uploaded to the vector store — NOT injected into context messages."""
|
|
mock_backend = self._make_mock_backend()
|
|
config = self._make_file_search_config(mock_backend)
|
|
provider = _make_provider(
|
|
mock_client=mock_cu_client,
|
|
file_search=config,
|
|
)
|
|
|
|
# Simulate a completed background analysis via continuation token
|
|
mock_poller = _make_mock_poller(pdf_analysis_result)
|
|
mock_poller.done = MagicMock(return_value=True)
|
|
mock_cu_client.begin_analyze = AsyncMock(return_value=mock_poller)
|
|
|
|
state: dict[str, Any] = {
|
|
"_pending_tokens": {
|
|
"report.pdf": {"continuation_token": "tok_fs", "analyzer_id": "prebuilt-documentSearch"}
|
|
},
|
|
"documents": {
|
|
"report.pdf": {
|
|
"status": DocumentStatus.ANALYZING,
|
|
"filename": "report.pdf",
|
|
"media_type": "application/pdf",
|
|
"analyzer_id": "prebuilt-documentSearch",
|
|
"analyzed_at": None,
|
|
"analysis_duration_s": None,
|
|
"upload_duration_s": None,
|
|
"result": None,
|
|
"error": None,
|
|
},
|
|
},
|
|
}
|
|
|
|
msg = Message(role="user", contents=[Content.from_text("Is the report ready?")])
|
|
context = _make_context([msg])
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
|
|
|
|
# Document should be ready
|
|
assert state["documents"]["report.pdf"]["status"] == DocumentStatus.READY
|
|
|
|
# Content should NOT be injected into context messages
|
|
for msgs in context.context_messages.values():
|
|
for m in msgs:
|
|
assert "Document Content" not in m.text
|
|
|
|
# Should be uploaded to vector store via backend
|
|
mock_backend.upload_file.assert_called_once()
|
|
|
|
# Context messages should mention file_search, not "provided above"
|
|
all_msg_text = " ".join(m.text for msgs in context.context_messages.values() for m in msgs)
|
|
assert "file_search" in all_msg_text or any("file_search" in instr for instr in context.instructions)
|
|
assert "provided above" not in all_msg_text
|
|
|
|
|
|
class TestCloseCancel:
|
|
async def test_close_cleans_up(self) -> None:
|
|
"""close() should close the CU client."""
|
|
provider = _make_provider(mock_client=AsyncMock())
|
|
|
|
await provider.close()
|
|
|
|
# Client should be closed (no tasks to cancel — tokens are just strings)
|
|
cast(Any, provider._client.close).assert_called_once()
|
|
|
|
|
|
class TestSessionIsolation:
|
|
"""Verify that per-session state (pending tasks, uploads) is isolated between sessions."""
|
|
|
|
async def test_background_task_isolated_per_session(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
pdf_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
"""A background task from session A must not leak into session B."""
|
|
mock_cu_client.begin_analyze_binary = AsyncMock(return_value=_make_slow_poller(pdf_analysis_result, delay=10.0))
|
|
provider = _make_provider(mock_client=mock_cu_client, max_wait=0.1)
|
|
|
|
# Session A: upload a file that times out → defers to background
|
|
msg_a = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("Analyze this"),
|
|
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "report.pdf"),
|
|
],
|
|
)
|
|
state_a: dict[str, Any] = {}
|
|
context_a = _make_context([msg_a])
|
|
await provider.before_run(agent=_make_mock_agent(), session=AgentSession(), context=context_a, state=state_a)
|
|
|
|
# Session A should have a pending token
|
|
assert "report.pdf" in state_a.get("_pending_tokens", {})
|
|
|
|
# Session B: separate state, no pending tokens
|
|
state_b: dict[str, Any] = {}
|
|
msg_b = Message(role="user", contents=[Content.from_text("Hello")])
|
|
context_b = _make_context([msg_b])
|
|
await provider.before_run(agent=_make_mock_agent(), session=AgentSession(), context=context_b, state=state_b)
|
|
|
|
# Session B must NOT see session A's pending token
|
|
assert "_pending_tokens" not in state_b or "report.pdf" not in state_b.get("_pending_tokens", {})
|
|
# Session B must NOT have session A's documents
|
|
assert "report.pdf" not in state_b.get("documents", {})
|
|
|
|
async def test_completed_task_resolves_in_correct_session(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
pdf_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
"""A completed background task should only inject content into its own session."""
|
|
provider = _make_provider(mock_client=mock_cu_client)
|
|
|
|
# Simulate completed analysis in session A via continuation token
|
|
mock_poller = _make_mock_poller(pdf_analysis_result)
|
|
mock_poller.done = MagicMock(return_value=True)
|
|
mock_cu_client.begin_analyze = AsyncMock(return_value=mock_poller)
|
|
|
|
state_a: dict[str, Any] = {
|
|
"_pending_tokens": {
|
|
"report.pdf": {"continuation_token": "tok_a", "analyzer_id": "prebuilt-documentSearch"}
|
|
},
|
|
"documents": {
|
|
"report.pdf": {
|
|
"status": DocumentStatus.ANALYZING,
|
|
"filename": "report.pdf",
|
|
"media_type": "application/pdf",
|
|
"analyzer_id": "prebuilt-documentSearch",
|
|
"analyzed_at": None,
|
|
"analysis_duration_s": None,
|
|
"upload_duration_s": None,
|
|
"result": None,
|
|
"error": None,
|
|
},
|
|
},
|
|
}
|
|
state_b: dict[str, Any] = {}
|
|
|
|
# Run session A — should resolve the task
|
|
context_a = _make_context([Message(role="user", contents=[Content.from_text("Is it ready?")])])
|
|
await provider.before_run(agent=_make_mock_agent(), session=AgentSession(), context=context_a, state=state_a)
|
|
assert state_a["documents"]["report.pdf"]["status"] == DocumentStatus.READY
|
|
|
|
# Run session B — must NOT have any documents or resolved content
|
|
context_b = _make_context([Message(role="user", contents=[Content.from_text("Hello")])])
|
|
await provider.before_run(agent=_make_mock_agent(), session=AgentSession(), context=context_b, state=state_b)
|
|
assert "report.pdf" not in state_b.get("documents", {})
|
|
# Session B context should have no document-related messages
|
|
assert not any("report.pdf" in m.text for msgs in context_b.context_messages.values() for m in msgs)
|
|
|
|
|
|
class TestAnalyzerAutoDetectionE2E:
|
|
"""End-to-end: verify _analyze_file stores the resolved analyzer in DocumentEntry."""
|
|
|
|
async def test_audio_file_uses_audio_analyzer(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
audio_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
mock_cu_client.begin_analyze_binary = AsyncMock(
|
|
return_value=_make_mock_poller(audio_analysis_result),
|
|
)
|
|
provider = _make_provider(mock_client=mock_cu_client) # analyzer_id=None
|
|
|
|
msg = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("Transcribe this"),
|
|
_make_content_from_data(b"\x00audio", "audio/mp3", "call.mp3"),
|
|
],
|
|
)
|
|
context = _make_context([msg])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
|
|
|
|
assert state["documents"]["call.mp3"]["analyzer_id"] == "prebuilt-audioSearch"
|
|
# CU client should have been called with the audio analyzer
|
|
mock_cu_client.begin_analyze_binary.assert_called_once()
|
|
call_args = mock_cu_client.begin_analyze_binary.call_args
|
|
assert call_args[0][0] == "prebuilt-audioSearch"
|
|
|
|
async def test_video_file_uses_video_analyzer(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
video_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
mock_cu_client.begin_analyze_binary = AsyncMock(
|
|
return_value=_make_mock_poller(video_analysis_result),
|
|
)
|
|
provider = _make_provider(mock_client=mock_cu_client)
|
|
|
|
msg = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("Analyze this video"),
|
|
_make_content_from_data(b"\x00video", "video/mp4", "demo.mp4"),
|
|
],
|
|
)
|
|
context = _make_context([msg])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
|
|
|
|
assert state["documents"]["demo.mp4"]["analyzer_id"] == "prebuilt-videoSearch"
|
|
call_args = mock_cu_client.begin_analyze_binary.call_args
|
|
assert call_args[0][0] == "prebuilt-videoSearch"
|
|
|
|
async def test_pdf_file_uses_document_analyzer(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
pdf_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
mock_cu_client.begin_analyze_binary = AsyncMock(
|
|
return_value=_make_mock_poller(pdf_analysis_result),
|
|
)
|
|
provider = _make_provider(mock_client=mock_cu_client)
|
|
|
|
msg = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("Read this"),
|
|
_make_content_from_data(_SAMPLE_PDF_BYTES, "application/pdf", "report.pdf"),
|
|
],
|
|
)
|
|
context = _make_context([msg])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
|
|
|
|
assert state["documents"]["report.pdf"]["analyzer_id"] == "prebuilt-documentSearch"
|
|
call_args = mock_cu_client.begin_analyze_binary.call_args
|
|
assert call_args[0][0] == "prebuilt-documentSearch"
|
|
|
|
async def test_explicit_override_ignores_media_type(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
audio_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
"""Explicit analyzer_id should override auto-detection even for audio."""
|
|
mock_cu_client.begin_analyze_binary = AsyncMock(
|
|
return_value=_make_mock_poller(audio_analysis_result),
|
|
)
|
|
provider = _make_provider(mock_client=mock_cu_client, analyzer_id="prebuilt-invoice")
|
|
|
|
msg = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("Analyze"),
|
|
_make_content_from_data(b"\x00audio", "audio/mp3", "call.mp3"),
|
|
],
|
|
)
|
|
context = _make_context([msg])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
|
|
|
|
assert state["documents"]["call.mp3"]["analyzer_id"] == "prebuilt-invoice"
|
|
call_args = mock_cu_client.begin_analyze_binary.call_args
|
|
assert call_args[0][0] == "prebuilt-invoice"
|
|
|
|
async def test_per_file_analyzer_overrides_provider_default(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
pdf_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
"""Per-file analyzer_id in additional_properties overrides provider-level default."""
|
|
mock_cu_client.begin_analyze_binary = AsyncMock(
|
|
return_value=_make_mock_poller(pdf_analysis_result),
|
|
)
|
|
# Provider default is prebuilt-documentSearch
|
|
provider = _make_provider(
|
|
mock_client=mock_cu_client,
|
|
analyzer_id="prebuilt-documentSearch",
|
|
)
|
|
|
|
msg = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("Process this invoice"),
|
|
Content.from_data(
|
|
_SAMPLE_PDF_BYTES,
|
|
"application/pdf",
|
|
# Per-file override to prebuilt-invoice
|
|
additional_properties={
|
|
"filename": "invoice.pdf",
|
|
"analyzer_id": "prebuilt-invoice",
|
|
},
|
|
),
|
|
],
|
|
)
|
|
context = _make_context([msg])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
|
|
|
|
# Per-file override should win
|
|
assert state["documents"]["invoice.pdf"]["analyzer_id"] == "prebuilt-invoice"
|
|
call_args = mock_cu_client.begin_analyze_binary.call_args
|
|
assert call_args[0][0] == "prebuilt-invoice"
|
|
|
|
|
|
class TestWarningsExtraction:
|
|
"""Verify that CU RAI warnings are surfaced via ``to_llm_input`` rendering.
|
|
|
|
The SDK serializes ``result.warnings`` under the reserved ``rai_warnings``
|
|
YAML front-matter key. Telemetry filtering of stray ``LLMStats:`` lines is
|
|
handled by the SDK helper (azure-ai-contentunderstanding >= 1.2.0b2).
|
|
"""
|
|
|
|
def test_warnings_included_when_present(self) -> None:
|
|
"""Non-empty warnings should appear under ``rai_warnings`` front-matter key."""
|
|
provider = _make_provider()
|
|
fixture = {
|
|
"contents": [
|
|
{
|
|
"path": "input1",
|
|
"markdown": "Some content",
|
|
"kind": "document",
|
|
}
|
|
],
|
|
"warnings": [
|
|
{
|
|
"code": "ContentFiltered",
|
|
"message": "Content was filtered due to Responsible AI policy.",
|
|
"target": "contents/0/markdown",
|
|
},
|
|
{
|
|
"code": "ContentFiltered",
|
|
"message": "Violence content detected and filtered.",
|
|
},
|
|
],
|
|
}
|
|
result_obj = AnalysisResult(fixture)
|
|
rendered = provider._render_for_llm(result_obj, "doc.pdf")
|
|
|
|
assert "rai_warnings:" in rendered
|
|
assert "ContentFiltered" in rendered
|
|
assert "Content was filtered due to Responsible AI policy." in rendered
|
|
assert "Violence content detected and filtered." in rendered
|
|
|
|
def test_warnings_omitted_when_empty(self, pdf_analysis_result: AnalysisResult) -> None:
|
|
"""The PDF fixture has no warnings, so ``rai_warnings:`` should not appear."""
|
|
provider = _make_provider()
|
|
rendered = provider._render_for_llm(pdf_analysis_result, "report.pdf")
|
|
assert "rai_warnings:" not in rendered
|
|
|
|
|
|
class TestCategoryExtraction:
|
|
"""Verify category metadata (from classifier analyzers) is rendered into output."""
|
|
|
|
def test_category_included_single_segment(self) -> None:
|
|
"""Category from classifier should appear under the ``category`` front-matter key."""
|
|
provider = _make_provider()
|
|
fixture = {
|
|
"contents": [
|
|
{
|
|
"path": "input1",
|
|
"markdown": "Contract text...",
|
|
"kind": "document",
|
|
"category": "Legal Contract",
|
|
}
|
|
],
|
|
}
|
|
result_obj = AnalysisResult(fixture)
|
|
rendered = provider._render_for_llm(result_obj, "contract.pdf")
|
|
assert "category:" in rendered
|
|
assert "Legal Contract" in rendered
|
|
|
|
def test_category_in_multi_segment_video(self) -> None:
|
|
"""Each segment's category should be visible in the rendered text."""
|
|
provider = _make_provider()
|
|
fixture = {
|
|
"contents": [
|
|
{
|
|
"path": "input1",
|
|
"kind": "audioVisual",
|
|
"startTimeMs": 0,
|
|
"endTimeMs": 30000,
|
|
"markdown": "Opening scene with product showcase.",
|
|
"category": "ProductDemo",
|
|
"fields": {
|
|
"Summary": {
|
|
"type": "string",
|
|
"valueString": "Product demo intro",
|
|
}
|
|
},
|
|
},
|
|
{
|
|
"path": "input1",
|
|
"kind": "audioVisual",
|
|
"startTimeMs": 30000,
|
|
"endTimeMs": 60000,
|
|
"markdown": "Customer testimonial segment.",
|
|
"category": "Testimonial",
|
|
"fields": {
|
|
"Summary": {
|
|
"type": "string",
|
|
"valueString": "Customer feedback",
|
|
}
|
|
},
|
|
},
|
|
],
|
|
}
|
|
result_obj = AnalysisResult(fixture)
|
|
rendered = provider._render_for_llm(result_obj, "promo.mp4")
|
|
|
|
# Both segments' markdown content visible.
|
|
assert "Opening scene with product showcase." in rendered
|
|
assert "Customer testimonial segment." in rendered
|
|
# Both categories visible.
|
|
assert "ProductDemo" in rendered
|
|
assert "Testimonial" in rendered
|
|
# Segments must be rendered in source order, not arbitrary.
|
|
assert rendered.index("Opening scene with product showcase.") < rendered.index("Customer testimonial segment.")
|
|
# Category-to-segment mapping must be correct. The SDK separates segments
|
|
# with a ``*****`` line, so split on it and verify each block carries the
|
|
# right category alongside the right markdown body.
|
|
blocks = rendered.split("*****")
|
|
assert len(blocks) == 2, f"expected 2 segment blocks, got {len(blocks)}"
|
|
assert "Opening scene with product showcase." in blocks[0]
|
|
assert "category: ProductDemo" in blocks[0]
|
|
assert "Customer testimonial segment." in blocks[1]
|
|
assert "category: Testimonial" in blocks[1]
|
|
|
|
def test_category_omitted_when_none(self, pdf_analysis_result: AnalysisResult) -> None:
|
|
"""No category should be in output when the analyzer doesn't classify."""
|
|
provider = _make_provider()
|
|
rendered = provider._render_for_llm(pdf_analysis_result, "report.pdf")
|
|
assert "category:" not in rendered
|
|
|
|
|
|
class TestContentRangeSupport:
|
|
"""Verify that content_range from additional_properties is passed to CU."""
|
|
|
|
async def test_content_range_passed_to_begin_analyze(
|
|
self,
|
|
mock_cu_client: AsyncMock,
|
|
pdf_analysis_result: AnalysisResult,
|
|
) -> None:
|
|
"""content_range in additional_properties should be forwarded to AnalysisInput."""
|
|
from azure.ai.contentunderstanding.models import AnalysisInput
|
|
|
|
mock_cu_client.begin_analyze = AsyncMock(return_value=_make_mock_poller(pdf_analysis_result))
|
|
provider = _make_provider(mock_client=mock_cu_client)
|
|
|
|
msg = Message(
|
|
role="user",
|
|
contents=[
|
|
Content.from_text("Analyze pages 1-3"),
|
|
Content.from_uri(
|
|
"https://example.com/report.pdf",
|
|
media_type="application/pdf",
|
|
additional_properties={"filename": "report.pdf", "content_range": "1-3"},
|
|
),
|
|
],
|
|
)
|
|
context = _make_context([msg])
|
|
state: dict[str, Any] = {}
|
|
session = AgentSession()
|
|
|
|
await provider.before_run(agent=_make_mock_agent(), session=session, context=context, state=state)
|
|
|
|
# Verify begin_analyze was called with AnalysisInput containing content_range
|
|
mock_cu_client.begin_analyze.assert_called_once()
|
|
call_kwargs = mock_cu_client.begin_analyze.call_args
|
|
inputs_arg = call_kwargs.kwargs.get("inputs") or call_kwargs[1].get("inputs")
|
|
assert inputs_arg is not None
|
|
assert len(inputs_arg) == 1
|
|
assert isinstance(inputs_arg[0], AnalysisInput)
|
|
assert inputs_arg[0].content_range == "1-3"
|
|
assert inputs_arg[0].url == "https://example.com/report.pdf"
|