Azure AI Search Agentic Retrieval Returning More than Just Selected Fields

[Wilsven]

• This request is referencing PR Add support for agentic retrieval #2537 @mattgotteiner
• Fields below will be pertaining to my data structure indexed into the search index
• Within the source_data, I realised agentic retrieval is only returning the id, title and content fields in the references (see image below).
• Is there currently a way to retrieve more than just these fields?
  • Within my index, I have more fields like full_url which I want to retrieve so when I display the sources on the frontend, user may click into it and a new browser tab will be opened and they will be directed to the full_url
• It is my understanding that the fields returned are controlled within the searchmanager.py module? (see second screenshot below)
  • It seems my title and content fields are listed as prioritized_fields which may be why they are only returned
  • The other remaining unset field is the keywords_field
  • Now, if this is where I define what the agentic retrieval returns, where should a field like full_url go? Surely does not make sense to place it within any of the three prioritized_fields options — title_field, content_field or keywords_field

[mattgotteiner]

Thanks for raising this issue - currently, only semantic configuration fields are returned. In a future release of agentic retrieval we will address this gap

In the meantime - unfortunately you'll have to retrieve fields you are missing by fetching individual documents or searching for them

[cforce]

Voting for this

[taylorn-ai]

I too was a little frustrated with this, so I created a new method to retrieve the extra fields, it doesn't seem to effect performance. For anyone interested...

    async def run_agentic_retrieval(
        self,
        messages: list[ChatCompletionMessageParam],
        agent_client: KnowledgeAgentRetrievalClient,
        search_index_name: str,
        top: int | None = None,
        filter_add_on: str | None = None,
        minimum_reranker_score: float | None = None,
        max_docs_for_reranker: int | None = None,
        results_merge_strategy: str | None = None,
    ) -> tuple[KnowledgeAgentRetrievalResponse, list[Document]]:
        response = await agent_client.retrieve(
            retrieval_request=KnowledgeAgentRetrievalRequest(
                messages=[
                    KnowledgeAgentMessage(
                        role=str(msg["role"]),
                        content=[KnowledgeAgentMessageTextContent(text=str(msg["content"]))],
                    )
                    for msg in messages
                    if msg["role"] != "system"
                ],
                target_index_params=[
                    KnowledgeAgentIndexParams(
                        index_name=search_index_name,
                        reranker_threshold=minimum_reranker_score,
                        max_docs_for_reranker=max_docs_for_reranker,
                        filter_add_on=filter_add_on,
                        include_reference_source_data=True,
                    )
                ],
            )
        )

        # Map activity id -> agent's internal search query
        activities = response.activity
        activity_mapping: dict[int, str] = (
            {
                activity.id: activity.query.search
                for activity in activities
                if (
                    isinstance(activity, KnowledgeAgentSearchActivityRecord)
                    and activity.query
                    and activity.query.search is not None
                )
            }
            if activities
            else {}
        )

        # No refs? we're done
        if not (response and response.references):
            return response, []

        if results_merge_strategy == "interleaved":  # Use interleaved reference order
            ordered_refs = sorted(
                (r for r in response.references if isinstance(r, KnowledgeAgentAzureSearchDocReference)),
                key=lambda r: int(r.id) if getattr(r, "id", None) is not None else 0,
            )
        else:  # Default to descending strategy
            ordered_refs = [r for r in response.references if isinstance(r, KnowledgeAgentAzureSearchDocReference)]

        documents = await self.hydrate_agent_references(
            references=ordered_refs, top=top, activity_mapping=activity_mapping
        )

        return response, documents

    async def hydrate_agent_references(
        self,
        references: list[KnowledgeAgentAzureSearchDocReference],
        top: int | None,
        activity_mapping: dict[int, str],
    ) -> list[Document]:
        doc_keys: list[str] = []
        key_to_agent_query: dict[str, str] = {}

        for ref in references:
            if not ref.doc_key:
                continue
            doc_keys.append(ref.doc_key)
            key_to_agent_query[ref.doc_key] = activity_mapping.get(ref.activity_source, "")
            if top and len(doc_keys) >= top:
                break

        if not doc_keys:
            return []

        # Build search filter only on doc IDs
        escaped = [k.replace("'", "''") for k in doc_keys]
        id_csv = ",".join(escaped)
        id_filter = f"search.in(id, '{id_csv}', ',')"

        # Fetch full documents
        hydrated_docs: list[Document] = await self.search(
            top=len(doc_keys),
            query_text=None,
            search_filter=id_filter,
            vectors=[],
            use_text_search=False,
            use_vector_search=False,
            use_semantic_ranker=False,
            use_semantic_captions=False,
            minimum_search_score=None,
            minimum_reranker_score=None,
            use_query_rewriting=False,
        )

        # Preserve agent order
        order_index = {k: i for i, k in enumerate(doc_keys)}
        hydrated_docs.sort(key=lambda d: order_index.get(d.id or "", 10**9))

        # Inject the agent's internal query string
        for d in hydrated_docs:
            d.search_agent_query = key_to_agent_query.get(d.id or "", "")

        return hydrated_docs

[cforce]

@taylorn-ai Please add a feature flag, add tests and create a merge request for this

[taylorn-ai]

@taylorn-ai Please add a feature flag, add tests and create a merge request for this

I will try to get to this tomorrow. Where do I send the bill? 😂

[taylorn-ai]

Done #2696. I considered adding another arg for which fields to retrieve (i.e. if the user only wants to return certain fields) but this seemed like it could be very fragile.