diff --git a/application/single_app/config.py b/application/single_app/config.py index a0c5b5b0a..09cae9063 100644 --- a/application/single_app/config.py +++ b/application/single_app/config.py @@ -96,7 +96,7 @@ EXECUTOR_TYPE = 'thread' EXECUTOR_MAX_WORKERS = 30 SESSION_TYPE = 'filesystem' -VERSION = "0.250.218" +VERSION = "0.250.219" IS_DEVELOPMENT = is_development_env_enabled() SESSION_COOKIE_SAMESITE = os.getenv('SESSION_COOKIE_SAMESITE', 'Lax') diff --git a/application/single_app/functions_agent_document_citations.py b/application/single_app/functions_agent_document_citations.py new file mode 100644 index 000000000..4693ce512 --- /dev/null +++ b/application/single_app/functions_agent_document_citations.py @@ -0,0 +1,571 @@ +# functions_agent_document_citations.py + +"""Derive document citations from agent document-search plugin invocations. + +Agents retrieve workspace documents through ``DocumentSearchPlugin``. Those calls are +recorded as agent tool citations, which describe the tool invocation rather than the +documents that were retrieved. This module converts the document payloads carried in +those invocations into the same document citation shape the route-level hybrid search +produces, so agent-discovered documents behave like any other retrieved source. + +Derived citations are *sources*, not cited references. They are intentionally not +capped or pre-filtered; ``functions_citation_tracking`` narrows sources down to the +subset a response actually cited. +""" + +import json +import logging +import os +from typing import Any, Dict, Iterable, List, Mapping, Optional, Tuple + +from functions_appinsights import log_event +from functions_citation_tracking import resolve_citation_location + +try: + # config.py builds live Azure clients on import, so it is unavailable to the + # standalone functional tests that exercise this module. Fall back to the same + # extension set config defines when it cannot be imported. + from config import TABULAR_EXTENSIONS +except Exception: + TABULAR_EXTENSIONS = frozenset({'csv', 'xlsx', 'xls', 'xlsm'}) + +AGENT_DOCUMENT_CITATION_SOURCE = 'agent_document_search' + +DOCUMENT_SEARCH_PLUGIN_NAMES = frozenset({ + 'documentsearchplugin', + 'document_search_plugin', + 'document_search', +}) + +DOCUMENT_SEARCH_RESULT_FUNCTIONS = frozenset({'search_documents'}) +DOCUMENT_CHUNK_FUNCTIONS = frozenset({'retrieve_document_chunks'}) +DOCUMENT_SUMMARY_FUNCTIONS = frozenset({'summarize_document'}) + +DOCUMENT_SEARCH_FUNCTIONS = ( + DOCUMENT_SEARCH_RESULT_FUNCTIONS + | DOCUMENT_CHUNK_FUNCTIONS + | DOCUMENT_SUMMARY_FUNCTIONS +) + +DOCUMENT_SEARCH_CITATION_INSTRUCTIONS = ( + 'When you use any excerpt below in your answer, cite it by copying that entry\'s ' + '"citation" value verbatim, including the bracketed reference id. Do not invent, ' + 'reformat, or renumber citation values.' +) + + +def _normalize_name(value: Any) -> str: + return str(value or '').strip().lower() + + +def is_document_search_plugin(plugin_name: Any) -> bool: + """Return True when the plugin name identifies the document search plugin.""" + return _normalize_name(plugin_name) in DOCUMENT_SEARCH_PLUGIN_NAMES + + +def is_document_search_invocation(plugin_name: Any, function_name: Any) -> bool: + """Return True when an invocation retrieved workspace documents.""" + return ( + is_document_search_plugin(plugin_name) + and _normalize_name(function_name) in DOCUMENT_SEARCH_FUNCTIONS + ) + + +def _is_tabular_file_name(file_name: Any) -> bool: + normalized_file_name = str(file_name or '').strip().lower() + if not normalized_file_name: + return False + + _, extension = os.path.splitext(normalized_file_name) + return extension.lstrip('.') in TABULAR_EXTENSIONS + + +def coerce_result_payload(value: Any) -> Optional[Dict[str, Any]]: + """Return a mapping payload from a plugin result, parsing JSON strings when needed.""" + if isinstance(value, Mapping): + return dict(value) + + if isinstance(value, str): + candidate = value.strip() + if not candidate.startswith('{'): + return None + try: + parsed_value = json.loads(candidate) + except (TypeError, ValueError): + return None + return dict(parsed_value) if isinstance(parsed_value, Mapping) else None + + return None + + +def build_inline_citation_marker( + file_name: Any, + location_label: Any, + location_value: Any, + citation_id: Any, +) -> str: + """Return the literal inline reference the citation tracker can match.""" + normalized_citation_id = str(citation_id or '').strip().lstrip('#').strip() + if not normalized_citation_id: + return '' + + normalized_file_name = str(file_name or 'Document').strip() or 'Document' + normalized_label = str(location_label or 'Page').strip() or 'Page' + normalized_value = str(location_value or '1').strip() or '1' + + return ( + f'(Source: {normalized_file_name}, {normalized_label}: {normalized_value}) ' + f'[#{normalized_citation_id}]' + ) + + +def _derive_document_id(result: Mapping[str, Any], citation_id: Any) -> str: + document_id = str(result.get('document_id') or '').strip() + if document_id: + return document_id + + normalized_citation_id = str(citation_id or '').strip() + if '_' in normalized_citation_id: + return '_'.join(normalized_citation_id.split('_')[:-1]) + + return normalized_citation_id + + +def _build_citation_record( + *, + file_name: Any, + document_id: Any, + citation_id: Any, + page_number: Any, + sheet_name: Any = None, + chunk_text: Any = None, + chunk_id: Any = None, + chunk_sequence: Any = None, + score: Any = None, + group_id: Any = None, + public_workspace_id: Any = None, + version: Any = None, + classification: Any = None, + plugin_name: Any = None, + function_name: Any = None, +) -> Dict[str, Any]: + resolved_file_name = str(file_name or 'Unknown').strip() or 'Unknown' + location_label, location_value = resolve_citation_location( + page_number=page_number, + chunk_text=chunk_text, + sheet_name=sheet_name, + is_tabular=_is_tabular_file_name(resolved_file_name), + ) + + return { + 'file_name': resolved_file_name, + 'document_id': str(document_id or '').strip(), + 'citation_id': str(citation_id or '').strip(), + 'page_number': page_number, + 'sheet_name': sheet_name, + 'location_label': location_label, + 'location_value': location_value, + 'chunk_id': chunk_id, + 'chunk_sequence': chunk_sequence, + 'score': score, + 'group_id': group_id, + 'public_workspace_id': public_workspace_id, + 'version': version, + 'classification': classification, + 'source': AGENT_DOCUMENT_CITATION_SOURCE, + 'agent_document_search': True, + 'plugin_name': plugin_name, + 'function_name': function_name, + } + + +def _resolve_location_number(*candidate_values): + """Return the first non-null candidate so a valid sequence of 0 is preserved.""" + for candidate_value in candidate_values: + if candidate_value is not None and candidate_value != '': + return candidate_value + return None + + +def _build_citation_from_search_result( + result: Mapping[str, Any], + plugin_name: Any, + function_name: Any, +) -> Optional[Dict[str, Any]]: + citation_id = str(result.get('id') or result.get('citation_id') or '').strip() + document_id = _derive_document_id(result, citation_id) + if not citation_id and not document_id: + return None + + chunk_sequence = result.get('chunk_sequence') + return _build_citation_record( + file_name=result.get('file_name') or result.get('title'), + document_id=document_id, + citation_id=citation_id or document_id, + page_number=_resolve_location_number(result.get('page_number'), chunk_sequence), + sheet_name=result.get('sheet_name'), + chunk_text=result.get('chunk_text'), + chunk_id=result.get('chunk_id'), + chunk_sequence=chunk_sequence, + score=result.get('score'), + group_id=result.get('group_id'), + public_workspace_id=result.get('public_workspace_id'), + version=result.get('version'), + classification=result.get('document_classification') or result.get('classification'), + plugin_name=plugin_name, + function_name=function_name, + ) + + +def _build_citations_from_chunk_payload( + payload: Mapping[str, Any], + plugin_name: Any, + function_name: Any, +) -> List[Dict[str, Any]]: + document_item = payload.get('document') if isinstance(payload.get('document'), Mapping) else {} + document_id = str(document_item.get('id') or '').strip() + citations: List[Dict[str, Any]] = [] + + for chunk in payload.get('chunks') or []: + if not isinstance(chunk, Mapping): + continue + + citation_id = str(chunk.get('id') or '').strip() + chunk_document_id = str(chunk.get('document_id') or '').strip() or document_id + if not citation_id and not chunk_document_id: + continue + + chunk_sequence = chunk.get('chunk_sequence') + citations.append(_build_citation_record( + file_name=chunk.get('file_name') or document_item.get('file_name') or document_item.get('title'), + document_id=chunk_document_id, + citation_id=citation_id or chunk_document_id, + page_number=_resolve_location_number(chunk.get('page_number'), chunk_sequence), + sheet_name=chunk.get('sheet_name'), + chunk_text=chunk.get('chunk_text'), + chunk_id=chunk.get('chunk_id'), + chunk_sequence=chunk_sequence, + score=chunk.get('score'), + group_id=chunk.get('group_id') or document_item.get('group_id'), + public_workspace_id=( + chunk.get('public_workspace_id') + or document_item.get('public_workspace_id') + ), + version=chunk.get('version') or document_item.get('version'), + classification=( + chunk.get('document_classification') + or document_item.get('document_classification') + ), + plugin_name=plugin_name, + function_name=function_name, + )) + + return citations + + +def _build_citation_from_document_payload( + payload: Mapping[str, Any], + plugin_name: Any, + function_name: Any, +) -> Optional[Dict[str, Any]]: + document_item = payload.get('document') if isinstance(payload.get('document'), Mapping) else {} + document_id = str(document_item.get('id') or '').strip() + citation_chunk = payload.get('citation_chunk') if isinstance(payload.get('citation_chunk'), Mapping) else {} + if not document_id: + document_id = str(citation_chunk.get('document_id') or '').strip() + if not document_id: + return None + + # Prefer a real indexed chunk so the citation resolves. Chunk ids are not always + # "_1" - video chunks are keyed by second and can start at zero - so a + # synthesized locator is never used when the summary reports its source chunk. + citation_id = str(citation_chunk.get('id') or '').strip() or document_id + chunk_sequence = citation_chunk.get('chunk_sequence') + + return _build_citation_record( + file_name=( + citation_chunk.get('file_name') + or document_item.get('file_name') + or document_item.get('title') + ), + document_id=document_id, + citation_id=citation_id, + page_number=_resolve_location_number(citation_chunk.get('page_number'), chunk_sequence), + chunk_id=citation_chunk.get('chunk_id'), + chunk_sequence=chunk_sequence, + group_id=document_item.get('group_id'), + public_workspace_id=document_item.get('public_workspace_id'), + version=citation_chunk.get('version') or document_item.get('version'), + classification=( + citation_chunk.get('document_classification') + or document_item.get('document_classification') + ), + plugin_name=plugin_name, + function_name=function_name, + ) + + +def build_document_citations_from_result_payload( + payload: Any, + plugin_name: Any = None, + function_name: Any = None, +) -> List[Dict[str, Any]]: + """Return document citations derived from one document-search result payload.""" + result_payload = coerce_result_payload(payload) + if not result_payload or result_payload.get('error'): + return [] + + normalized_function_name = _normalize_name(function_name) + + if normalized_function_name in DOCUMENT_SEARCH_RESULT_FUNCTIONS: + citations = [] + for result in result_payload.get('results') or []: + if not isinstance(result, Mapping): + continue + citation = _build_citation_from_search_result(result, plugin_name, function_name) + if citation: + citations.append(citation) + return citations + + if normalized_function_name in DOCUMENT_CHUNK_FUNCTIONS: + return _build_citations_from_chunk_payload(result_payload, plugin_name, function_name) + + if normalized_function_name in DOCUMENT_SUMMARY_FUNCTIONS: + citation = _build_citation_from_document_payload(result_payload, plugin_name, function_name) + return [citation] if citation else [] + + return [] + + +def _iter_document_search_entries( + entries: Optional[Iterable[Any]], + plugin_name_getter, + function_name_getter, + result_getter, + success_getter, +): + for entry in entries or []: + if entry is None: + continue + + plugin_name = plugin_name_getter(entry) + function_name = function_name_getter(entry) + if not is_document_search_invocation(plugin_name, function_name): + continue + + success_value = success_getter(entry) + if success_value is False: + continue + + yield plugin_name, function_name, result_getter(entry) + + +def build_document_citations_from_agent_citations( + agent_citations: Optional[Iterable[Any]], +) -> List[Dict[str, Any]]: + """Return document citations derived from agent tool citation records.""" + citations: List[Dict[str, Any]] = [] + entries = _iter_document_search_entries( + agent_citations, + lambda entry: entry.get('plugin_name') if isinstance(entry, Mapping) else None, + lambda entry: entry.get('function_name') if isinstance(entry, Mapping) else None, + lambda entry: entry.get('function_result') if isinstance(entry, Mapping) else None, + lambda entry: entry.get('success') if isinstance(entry, Mapping) else None, + ) + + for plugin_name, function_name, result_payload in entries: + citations.extend(build_document_citations_from_result_payload( + result_payload, + plugin_name=plugin_name, + function_name=function_name, + )) + + return citations + + +def build_document_citations_from_invocations( + invocations: Optional[Iterable[Any]], +) -> List[Dict[str, Any]]: + """Return document citations derived from raw plugin invocation records.""" + citations: List[Dict[str, Any]] = [] + entries = _iter_document_search_entries( + invocations, + lambda entry: getattr(entry, 'plugin_name', None), + lambda entry: getattr(entry, 'function_name', None), + lambda entry: getattr(entry, 'result', None), + lambda entry: getattr(entry, 'success', None), + ) + + for plugin_name, function_name, result_payload in entries: + citations.extend(build_document_citations_from_result_payload( + result_payload, + plugin_name=plugin_name, + function_name=function_name, + )) + + return citations + + +def _build_citation_identity(citation: Mapping[str, Any]) -> Tuple[str, str, str, str]: + citation_id = str(citation.get('citation_id') or '').strip() + if citation_id: + return ('citation_id', citation_id, '', '') + + return ( + 'locator', + str(citation.get('document_id') or '').strip(), + str(citation.get('chunk_id') or '').strip(), + str(citation.get('page_number') or '').strip(), + ) + + +def merge_agent_document_citations( + target_citations: Optional[List[Dict[str, Any]]], + derived_citations: Optional[Iterable[Mapping[str, Any]]], +) -> int: + """Append deduplicated derived citations into ``target_citations`` in place. + + Existing entries always win, so route-level citations keep their original metadata + when the same chunk was also retrieved by an agent. Returns the number appended. + """ + if target_citations is None: + return 0 + + seen_identities = { + _build_citation_identity(citation) + for citation in target_citations + if isinstance(citation, Mapping) + } + + appended_count = 0 + for citation in derived_citations or []: + if not isinstance(citation, Mapping): + continue + + identity = _build_citation_identity(citation) + if identity in seen_identities: + continue + + seen_identities.add(identity) + target_citations.append(dict(citation)) + appended_count += 1 + + return appended_count + + +def apply_agent_document_citations( + hybrid_citations: Optional[List[Dict[str, Any]]], + agent_citations: Optional[Iterable[Any]] = None, + sort_key=None, + conversation_id: Any = None, + plugin_invocations: Optional[Iterable[Any]] = None, +) -> int: + """Merge agent document-search results into hybrid citations in place. + + ``plugin_invocations`` lets callers supply raw invocation records in addition to + agent citation records. Streaming cancellation and error paths need this because + invocations are only folded into the agent citation list once a stream completes. + Merging is deduplicated, so passing both sources never double-counts a chunk. + + Returns the number of document citations added, which callers use for capability + usage metadata and telemetry. + """ + if hybrid_citations is None: + return 0 + + derived_citations = build_document_citations_from_agent_citations(agent_citations) + derived_citations.extend(build_document_citations_from_invocations(plugin_invocations)) + if not derived_citations: + return 0 + + appended_count = merge_agent_document_citations(hybrid_citations, derived_citations) + if appended_count and sort_key: + hybrid_citations.sort(key=sort_key, reverse=True) + + if appended_count: + log_event( + '[AGENT_DOCUMENT_CITATIONS] Added document sources from agent document search', + extra={ + 'conversation_id': conversation_id, + 'derived_citation_count': len(derived_citations), + 'added_citation_count': appended_count, + 'total_document_citation_count': len(hybrid_citations), + }, + level=logging.INFO, + ) + + return appended_count + + +def annotate_document_search_payload( + payload: Any, + function_name: Any, +) -> Any: + """Attach ready-to-copy inline citation markers to a document-search payload. + + The markers use the exact format the citation tracker matches, so a model that + copies them promotes the retrieved document into the cited references for the + response and into the conversation's used documents. + """ + if not isinstance(payload, dict) or payload.get('error'): + return payload + + normalized_function_name = _normalize_name(function_name) + annotated_count = 0 + + if normalized_function_name in DOCUMENT_SEARCH_RESULT_FUNCTIONS: + for result in payload.get('results') or []: + if not isinstance(result, dict): + continue + citation = _build_citation_from_search_result(result, None, function_name) + if not citation: + continue + marker = build_inline_citation_marker( + citation.get('file_name'), + citation.get('location_label'), + citation.get('location_value'), + citation.get('citation_id'), + ) + if marker: + result['citation'] = marker + annotated_count += 1 + + elif normalized_function_name in DOCUMENT_CHUNK_FUNCTIONS: + chunk_citations = _build_citations_from_chunk_payload(payload, None, function_name) + citations_by_id = { + str(citation.get('citation_id') or ''): citation + for citation in chunk_citations + } + for chunk in payload.get('chunks') or []: + if not isinstance(chunk, dict): + continue + citation = citations_by_id.get(str(chunk.get('id') or '')) + if not citation: + continue + marker = build_inline_citation_marker( + citation.get('file_name'), + citation.get('location_label'), + citation.get('location_value'), + citation.get('citation_id'), + ) + if marker: + chunk['citation'] = marker + annotated_count += 1 + + elif normalized_function_name in DOCUMENT_SUMMARY_FUNCTIONS: + citation = _build_citation_from_document_payload(payload, None, function_name) + if citation: + marker = build_inline_citation_marker( + citation.get('file_name'), + citation.get('location_label'), + citation.get('location_value'), + citation.get('citation_id'), + ) + if marker: + payload['citation'] = marker + annotated_count += 1 + + if annotated_count: + payload['citation_instructions'] = DOCUMENT_SEARCH_CITATION_INSTRUCTIONS + + return payload diff --git a/application/single_app/functions_citation_tracking.py b/application/single_app/functions_citation_tracking.py index b2da0ab74..c0c16846d 100644 --- a/application/single_app/functions_citation_tracking.py +++ b/application/single_app/functions_citation_tracking.py @@ -160,7 +160,12 @@ def resolve_citation_location( ): return "Location", "Workbook Schema" - return "Page", str(page_number or 1) + # Preserve a valid page or sequence of 0. Video chunks are keyed by second and + # legitimately start at zero, so truthiness would relabel them as page 1. + if page_number is None or page_number == "": + return "Page", "1" + + return "Page", str(page_number) def _source_reference_matches_citation( diff --git a/application/single_app/functions_search_service.py b/application/single_app/functions_search_service.py index 69cc22a59..7d4303b38 100644 --- a/application/single_app/functions_search_service.py +++ b/application/single_app/functions_search_service.py @@ -911,6 +911,27 @@ def get_document_chunks_payload( } +def _build_summary_citation_chunk(chunks): + """Return identifying fields for the first chunk so summaries stay citable.""" + first_chunk = next( + (chunk for chunk in chunks or [] if isinstance(chunk, dict)), + None, + ) + if not first_chunk: + return None + + return { + "id": first_chunk.get("id"), + "document_id": first_chunk.get("document_id"), + "file_name": first_chunk.get("file_name"), + "page_number": first_chunk.get("page_number"), + "chunk_id": first_chunk.get("chunk_id"), + "chunk_sequence": first_chunk.get("chunk_sequence"), + "version": first_chunk.get("version"), + "document_classification": first_chunk.get("document_classification"), + } + + def _render_window_source_text(window_payload): source_parts = [] for chunk in window_payload.get("chunks", []): @@ -1192,6 +1213,7 @@ def summarize_document_content( return { 'document': chunk_payload.get('document'), + 'citation_chunk': _build_summary_citation_chunk(chunk_payload.get('chunks')), 'scope': chunk_payload.get('scope'), 'scope_id': chunk_payload.get('scope_id'), 'chunk_count': chunk_payload.get('chunk_count'), diff --git a/application/single_app/functions_workflow_runner.py b/application/single_app/functions_workflow_runner.py index 2c36b99f0..485621773 100644 --- a/application/single_app/functions_workflow_runner.py +++ b/application/single_app/functions_workflow_runner.py @@ -55,6 +55,7 @@ build_conversation_context_system_message, serialize_conversation_context_snapshot, ) +from functions_agent_document_citations import apply_agent_document_citations from functions_citation_tracking import ( build_cited_source_subsets, initialize_conversation_used_document_tracking, @@ -5859,6 +5860,11 @@ def _create_assistant_message(conversation, workflow, result, trigger_source, ru generated_tabular_outputs.append(generated_file_output) web_search_citations = list(result.get('web_search_citations') or []) hybrid_citations = list(result.get('hybrid_citations') or []) + apply_agent_document_citations( + hybrid_citations, + raw_agent_citations, + conversation_id=conversation.get('id'), + ) citation_tracking = build_cited_source_subsets( result.get('reply', ''), hybrid_citations=hybrid_citations, diff --git a/application/single_app/route_backend_chats.py b/application/single_app/route_backend_chats.py index bad4c20dc..b4224c1a0 100644 --- a/application/single_app/route_backend_chats.py +++ b/application/single_app/route_backend_chats.py @@ -179,6 +179,7 @@ inject_conversation_context_message, serialize_conversation_context_snapshot, ) +from functions_agent_document_citations import apply_agent_document_citations from functions_citation_tracking import ( build_cited_source_subsets, initialize_conversation_used_document_tracking, @@ -2835,6 +2836,35 @@ def _append_new_plugin_invocation_citations( return len(new_invocations) +def _get_current_message_plugin_invocations(user_id, conversation_id): + """Return this message's plugin invocations. + + Invocations are cleared per chat request, so everything the logger holds for the + conversation belongs to the message being generated. Streaming cancellation and + error paths read this directly because invocations are only folded into the agent + citation list once a stream completes normally. + """ + if not user_id or not conversation_id: + return [] + + try: + return get_plugin_logger().get_invocations_for_conversation( + user_id, + conversation_id, + limit=1000, + ) + except Exception as e: + log_event( + '[AGENT_DOCUMENT_CITATIONS] Unable to read plugin invocations for document citations', + extra={ + 'conversation_id': conversation_id, + 'error_message': str(e), + }, + level=logging.WARNING, + ) + return [] + + def normalize_fact_memory_type(memory_type): normalized = str(memory_type or '').strip().lower() if normalized == FACT_MEMORY_TYPE_LEGACY_DESCRIBER: @@ -15314,6 +15344,12 @@ def execute_document_action_chat_request( document_action_agent_citations, document_action_context_json, ) + apply_agent_document_citations( + hybrid_citations_list, + document_action_agent_citations, + sort_key=_build_hybrid_citation_sort_key, + conversation_id=conversation_id, + ) prepared_agent_citations = [] document_generated_analysis_artifacts = list(execution_result.get('generated_analysis_artifacts') or []) document_generated_tabular_outputs = list(execution_result.get('generated_tabular_outputs') or []) @@ -19766,6 +19802,13 @@ def gpt_error(e): request_correlation_id=mixed_source_request_correlation_id, ) assistant_timestamp = datetime.utcnow().isoformat() + apply_agent_document_citations( + hybrid_citations_list, + agent_citations_list, + sort_key=_build_hybrid_citation_sort_key, + conversation_id=conversation_id, + plugin_invocations=_get_current_message_plugin_invocations(user_id, conversation_id), + ) prepared_agent_citations = persist_agent_citation_artifacts( conversation_id=conversation_id, assistant_message_id=assistant_message_id, @@ -19806,7 +19849,9 @@ def gpt_error(e): assistant_capability_usage = _build_capability_usage_metadata( workspace_search_enabled=mixed_source_document_context_active, workspace_search_used=bool( - search_results or mixed_source_has_authorized_evidence_sources + search_results + or mixed_source_has_authorized_evidence_sources + or hybrid_citations_list ), workspace_search_result_count=len(hybrid_citations_list or []), document_action_type=DOCUMENT_ACTION_TYPE_NONE, @@ -20648,7 +20693,9 @@ def build_streaming_capability_usage(): return _build_capability_usage_metadata( workspace_search_enabled=mixed_source_document_context_active, workspace_search_used=bool( - search_results or mixed_source_has_authorized_evidence_sources + search_results + or mixed_source_has_authorized_evidence_sources + or hybrid_citations_list ), workspace_search_result_count=len(hybrid_citations_list or []), document_action_type=DOCUMENT_ACTION_TYPE_NONE, @@ -23031,6 +23078,13 @@ def finalize_cancelled_stream_response(): if partial_content: assistant_timestamp = datetime.utcnow().isoformat() + apply_agent_document_citations( + hybrid_citations_list, + agent_citations_list, + sort_key=_build_hybrid_citation_sort_key, + conversation_id=conversation_id, + plugin_invocations=_get_current_message_plugin_invocations(user_id, conversation_id), + ) partial_citation_tracking = build_cited_source_subsets( partial_content, hybrid_citations=hybrid_citations_list, @@ -23687,6 +23741,13 @@ def finalize_cancelled_agent_stream_response(): request_correlation_id=mixed_source_request_correlation_id, ) assistant_timestamp = datetime.utcnow().isoformat() + apply_agent_document_citations( + hybrid_citations_list, + agent_citations_list, + sort_key=_build_hybrid_citation_sort_key, + conversation_id=conversation_id, + plugin_invocations=_get_current_message_plugin_invocations(user_id, conversation_id), + ) prepared_agent_citations = persist_agent_citation_artifacts( conversation_id=conversation_id, assistant_message_id=assistant_message_id, @@ -24027,6 +24088,13 @@ def finalize_cancelled_agent_stream_response(): if accumulated_content: current_assistant_thread_id = str(uuid.uuid4()) assistant_timestamp = datetime.utcnow().isoformat() + apply_agent_document_citations( + hybrid_citations_list, + agent_citations_list, + sort_key=_build_hybrid_citation_sort_key, + conversation_id=conversation_id, + plugin_invocations=_get_current_message_plugin_invocations(user_id, conversation_id), + ) interrupted_citation_tracking = build_cited_source_subsets( accumulated_content, hybrid_citations=hybrid_citations_list, diff --git a/application/single_app/semantic_kernel_plugins/document_search_plugin.py b/application/single_app/semantic_kernel_plugins/document_search_plugin.py index cf48ade2a..21a681a30 100644 --- a/application/single_app/semantic_kernel_plugins/document_search_plugin.py +++ b/application/single_app/semantic_kernel_plugins/document_search_plugin.py @@ -5,6 +5,7 @@ from semantic_kernel.functions import kernel_function from functions_authentication import get_current_user_id +from functions_agent_document_citations import annotate_document_search_payload from functions_search import SEARCH_DEFAULT_TOP_N, SEARCH_MAX_TOP_N, normalize_search_scope, normalize_search_top_n from functions_search_service import ( SUMMARY_DEFAULT_FINAL_TARGET, @@ -133,7 +134,10 @@ def _resolve_target_length(self, requested_value: str, manifest_key: str, fallba @plugin_function_logger('DocumentSearchPlugin') @kernel_function( name='search_documents', - description='Run hybrid document search over accessible workspaces and return chunk-level results with document ids.', + description=( + 'Run hybrid document search over accessible workspaces and return chunk-level results with document ids. ' + 'Every result carries a "citation" value; copy it verbatim into your answer whenever you use that excerpt.' + ), ) def search_documents( self, @@ -146,15 +150,18 @@ def search_documents( active_public_workspace_id: Annotated[str, 'Optional public workspace id when searching public content.'] = '', ) -> Annotated[dict, 'Search results and request metadata.']: try: - return run_document_search( - query=query, - user_id=self._get_user_id(), - top_n=self._resolve_top_n(top_n), - doc_scope=self._resolve_doc_scope(doc_scope), - document_ids=document_ids, - tags_filter=tags_filter, - active_group_ids=active_group_ids, - active_public_workspace_id=active_public_workspace_id, + return annotate_document_search_payload( + run_document_search( + query=query, + user_id=self._get_user_id(), + top_n=self._resolve_top_n(top_n), + doc_scope=self._resolve_doc_scope(doc_scope), + document_ids=document_ids, + tags_filter=tags_filter, + active_group_ids=active_group_ids, + active_public_workspace_id=active_public_workspace_id, + ), + 'search_documents', ) except Exception as e: return {'error': str(e)} @@ -163,7 +170,10 @@ def search_documents( @plugin_function_logger('DocumentSearchPlugin') @kernel_function( name='retrieve_document_chunks', - description='Retrieve ordered chunks for one accessible document, optionally selecting one window of chunks.', + description=( + 'Retrieve ordered chunks for one accessible document, optionally selecting one window of chunks. ' + 'Every chunk carries a "citation" value; copy it verbatim into your answer whenever you use that chunk.' + ), ) def retrieve_document_chunks( self, @@ -177,16 +187,19 @@ def retrieve_document_chunks( active_public_workspace_id: Annotated[str, 'Optional public workspace id when resolving public content.'] = '', ) -> Annotated[dict, 'Ordered chunks and window metadata for one document.']: try: - return get_document_chunks_payload( - document_id=document_id, - user_id=self._get_user_id(), - doc_scope=self._resolve_doc_scope(doc_scope), - active_group_ids=active_group_ids, - active_public_workspace_id=active_public_workspace_id, - window_unit=self._resolve_window_unit(window_unit), - window_size=self._resolve_optional_window_value(window_size, 'default_window_size'), - window_percent=self._resolve_optional_window_value(window_percent, 'default_window_percent'), - window_number=window_number if int(window_number or 0) > 0 else None, + return annotate_document_search_payload( + get_document_chunks_payload( + document_id=document_id, + user_id=self._get_user_id(), + doc_scope=self._resolve_doc_scope(doc_scope), + active_group_ids=active_group_ids, + active_public_workspace_id=active_public_workspace_id, + window_unit=self._resolve_window_unit(window_unit), + window_size=self._resolve_optional_window_value(window_size, 'default_window_size'), + window_percent=self._resolve_optional_window_value(window_percent, 'default_window_percent'), + window_number=window_number if int(window_number or 0) > 0 else None, + ), + 'retrieve_document_chunks', ) except Exception as e: return {'error': str(e)} @@ -195,7 +208,10 @@ def retrieve_document_chunks( @plugin_function_logger('DocumentSearchPlugin') @kernel_function( name='summarize_document', - description='Summarize one accessible document hierarchically across ordered chunk windows, with optional focus guidance.', + description=( + 'Summarize one accessible document hierarchically across ordered chunk windows, with optional focus guidance. ' + 'The payload carries a "citation" value; copy it verbatim into your answer when you use the summary.' + ), ) def summarize_document( self, @@ -211,18 +227,21 @@ def summarize_document( active_public_workspace_id: Annotated[str, 'Optional public workspace id when resolving public content.'] = '', ) -> Annotated[dict, 'Final summary text plus stage and window metadata.']: try: - return summarize_document_content( - document_id=document_id, - user_id=self._get_user_id(), - doc_scope=self._resolve_doc_scope(doc_scope), - active_group_ids=active_group_ids, - active_public_workspace_id=active_public_workspace_id, - focus_instructions=self._resolve_focus_instructions(focus_instructions), - final_target_length=self._resolve_target_length(final_target_length, 'default_final_target_length', SUMMARY_DEFAULT_FINAL_TARGET), - window_target_length=self._resolve_target_length(window_target_length, 'default_window_target_length', SUMMARY_DEFAULT_WINDOW_SUMMARY_TARGET), - window_unit=self._resolve_window_unit(window_unit), - window_size=self._resolve_optional_window_value(window_size, 'default_window_size'), - window_percent=self._resolve_optional_window_value(window_percent, 'default_window_percent'), + return annotate_document_search_payload( + summarize_document_content( + document_id=document_id, + user_id=self._get_user_id(), + doc_scope=self._resolve_doc_scope(doc_scope), + active_group_ids=active_group_ids, + active_public_workspace_id=active_public_workspace_id, + focus_instructions=self._resolve_focus_instructions(focus_instructions), + final_target_length=self._resolve_target_length(final_target_length, 'default_final_target_length', SUMMARY_DEFAULT_FINAL_TARGET), + window_target_length=self._resolve_target_length(window_target_length, 'default_window_target_length', SUMMARY_DEFAULT_WINDOW_SUMMARY_TARGET), + window_unit=self._resolve_window_unit(window_unit), + window_size=self._resolve_optional_window_value(window_size, 'default_window_size'), + window_percent=self._resolve_optional_window_value(window_percent, 'default_window_percent'), + ), + 'summarize_document', ) except Exception as e: return {'error': str(e)} \ No newline at end of file diff --git a/application/single_app/static/js/chat/chat-citations.js b/application/single_app/static/js/chat/chat-citations.js index edd6ac293..077cc1396 100644 --- a/application/single_app/static/js/chat/chat-citations.js +++ b/application/single_app/static/js/chat/chat-citations.js @@ -873,6 +873,38 @@ export function showPdfModal(docId, pageNumber, citationId) { } // -------------------------------------------------------------------- +function toggleCitationOverflowGroup(toggleButton) { + const citationsContainer = toggleButton.closest(".citations-container"); + const overflowGroup = citationsContainer?.querySelector(".citation-overflow-group"); + if (!overflowGroup) { + return; + } + + const isCollapsed = overflowGroup.classList.contains("d-none"); + overflowGroup.classList.toggle("d-none", !isCollapsed); + toggleButton.setAttribute("aria-expanded", String(isCollapsed)); + + const label = isCollapsed + ? toggleButton.dataset.expandedLabel || "Show fewer sources" + : toggleButton.dataset.collapsedLabel || "Show more sources"; + const icon = document.createElement("i"); + icon.className = `bi ${isCollapsed ? "bi-dash-circle" : "bi-plus-circle"} me-1`; + + toggleButton.replaceChildren(icon, document.createTextNode(label)); + toggleButton.title = label; +} + +document.addEventListener("click", (event) => { + const eventTarget = event.target instanceof Element ? event.target : null; + const toggleButton = eventTarget?.closest("button.citation-overflow-toggle"); + if (!toggleButton) { + return; + } + + event.preventDefault(); + toggleCitationOverflowGroup(toggleButton); +}); + // --- MODIFIED: Event Listener Logic --- if (chatboxEl) { chatboxEl.addEventListener("click", (event) => { diff --git a/application/single_app/static/js/chat/chat-messages.js b/application/single_app/static/js/chat/chat-messages.js index da6310144..90f9487b5 100644 --- a/application/single_app/static/js/chat/chat-messages.js +++ b/application/single_app/static/js/chat/chat-messages.js @@ -1862,6 +1862,23 @@ function resolveHybridCitationId(cite, index) { return `${cite?.chunk_id || ''}_${cite?.page_number || index}`; } +// Agent document search can return hundreds of source chunks for one answer, so the +// source list is collapsed past this many entries instead of being capped server-side. +const DOCUMENT_CITATION_VISIBLE_LIMIT = 25; + +function buildDocumentCitationGroupHtml(citationParts) { + if (!Array.isArray(citationParts) || citationParts.length <= DOCUMENT_CITATION_VISIBLE_LIMIT) { + return Array.isArray(citationParts) ? citationParts.join("") : ""; + } + + const visibleParts = citationParts.slice(0, DOCUMENT_CITATION_VISIBLE_LIMIT); + const overflowParts = citationParts.slice(DOCUMENT_CITATION_VISIBLE_LIMIT); + const collapsedLabel = `Show ${overflowParts.length} more sources`; + const expandedLabel = "Show fewer sources"; + + return `${visibleParts.join("")}${overflowParts.join("")}`; +} + function createCitationsHtml( hybridCitations = [], webCitations = [], @@ -1874,6 +1891,7 @@ function createCitationsHtml( if (hybridCitations && hybridCitations.length > 0) { hasCitations = true; + const documentCitationParts = []; hybridCitations.forEach((cite, index) => { const citationId = resolveHybridCitationId(cite, index); const fileName = cite.file_name || 'Document'; @@ -1907,7 +1925,7 @@ function createCitationsHtml( if (isMetadata && documentId) { const summaryText = `${escapeHtml(locationLabel)}: ${escapeHtml(locationValue)}`; - citationsHtml += ` + documentCitationParts.push(` ${summaryText} - `; + `); return; } - citationsHtml += ` + documentCitationParts.push(` ${displayText} - `; + `); }); + citationsHtml += buildDocumentCitationGroupHtml(documentCitationParts); } if (webCitations && webCitations.length > 0) { diff --git a/docs/explanation/fixes/AGENT_DOCUMENT_SEARCH_CITATION_FIX.md b/docs/explanation/fixes/AGENT_DOCUMENT_SEARCH_CITATION_FIX.md new file mode 100644 index 000000000..8d3a41f83 --- /dev/null +++ b/docs/explanation/fixes/AGENT_DOCUMENT_SEARCH_CITATION_FIX.md @@ -0,0 +1,129 @@ +# Agent Document Search Citation Fix + +## Header information + +**Issue:** Documents an agent retrieved through the document search action never became document citations. They were recorded only as agent tool citations, so the retrieved documents did not appear in the message Sources disclosure, were not clickable, never reached the enhanced citation viewer, and could not be promoted into cited references or the conversation's Used documents drawer. + +**Root cause:** The route-level hybrid search built document citation records directly into `hybrid_citations`. `DocumentSearchPlugin` results took a different path: `_build_plugin_invocation_agent_citation()` wrapped the entire invocation into an `agent_citations` entry describing the tool call (`tool_name`, `function_arguments`, `function_result`). The retrieved document payload was present inside `function_result`, but nothing converted it into the document citation shape, and the plugin returned no citation markers for the model to reference. `functions_workflow_runner.py` had the same gap through `_build_agent_citations_from_plugin_invocations()`. + +Fixed/Implemented in version: **0.250.219** + +Related config.py update: `VERSION = "0.250.219"` + +Associated issue: `microsoft/simplechat#1239` + +## Technical details + +### Behavior comparison + +| | Standard document search | Agent document search (before) | Agent document search (after) | +|---|---|---|---| +| Retrieval | route-level `hybrid_search()` | `DocumentSearchPlugin.search_documents()` | unchanged | +| Stored as | `hybrid_citations` | `agent_citations` only | `agent_citations` **and** `hybrid_citations` | +| Sources disclosure | document buttons | raw JSON tool modal | document buttons plus the tool modal | +| Enhanced citation / PDF viewer | yes | no | yes | +| Eligible for `cited_hybrid_citations` | yes | no | yes | +| Counted in `used_documents` | yes | no | yes | +| Capability usage `workspace.search_used` | yes | no | yes | + +### Sources are not cited references + +Derived citations are *sources*. They are intentionally not capped or pre-filtered, so an agent that sources 500 chunks records 500 sources. `functions_citation_tracking.build_cited_source_subsets()` remains solely responsible for narrowing sources down to the subset the response actually cited. + +### Derivation + +`functions_agent_document_citations.py` converts document-search payloads into the same record shape the route-level search produces (`file_name`, `document_id`, `citation_id`, `page_number`, `sheet_name`, `location_label`, `location_value`, `chunk_id`, `chunk_sequence`, `score`, `group_id`, `public_workspace_id`, `version`, `classification`), using the shared `resolve_citation_location()` helper so tabular sheets and page locations match exactly. + +| Plugin function | Payload shape | Derived citations | +|---|---|---| +| `search_documents` | `results[]` | one per retrieved chunk | +| `retrieve_document_chunks` | `document` + `chunks[]` | one per returned chunk | +| `summarize_document` | `document` + `citation_chunk` | one citation anchored to a real source chunk | + +Each derived record is tagged with `source: "agent_document_search"`, plus `plugin_name` and `function_name`, so its provenance is auditable. Invocations from other plugins, failed invocations, and payloads containing `error` are ignored. + +Merging deduplicates by `citation_id`, falling back to `document_id` + `chunk_id` + `page_number`. Existing route-level records always win, so a chunk retrieved by both the document search toggle and an agent keeps its original metadata and is listed once. + +### Locator accuracy + +Chunk keys are not always `_1`. Video chunks are keyed by second and legitimately start at `_0`, and some documents have no page 1. Three rules keep derived links resolvable: + +- `summarize_document_content()` now reports a `citation_chunk` describing its first source chunk, and the summary citation uses that chunk's real id, page, and sequence. When a payload reports no source chunk, the citation falls back to the document id rather than synthesizing a `_1` locator that may not exist. +- Page and sequence resolution uses explicit null checks instead of truthiness, so a valid sequence of `0` is preserved rather than being rewritten to `1`. +- `resolve_citation_location()` in `functions_citation_tracking.py` also preserves `0` instead of relabelling it as page 1. Every pre-existing caller already coerced `0` to `1` before calling, so this only affects the new agent-derived path and keeps the displayed location consistent with the citation id. + +### Cancelled and interrupted streams + +In the streaming path, plugin invocations are only folded into the agent citation list once a stream completes normally. A stream cancelled or interrupted after a document search would otherwise persist an empty citation list. Those two paths therefore also pass the raw invocation records from the plugin logger, which is cleared per chat request so it holds only the current message's invocations. Merging is deduplicated, so supplying both sources never double-counts a chunk. + +### Inline citation markers + +Every document-search payload now carries a ready-to-copy `citation` value formatted exactly as the citation tracker matches: + +``` +(Source: Policy.pdf, Page: 3) [#doc-1_3] +``` + +The payload also carries `citation_instructions`, and the three kernel function descriptions instruct the model to copy the `citation` value verbatim. When the model does, `build_cited_source_subsets()` promotes that document into `cited_hybrid_citations`, and `merge_cited_documents_into_conversation()` records it in the conversation's `used_documents`. + +### Large source lists + +Because sources are uncapped, the per-message Sources disclosure now renders the first 25 document sources and collapses the remainder behind a **Show N more sources** control (`citation-overflow-group` / `citation-overflow-toggle`). The control uses the Bootstrap `d-none` class and a delegated click handler; no stored data is truncated. The Used documents drawer is unaffected because it lists cited documents, not source chunks. + +### Files modified + +- `application/single_app/functions_agent_document_citations.py` (new) +- `application/single_app/route_backend_chats.py` +- `application/single_app/functions_workflow_runner.py` +- `application/single_app/functions_search_service.py` +- `application/single_app/functions_citation_tracking.py` +- `application/single_app/semantic_kernel_plugins/document_search_plugin.py` +- `application/single_app/static/js/chat/chat-messages.js` +- `application/single_app/static/js/chat/chat-citations.js` +- `application/single_app/config.py` + +### Integration points + +The merge runs before `build_cited_source_subsets()` and before `persist_agent_citation_artifacts()` in every path that writes an assistant message: + +- `route_backend_chats.py` — document action path +- `route_backend_chats.py` — non-streaming chat path +- `route_backend_chats.py` — streaming completion path +- `route_backend_chats.py` — streaming cancellation path (partial content) +- `route_backend_chats.py` — streaming interruption path +- `functions_workflow_runner.py` — `_create_assistant_message()` + +### Access control + +No new access surface is introduced. Every derived record comes from payloads produced by `functions_search_service`, which already resolves personal, group, and public scope against the current user before returning any content. Derivation is a read of data the requesting user was already authorized to receive. + +## Validation + +`functional_tests/test_agent_document_search_citations.py` — 12/12 passing: + +1. `config.py` VERSION is at or above the implementation version +2. `search_documents` results map to the route-level document citation shape +3. `retrieve_document_chunks` and `summarize_document` produce citations anchored to real chunks +4. Zero-indexed sequences survive and summaries never invent a `_1` locator +5. Raw plugin invocations produce citations for cancelled and interrupted streams, deduplicated against agent citations +6. Unrelated plugins, failed invocations, and errored payloads are ignored +7. Merging dedupes against route citations, preserves them, and does not truncate a 500-result set +8. Payload markers are matched by the citation tracker and promote the correct document into `cited_hybrid_citations` +9. Tabular sheet locations and JSON-string payloads are handled +10. All chat paths, the workflow runner, and all three plugin functions are wired +11. Every function that builds cited subsets merges agent document citations first on **every branch**, verified by AST walking so a merge inside one finalization branch cannot vouch for another +12. Large source lists collapse behind a show-more control + +The ordering test was confirmed to have teeth: removing the interrupted-stream merge makes it fail with the exact branch and line, and it independently caught a path that was missing a merge during development. + +Regression coverage confirmed passing: `test_chat_cited_source_tracking.py`, `test_agent_citations_fix.py`, `test_agent_citations_per_message_fix.py`, `test_chat_capability_usage_metadata.py`, `test_markdown_citation_lookup_fallback.py`, `test_stored_xss_chat_workspace_rendering_fix.py`, `test_mixed_source_chat_search_consistency.py`, and the three `route_tests` policy suites. + +## Known limitation + +`PluginInvocationLogger` is process-global and filters only by user and conversation. If a user supersedes a request while the prior one is still finalizing, the newer request's `clear_invocations_for_conversation()` can remove invocations the older request has not yet read. This predates the change and affects existing agent tool citation capture identically; making it exact requires tagging invocations with a request or run id and filtering on it, which is a broader change to shared plugin infrastructure and is better handled separately. + +### Before and after + +**Before** — an agent answering from workspace documents produced a message with zero document sources. The only evidence of retrieval was an "Agent" tool citation that opened a JSON modal, and the conversation's Used documents drawer stayed empty. + +**After** — the same answer lists every retrieved document chunk as a source, each source opens the document at the correct page or sheet, cited documents are separated from retrieved sources, and the Used documents drawer reflects the documents the answer actually cited. diff --git a/docs/explanation/release_notes.md b/docs/explanation/release_notes.md index 248d506a9..c661ac6a5 100644 --- a/docs/explanation/release_notes.md +++ b/docs/explanation/release_notes.md @@ -2,6 +2,27 @@ For feature-focused and fix-focused drill-downs by version, see [Features by Version](/explanation/features/) and [Fixes by Version](/explanation/fixes/). +### **(v0.250.219)** + +#### Bug Fixes + +* **Agent Document Search Now Produces Real Document Citations** + * Documents an agent retrieved through the document search action are now recorded as document sources instead of only as an agent tool call. Previously they appeared solely as a raw JSON tool modal, so the documents were missing from the message Sources disclosure, were not clickable, never opened in the enhanced citation viewer, and could never reach the Used documents drawer. + * Covers all three document search functions — relevance-ranked search, ordered chunk retrieval, and document summarization — across personal, group, and public workspaces. + * Document search results now carry a ready-to-copy citation value, and the action instructs the model to reuse it verbatim. When the answer cites a document, it is correctly separated from the retrieved sources and recorded in the conversation's used documents. + * Applies to streaming and non-streaming chat, document actions, cancelled and interrupted streams, and scheduled workflow runs. + * Retrieved sources are deliberately not capped, so a search that sources hundreds of chunks records all of them. Chunks retrieved by both the document search toggle and an agent are listed once. + * Cancelled and interrupted streams keep the documents the agent had already retrieved, and citation locations no longer relabel a valid page or sequence of `0` as page 1, which affected video chunks keyed by second. + * Workspace capability metadata now reports document usage for agent-only document turns, which previously under-reported as unused. + * (Ref: #1239, `functions_agent_document_citations.py`, `route_backend_chats.py`, `functions_workflow_runner.py`, `document_search_plugin.py`, `AGENT_DOCUMENT_SEARCH_CITATION_FIX.md`) + +#### User Interface Enhancements + +* **Collapsed Long Source Lists** + * The per-message Sources disclosure now shows the first 25 document sources and collapses the rest behind a **Show N more sources** control, so an agent that retrieves hundreds of chunks no longer floods the panel. + * No source data is discarded — the full set is still stored, exported, and available for citation matching. + * (Ref: #1239, `chat-messages.js`, `chat-citations.js`) + ### **(v0.250.218)** #### Bug Fixes diff --git a/functional_tests/test_agent_document_search_citations.py b/functional_tests/test_agent_document_search_citations.py new file mode 100644 index 000000000..807e7b18b --- /dev/null +++ b/functional_tests/test_agent_document_search_citations.py @@ -0,0 +1,819 @@ +#!/usr/bin/env python3 +# test_agent_document_search_citations.py +""" +Functional test for agent document-search citations. +Version: 0.250.219 +Implemented in: 0.250.219 + +This test ensures that documents an agent retrieves through DocumentSearchPlugin +produce the same document citation shape as the route-level hybrid search, so they +appear as message sources, feed the sources-vs-cited-reference tracking, and can be +promoted into the conversation's used documents when the response cites them. + +Related issue: microsoft/simplechat#1239 +""" + +import ast +import os +import sys + +sys.path.append(os.path.dirname(os.path.abspath(__file__))) + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +APP_ROOT = os.path.join(REPO_ROOT, 'application', 'single_app') + +if APP_ROOT not in sys.path: + sys.path.insert(0, APP_ROOT) + +from test_support.versioning import assert_app_version_at_least + +import functions_agent_document_citations as agent_document_citations +from functions_citation_tracking import build_cited_source_subsets + +IMPLEMENTED_VERSION = '0.250.219' + + +def read_source(*relative_parts): + file_path = os.path.join(REPO_ROOT, *relative_parts) + with open(file_path, 'r', encoding='utf-8') as source_file: + return source_file.read() + + +def build_search_agent_citation(results): + return { + 'plugin_name': 'DocumentSearchPlugin', + 'function_name': 'search_documents', + 'function_result': {'query': 'policy', 'scope': 'all', 'results': results}, + 'success': True, + } + + +def test_version_is_at_least_implementation_version(): + """The fix must be present in at least the version it shipped in.""" + print('🔍 Validating application version...') + try: + app_version = assert_app_version_at_least( + IMPLEMENTED_VERSION, + reason='Agent document-search citations were added in this version.', + ) + print(f'✅ config.py VERSION {app_version} is at or above {IMPLEMENTED_VERSION}') + return True + except Exception as e: + print(f'❌ Version check failed: {e}') + return False + + +def test_search_results_become_document_citations(): + """search_documents results must map to the route-level document citation shape.""" + print('🔍 Validating search_documents citation derivation...') + try: + agent_citation = build_search_agent_citation([ + { + 'id': 'doc-1_3', + 'document_id': 'doc-1', + 'file_name': 'Policy.pdf', + 'page_number': 3, + 'chunk_id': 'chunk-3', + 'chunk_sequence': 3, + 'score': 2.5, + 'version': 4, + 'document_classification': 'internal', + 'group_id': 'group-9', + }, + ]) + + citations = agent_document_citations.build_document_citations_from_agent_citations( + [agent_citation] + ) + if len(citations) != 1: + print(f'❌ Expected 1 derived citation, found {len(citations)}') + return False + + citation = citations[0] + required_fields = { + 'file_name': 'Policy.pdf', + 'document_id': 'doc-1', + 'citation_id': 'doc-1_3', + 'page_number': 3, + 'location_label': 'Page', + 'location_value': '3', + 'chunk_id': 'chunk-3', + 'chunk_sequence': 3, + 'score': 2.5, + 'version': 4, + 'classification': 'internal', + 'group_id': 'group-9', + } + for field_name, expected_value in required_fields.items(): + if citation.get(field_name) != expected_value: + print( + f'❌ Citation field {field_name} was {citation.get(field_name)!r}, ' + f'expected {expected_value!r}' + ) + return False + + if citation.get('source') != agent_document_citations.AGENT_DOCUMENT_CITATION_SOURCE: + print('❌ Derived citation is missing agent document-search provenance') + return False + + print('✅ search_documents results produce document citations') + return True + except Exception as e: + print(f'❌ Test failed: {e}') + import traceback + traceback.print_exc() + return False + + +def test_chunk_and_summary_functions_produce_citations(): + """retrieve_document_chunks and summarize_document must also produce citations.""" + print('🔍 Validating chunk and summary citation derivation...') + try: + chunk_citations = agent_document_citations.build_document_citations_from_agent_citations([ + { + 'plugin_name': 'DocumentSearchPlugin', + 'function_name': 'retrieve_document_chunks', + 'function_result': { + 'document': {'id': 'doc-2', 'file_name': 'Guide.pdf'}, + 'chunks': [ + {'id': 'doc-2_5', 'document_id': 'doc-2', 'page_number': 5, 'chunk_id': 'chunk-5'}, + {'id': 'doc-2_6', 'document_id': 'doc-2', 'page_number': 6, 'chunk_id': 'chunk-6'}, + ], + }, + 'success': True, + }, + ]) + if len(chunk_citations) != 2: + print(f'❌ Expected 2 chunk citations, found {len(chunk_citations)}') + return False + if chunk_citations[0].get('citation_id') != 'doc-2_5': + print('❌ Chunk citation id was not taken from the indexed chunk id') + return False + + summary_citations = agent_document_citations.build_document_citations_from_agent_citations([ + { + 'plugin_name': 'DocumentSearchPlugin', + 'function_name': 'summarize_document', + 'function_result': { + 'document': {'id': 'doc-3', 'file_name': 'Spec.docx'}, + 'citation_chunk': { + 'id': 'doc-3_4', + 'document_id': 'doc-3', + 'file_name': 'Spec.docx', + 'page_number': 4, + 'chunk_id': 'chunk-4', + }, + 'summary': 'A summary.', + }, + 'success': True, + }, + ]) + if len(summary_citations) != 1: + print(f'❌ Expected 1 summary citation, found {len(summary_citations)}') + return False + if summary_citations[0].get('document_id') != 'doc-3': + print('❌ Summary citation did not carry the summarized document id') + return False + if summary_citations[0].get('citation_id') != 'doc-3_4': + print('❌ Summary citation did not use the real source chunk id') + return False + + print('✅ Chunk retrieval and summarization produce document citations') + return True + except Exception as e: + print(f'❌ Test failed: {e}') + import traceback + traceback.print_exc() + return False + + +def test_zero_indexed_and_missing_locators_are_not_faked(): + """Zero sequences must survive and summaries must never invent a chunk locator.""" + print('🔍 Validating locator handling...') + try: + # Video chunks are keyed by second and legitimately start at zero. + video_citations = agent_document_citations.build_document_citations_from_agent_citations([ + build_search_agent_citation([ + { + 'id': 'video-1_0', + 'document_id': 'video-1', + 'file_name': 'Briefing.mp4', + 'chunk_sequence': 0, + }, + ]), + ]) + if len(video_citations) != 1: + print(f'❌ Expected 1 video citation, found {len(video_citations)}') + return False + if video_citations[0].get('page_number') != 0: + print( + '❌ A valid chunk sequence of 0 was replaced with ' + f'{video_citations[0].get("page_number")!r}' + ) + return False + if video_citations[0].get('citation_id') != 'video-1_0': + print('❌ Video citation id was not preserved') + return False + if video_citations[0].get('location_value') != '0': + print( + '❌ A zero location was displayed as ' + f'{video_citations[0].get("location_value")!r} instead of "0"' + ) + return False + + # The emitted marker must point at the same location the id encodes, otherwise + # the rendered inline link cannot resolve. + zero_payload = agent_document_citations.annotate_document_search_payload( + { + 'results': [{ + 'id': 'video-1_0', + 'document_id': 'video-1', + 'file_name': 'Briefing.mp4', + 'chunk_sequence': 0, + }], + }, + 'search_documents', + ) + zero_marker = zero_payload['results'][0].get('citation') + if zero_marker != '(Source: Briefing.mp4, Page: 0) [#video-1_0]': + print(f'❌ Unexpected zero-location citation marker: {zero_marker!r}') + return False + + # Without a reported source chunk the summary must not synthesize "_1". + summary_citations = agent_document_citations.build_document_citations_from_agent_citations([ + { + 'plugin_name': 'DocumentSearchPlugin', + 'function_name': 'summarize_document', + 'function_result': { + 'document': {'id': 'doc-9', 'file_name': 'Legacy.pdf'}, + 'summary': 'A summary.', + }, + 'success': True, + }, + ]) + if len(summary_citations) != 1: + print(f'❌ Expected 1 summary citation, found {len(summary_citations)}') + return False + if summary_citations[0].get('citation_id') == 'doc-9_1': + print('❌ Summary citation synthesized a chunk locator that may not exist') + return False + if summary_citations[0].get('citation_id') != 'doc-9': + print( + '❌ Summary fallback citation id should be the document id, found ' + f'{summary_citations[0].get("citation_id")!r}' + ) + return False + + print('✅ Zero sequences survive and summaries do not invent locators') + return True + except Exception as e: + print(f'❌ Test failed: {e}') + import traceback + traceback.print_exc() + return False + + +def test_plugin_invocations_are_supported_for_cancelled_streams(): + """Raw invocations must also produce citations for cancelled or errored streams.""" + print('🔍 Validating plugin invocation derivation...') + try: + class StubInvocation: + def __init__(self): + self.plugin_name = 'DocumentSearchPlugin' + self.function_name = 'search_documents' + self.success = True + self.result = { + 'results': [{ + 'id': 'doc-7_2', + 'document_id': 'doc-7', + 'file_name': 'Draft.pdf', + 'page_number': 2, + 'chunk_id': 'chunk-2', + }], + } + + invocation = StubInvocation() + + # A cancelled stream has invocations but an empty agent citation list. + hybrid_citations = [] + added_count = agent_document_citations.apply_agent_document_citations( + hybrid_citations, + [], + plugin_invocations=[invocation], + ) + if added_count != 1 or len(hybrid_citations) != 1: + print(f'❌ Expected 1 citation from raw invocations, found {added_count}') + return False + if hybrid_citations[0].get('citation_id') != 'doc-7_2': + print('❌ Invocation-derived citation carried the wrong id') + return False + + # Passing both sources must not double count the same chunk. + both_sources = [] + agent_document_citations.apply_agent_document_citations( + both_sources, + [build_search_agent_citation([{ + 'id': 'doc-7_2', + 'document_id': 'doc-7', + 'file_name': 'Draft.pdf', + 'page_number': 2, + 'chunk_id': 'chunk-2', + }])], + plugin_invocations=[invocation], + ) + if len(both_sources) != 1: + print(f'❌ Expected dedupe across both sources, found {len(both_sources)}') + return False + + print('✅ Raw invocations are supported and deduplicated') + return True + except Exception as e: + print(f'❌ Test failed: {e}') + import traceback + traceback.print_exc() + return False + + +def test_non_document_and_failed_invocations_are_ignored(): + """Only successful document-search invocations may produce document citations.""" + print('🔍 Validating invocation filtering...') + try: + ignored_citations = agent_document_citations.build_document_citations_from_agent_citations([ + { + 'plugin_name': 'SmartHttpPlugin', + 'function_name': 'search_documents', + 'function_result': {'results': [{'id': 'x_1', 'document_id': 'x', 'file_name': 'a.pdf'}]}, + 'success': True, + }, + { + 'plugin_name': 'DocumentSearchPlugin', + 'function_name': 'search_documents', + 'function_result': {'results': [{'id': 'y_1', 'document_id': 'y', 'file_name': 'b.pdf'}]}, + 'success': False, + }, + { + 'plugin_name': 'DocumentSearchPlugin', + 'function_name': 'search_documents', + 'function_result': {'error': 'Access denied'}, + 'success': True, + }, + ]) + if ignored_citations: + print(f'❌ Expected no derived citations, found {len(ignored_citations)}') + return False + + print('✅ Unrelated, failed, and errored invocations are ignored') + return True + except Exception as e: + print(f'❌ Test failed: {e}') + import traceback + traceback.print_exc() + return False + + +def test_merge_deduplicates_without_truncating(): + """Merging must dedupe against route citations and never cap the source list.""" + print('🔍 Validating merge behavior...') + try: + existing_citations = [{ + 'file_name': 'Policy.pdf', + 'document_id': 'doc-1', + 'citation_id': 'doc-1_3', + 'page_number': 3, + 'chunk_id': 'chunk-3', + }] + + large_result_set = [ + { + 'id': f'doc-1_{index}', + 'document_id': 'doc-1', + 'file_name': 'Policy.pdf', + 'page_number': index, + 'chunk_id': f'chunk-{index}', + } + for index in range(1, 501) + ] + added_count = agent_document_citations.apply_agent_document_citations( + existing_citations, + [build_search_agent_citation(large_result_set)], + ) + + if added_count != 499: + print(f'❌ Expected 499 added citations, found {added_count}') + return False + if len(existing_citations) != 500: + print(f'❌ Expected 500 total source citations, found {len(existing_citations)}') + return False + + citation_ids = [citation.get('citation_id') for citation in existing_citations] + if len(citation_ids) != len(set(citation_ids)): + print('❌ Merged citations contain duplicates') + return False + + route_citation = next( + citation for citation in existing_citations + if citation.get('citation_id') == 'doc-1_3' + ) + if route_citation.get('source') == agent_document_citations.AGENT_DOCUMENT_CITATION_SOURCE: + print('❌ Existing route-level citation was overwritten by the agent-derived copy') + return False + + print('✅ Merge dedupes, preserves route citations, and does not truncate') + return True + except Exception as e: + print(f'❌ Test failed: {e}') + import traceback + traceback.print_exc() + return False + + +def test_inline_markers_promote_documents_into_cited_references(): + """Payload citation markers must be matched by the citation tracker.""" + print('🔍 Validating inline citation markers...') + try: + payload = agent_document_citations.annotate_document_search_payload( + { + 'results': [ + { + 'id': 'doc-1_3', + 'document_id': 'doc-1', + 'file_name': 'Policy.pdf', + 'page_number': 3, + 'chunk_id': 'chunk-3', + }, + { + 'id': 'doc-2_7', + 'document_id': 'doc-2', + 'file_name': 'Other.pdf', + 'page_number': 7, + 'chunk_id': 'chunk-7', + }, + ], + }, + 'search_documents', + ) + + marker = payload['results'][0].get('citation') + if marker != '(Source: Policy.pdf, Page: 3) [#doc-1_3]': + print(f'❌ Unexpected citation marker: {marker!r}') + return False + if not payload.get('citation_instructions'): + print('❌ Annotated payload is missing citation instructions') + return False + + hybrid_citations = [] + agent_document_citations.apply_agent_document_citations( + hybrid_citations, + [{ + 'plugin_name': 'DocumentSearchPlugin', + 'function_name': 'search_documents', + 'function_result': payload, + 'success': True, + }], + ) + if len(hybrid_citations) != 2: + print(f'❌ Expected 2 source citations, found {len(hybrid_citations)}') + return False + + citation_tracking = build_cited_source_subsets( + f'Double dipping is prohibited {marker}', + hybrid_citations=hybrid_citations, + web_search_citations=[], + ) + cited_citations = citation_tracking.get('cited_hybrid_citations') or [] + if len(cited_citations) != 1: + print(f'❌ Expected 1 cited reference, found {len(cited_citations)}') + return False + if cited_citations[0].get('document_id') != 'doc-1': + print('❌ The wrong document was promoted into cited references') + return False + + print('✅ Markers promote agent-discovered documents into cited references') + return True + except Exception as e: + print(f'❌ Test failed: {e}') + import traceback + traceback.print_exc() + return False + + +def test_sheet_and_json_payload_handling(): + """Tabular sheets and JSON-string payloads must be handled.""" + print('🔍 Validating tabular and JSON payload handling...') + try: + sheet_payload = agent_document_citations.annotate_document_search_payload( + { + 'results': [{ + 'id': 'doc-4_1', + 'document_id': 'doc-4', + 'file_name': 'Budget.xlsx', + 'sheet_name': 'Q1', + 'page_number': 1, + }], + }, + 'search_documents', + ) + sheet_marker = sheet_payload['results'][0].get('citation') + if sheet_marker != '(Source: Budget.xlsx, Sheet: Q1) [#doc-4_1]': + print(f'❌ Unexpected tabular citation marker: {sheet_marker!r}') + return False + + json_citations = agent_document_citations.build_document_citations_from_agent_citations([ + { + 'plugin_name': 'DocumentSearchPlugin', + 'function_name': 'search_documents', + 'function_result': ( + '{"results": [{"id": "doc-5_2", "document_id": "doc-5", ' + '"file_name": "Notes.pdf", "page_number": 2}]}' + ), + 'success': True, + }, + ]) + if len(json_citations) != 1 or json_citations[0].get('citation_id') != 'doc-5_2': + print('❌ JSON-string plugin results were not parsed into citations') + return False + + print('✅ Tabular sheets and JSON-string payloads are handled') + return True + except Exception as e: + print(f'❌ Test failed: {e}') + import traceback + traceback.print_exc() + return False + + +def test_chat_and_workflow_paths_apply_the_helper(): + """Every assistant-message path must merge agent document citations.""" + print('🔍 Validating chat and workflow wiring...') + try: + chat_source = read_source('application', 'single_app', 'route_backend_chats.py') + workflow_source = read_source('application', 'single_app', 'functions_workflow_runner.py') + plugin_source = read_source( + 'application', 'single_app', 'semantic_kernel_plugins', 'document_search_plugin.py' + ) + + chat_call_count = chat_source.count('apply_agent_document_citations(') + if chat_call_count < 5: + print( + '❌ Expected apply_agent_document_citations in the document action, ' + f'non-streaming, and all three streaming branches, found {chat_call_count} calls' + ) + return False + if 'from functions_agent_document_citations import apply_agent_document_citations' not in chat_source: + print('❌ route_backend_chats.py does not import the shared helper') + return False + + if 'apply_agent_document_citations(' not in workflow_source: + print('❌ functions_workflow_runner.py does not merge agent document citations') + return False + + annotate_count = plugin_source.count('annotate_document_search_payload(') + if annotate_count < 3: + print( + '❌ Expected the document search plugin to annotate all three functions, ' + f'found {annotate_count} call sites' + ) + return False + + # Cancelled and errored streams must supply raw invocations, because streaming + # invocations only reach the agent citation list after a stream completes. + invocation_call_count = chat_source.count( + 'plugin_invocations=_get_current_message_plugin_invocations(' + ) + if invocation_call_count < 4: + print( + '❌ Expected all four agent chat paths to pass raw plugin invocations, ' + f'found {invocation_call_count}' + ) + return False + + print('✅ Chat, workflow, and plugin wiring are in place') + return True + except Exception as e: + print(f'❌ Test failed: {e}') + import traceback + traceback.print_exc() + return False + + +def collect_direct_call_lines(function_node, call_names): + """Return line numbers of calls made directly in a function, skipping nested defs.""" + call_lines = {call_name: [] for call_name in call_names} + nested_function_types = (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda) + + def visit(node): + for child in ast.iter_child_nodes(node): + if isinstance(child, nested_function_types): + continue + if isinstance(child, ast.Call): + called = child.func + called_name = getattr(called, 'id', None) or getattr(called, 'attr', None) + if called_name in call_lines: + call_lines[called_name].append(child.lineno) + visit(child) + + for statement in function_node.body: + if isinstance(statement, nested_function_types): + continue + visit(statement) + + return call_lines + + +NESTED_FUNCTION_TYPES = (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda) + +BLOCK_FIELDS = ('body', 'orelse', 'finalbody', 'handlers') + + +def collect_statement_own_calls(statement, call_names): + """Return (line, name) for calls in a statement, excluding its nested blocks.""" + own_calls = [] + nested_block_nodes = set() + + for field_name in BLOCK_FIELDS: + for child in getattr(statement, field_name, None) or []: + nested_block_nodes.add(id(child)) + + def visit(node): + for child in ast.iter_child_nodes(node): + if isinstance(child, NESTED_FUNCTION_TYPES) or id(child) in nested_block_nodes: + continue + if isinstance(child, ast.Call): + called = child.func + called_name = getattr(called, 'id', None) or getattr(called, 'attr', None) + if called_name in call_names: + own_calls.append((child.lineno, called_name)) + visit(child) + + visit(statement) + return sorted(own_calls) + + +def get_statement_blocks(statement): + """Return the nested statement blocks a compound statement introduces.""" + blocks = [] + for field_name in BLOCK_FIELDS: + block = getattr(statement, field_name, None) + if not block: + continue + for entry in block: + if isinstance(entry, ast.ExceptHandler): + blocks.append(entry.body) + else: + blocks.append(block) + break + return blocks + + +def find_unmerged_snapshot(statements, merge_name, snapshot_names, merge_seen=False): + """Return the first snapshot line reachable without a preceding merge on that path. + + Merges inside a branch do not leak to sibling branches or to code after the branch, + so a merge moved into one finalization branch cannot vouch for another. + """ + for statement in statements: + if isinstance(statement, NESTED_FUNCTION_TYPES): + continue + + for line, called_name in collect_statement_own_calls( + statement, + [merge_name] + snapshot_names, + ): + if called_name == merge_name: + merge_seen = True + elif not merge_seen: + return line + + for block in get_statement_blocks(statement): + unmerged_line = find_unmerged_snapshot( + block, + merge_name, + snapshot_names, + merge_seen=merge_seen, + ) + if unmerged_line is not None: + return unmerged_line + + return None + + +def test_merge_runs_before_tracking_and_persistence(): + """The merge must precede cited-subset tracking and persistence on every branch.""" + print('🔍 Validating merge ordering...') + try: + merge_name = 'apply_agent_document_citations' + tracking_name = 'build_cited_source_subsets' + persistence_names = [ + 'persist_agent_citation_artifacts', + '_persist_agent_citation_artifacts', + ] + snapshot_names = [tracking_name] + persistence_names + + module_targets = [ + ('route_backend_chats.py', ['application', 'single_app', 'route_backend_chats.py']), + ( + 'functions_workflow_runner.py', + ['application', 'single_app', 'functions_workflow_runner.py'], + ), + ] + + for module_label, module_parts in module_targets: + module_ast = ast.parse(read_source(*module_parts)) + verified_function_count = 0 + + for node in ast.walk(module_ast): + if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + continue + + call_lines = collect_direct_call_lines(node, [merge_name] + snapshot_names) + if not call_lines.get(tracking_name): + # Functions that only mirror an already-tracked assistant document + # copy its citations and must not re-merge. + continue + + unmerged_line = find_unmerged_snapshot(node.body, merge_name, snapshot_names) + if unmerged_line is not None: + print( + f'❌ {module_label}:{node.lineno} function {node.name!r} snapshots ' + f'citations at line {unmerged_line} without a merge on that branch' + ) + return False + + verified_function_count += 1 + + if verified_function_count == 0: + print(f'❌ No cited-subset functions were found in {module_label}') + return False + + print(f' {module_label}: {verified_function_count} function(s) verified') + + print('✅ Merges run before cited-subset tracking and persistence on every branch') + return True + except Exception as e: + print(f'❌ Test failed: {e}') + import traceback + traceback.print_exc() + return False + + +def test_large_source_lists_collapse_in_the_ui(): + """The sources panel must collapse large source lists instead of capping data.""" + print('🔍 Validating source list rendering...') + try: + messages_source = read_source( + 'application', 'single_app', 'static', 'js', 'chat', 'chat-messages.js' + ) + citations_source = read_source( + 'application', 'single_app', 'static', 'js', 'chat', 'chat-citations.js' + ) + + required_message_snippets = [ + 'const DOCUMENT_CITATION_VISIBLE_LIMIT', + 'function buildDocumentCitationGroupHtml(', + 'citation-overflow-group d-none', + 'citation-overflow-toggle', + ] + for snippet in required_message_snippets: + if snippet not in messages_source: + print(f'❌ chat-messages.js is missing {snippet!r}') + return False + + if 'button.citation-overflow-toggle' not in citations_source: + print('❌ chat-citations.js does not handle the source overflow toggle') + return False + if 'function toggleCitationOverflowGroup(' not in citations_source: + print('❌ chat-citations.js is missing the overflow toggle handler') + return False + + print('✅ Large source lists collapse behind a show-more control') + return True + except Exception as e: + print(f'❌ Test failed: {e}') + import traceback + traceback.print_exc() + return False + + +if __name__ == '__main__': + tests = [ + test_version_is_at_least_implementation_version, + test_search_results_become_document_citations, + test_chunk_and_summary_functions_produce_citations, + test_zero_indexed_and_missing_locators_are_not_faked, + test_plugin_invocations_are_supported_for_cancelled_streams, + test_non_document_and_failed_invocations_are_ignored, + test_merge_deduplicates_without_truncating, + test_inline_markers_promote_documents_into_cited_references, + test_sheet_and_json_payload_handling, + test_chat_and_workflow_paths_apply_the_helper, + test_merge_runs_before_tracking_and_persistence, + test_large_source_lists_collapse_in_the_ui, + ] + + results = [] + for test in tests: + print(f'\n🧪 Running {test.__name__}...') + results.append(test()) + + print(f'\n📊 Results: {sum(results)}/{len(results)} tests passed') + sys.exit(0 if all(results) else 1)