Remove max subqueries option, fix sort order

pamelafox · pamelafox · commit ccc28e5fff93 · 2025-09-09T11:50:41.000-07:00
diff --git a/app/backend/approaches/approach.py b/app/backend/approaches/approach.py
@@ -10,6 +10,7 @@
     KnowledgeAgentMessageTextContent,
     KnowledgeAgentRetrievalRequest,
     KnowledgeAgentRetrievalResponse,
+    KnowledgeAgentSearchIndexActivityRecord,
     KnowledgeAgentSearchIndexReference,
     SearchIndexKnowledgeSourceParams,
 )
@@ -287,21 +288,37 @@ async def run_agentic_retrieval(
                 ],
                 knowledge_source_params=[
                     SearchIndexKnowledgeSourceParams(
-                        knowledge_source_name="default-knowledge-source",
+                        knowledge_source_name=search_index_name,
                         filter_add_on=filter_add_on,
                     )
                 ],
             )
         )
 
+        # Map activity id -> agent's internal search query
+        activities = response.activity
+        activity_mapping: dict[int, str] = (
+            {
+                activity.id: activity.search_index_arguments.search
+                for activity in activities
+                if (
+                    isinstance(activity, KnowledgeAgentSearchIndexActivityRecord)
+                    and activity.search_index_arguments
+                    and activity.search_index_arguments.search is not None
+                )
+            }
+            if activities
+            else {}
+        )
+
         # No refs? we're done
         if not (response and response.references):
             return response, []
 
         # Extract references
         refs = [r for r in response.references if isinstance(r, KnowledgeAgentSearchIndexReference)]
-
         documents: list[Document] = []
+        doc_to_ref_id: dict[str, str] = {}
 
         # Create documents from reference source data
         for ref in refs:
@@ -317,11 +334,18 @@ async def run_agentic_retrieval(
                         groups=ref.source_data.get("groups"),
                         reranker_score=ref.reranker_score,
                         images=ref.source_data.get("images"),
+                        search_agent_query=activity_mapping[ref.activity_source],
                     )
                 )
+                doc_to_ref_id[ref.source_data.get("id")] = ref.id
                 if top and len(documents) >= top:
                     break
 
+        if results_merge_strategy == "interleaved":
+            documents = sorted(
+                documents,
+                key=lambda d: int(doc_to_ref_id.get(d.id, 0)) if d.id and doc_to_ref_id.get(d.id) else 0,
+            )
         return response, documents
 
     async def get_sources_content(
diff --git a/app/backend/prepdocslib/searchmanager.py b/app/backend/prepdocslib/searchmanager.py
@@ -470,7 +470,7 @@ async def create_agent(self):
                 field_names.append("images")
             async with self.search_info.create_search_index_client() as search_index_client:
                 knowledge_source = SearchIndexKnowledgeSource(
-                    name="default-knowledge-source",
+                    name=self.search_info.index_name,  # Use the same name for convenience
                     description="Default knowledge source using the main search index",
                     search_index_parameters=SearchIndexKnowledgeSourceParameters(
                         search_index_name=self.search_info.index_name,
diff --git a/app/frontend/src/api/models.ts b/app/frontend/src/api/models.ts
@@ -14,7 +14,6 @@ export type ChatAppRequestOverrides = {
     exclude_category?: string;
     seed?: number;
     top?: number;
-    max_subqueries?: number;
     results_merge_strategy?: string;
     temperature?: number;
     minimum_search_score?: number;
diff --git a/app/frontend/src/components/AnalysisPanel/AgentPlan.tsx b/app/frontend/src/components/AnalysisPanel/AgentPlan.tsx
@@ -8,16 +8,16 @@ SyntaxHighlighter.registerLanguage("json", json);
 
 type ModelQueryPlanningStep = {
     id: number;
-    type: "ModelQueryPlanning";
+    type: "modelQueryPlanning";
     input_tokens: number;
     output_tokens: number;
 };
 
 type AzureSearchQueryStep = {
     id: number;
-    type: "AzureSearchQuery";
-    target_index: string;
-    query: { search: string };
+    type: "searchIndex";
+    knowledge_source_name: string;
+    search_index_arguments: { search: string };
     query_time: string;
     count: number;
     elapsed_ms: number;
@@ -32,10 +32,10 @@ interface Props {
 
 export const AgentPlan: React.FC<Props> = ({ query_plan, description }) => {
     // find the planning step
-    const planning = query_plan.find((step): step is ModelQueryPlanningStep => step.type === "ModelQueryPlanning");
+    const planning = query_plan.find((step): step is ModelQueryPlanningStep => step.type === "modelQueryPlanning");
 
     // collect all search query steps
-    const queries = query_plan.filter((step): step is AzureSearchQueryStep => step.type === "AzureSearchQuery");
+    const queries = query_plan.filter((step): step is AzureSearchQueryStep => step.type === "searchIndex");
 
     return (
         <div>
@@ -65,7 +65,7 @@ export const AgentPlan: React.FC<Props> = ({ query_plan, description }) => {
                     <tbody>
                         {queries.map(q => (
                             <tr key={q.id}>
-                                <td>{q.query.search}</td>
+                                <td>{q.search_index_arguments.search}</td>
                                 <td>{q.count}</td>
                                 <td>{q.elapsed_ms}</td>
                             </tr>
diff --git a/app/frontend/src/components/Settings/Settings.tsx b/app/frontend/src/components/Settings/Settings.tsx
@@ -13,7 +13,6 @@ export interface SettingsProps {
     promptTemplate: string;
     temperature: number;
     retrieveCount: number;
-    maxSubqueryCount: number;
     resultsMergeStrategy: string;
     seed: number | null;
     minimumSearchScore: number;
@@ -55,7 +54,6 @@ export const Settings = ({
     promptTemplate,
     temperature,
     retrieveCount,
-    maxSubqueryCount,
     resultsMergeStrategy,
     seed,
     minimumSearchScore,
@@ -109,8 +107,6 @@ export const Settings = ({
     const rerankerScoreFieldId = useId("rerankerScoreField");
     const retrieveCountId = useId("retrieveCount");
     const retrieveCountFieldId = useId("retrieveCountField");
-    const maxSubqueryCountId = useId("maxSubqueryCount");
-    const maxSubqueryCountFieldId = useId("maxSubqueryCountField");
     const resultsMergeStrategyFieldId = useId("resultsMergeStrategy");
     const includeCategoryId = useId("includeCategory");
     const includeCategoryFieldId = useId("includeCategoryField");
@@ -209,20 +205,6 @@ export const Settings = ({
                     onRenderLabel={props => renderLabel(props, rerankerScoreId, rerankerScoreFieldId, t("helpTexts.rerankerScore"))}
                 />
             )}
-            {showAgenticRetrievalOption && useAgenticRetrieval && (
-                <TextField
-                    id={maxSubqueryCountFieldId}
-                    className={styles.settingsSeparator}
-                    label={t("labels.maxSubqueryCount")}
-                    type="number"
-                    min={2}
-                    max={40}
-                    defaultValue={maxSubqueryCount.toString()}
-                    onChange={(_ev, val) => onChange("maxSubqueryCount", parseInt(val || "10"))}
-                    aria-labelledby={maxSubqueryCountId}
-                    onRenderLabel={props => renderLabel(props, maxSubqueryCountId, maxSubqueryCountFieldId, t("helpTexts.maxSubqueryCount"))}
-                />
-            )}
             {showAgenticRetrievalOption && useAgenticRetrieval && (
                 <Dropdown
                     id={resultsMergeStrategyFieldId}
diff --git a/app/frontend/src/locales/en/translation.json b/app/frontend/src/locales/en/translation.json
@@ -83,7 +83,6 @@
         "minimumSearchScore": "Minimum search score",
         "minimumRerankerScore": "Minimum reranker score",
         "retrieveCount": "Retrieve this many search results:",
-        "maxSubqueryCount": "Max subqueries",
         "resultsMergeStrategy": "Results merge strategy",
         "resultsMergeStrategyOptions": {
             "interleaved": "Interleaved",
@@ -144,8 +143,6 @@
              "Sets a minimum score for search results coming back from the semantic reranker. The score always ranges between 0-4. The higher the score, the more semantically relevant the result is to the question.",
         "retrieveNumber":
              "Sets the number of search results to retrieve from Azure AI search. More results may increase the likelihood of finding the correct answer, but may lead to the model getting 'lost in the middle'.",
-        "maxSubqueryCount":
-            "Sets the maximum number of subqueries to use for the agentic retrieval query plan.",
         "resultsMergeStrategy":
             "Sets the strategy for merging results from multiple subqueries. Interleaved picks the top results from each subquery in round robin order, while descending sorts all results by reranker score.",
         "includeCategory":
diff --git a/app/frontend/src/locales/ja/translation.json b/app/frontend/src/locales/ja/translation.json
@@ -83,7 +83,6 @@
         "minimumSearchScore": "最小検索スコア",
         "minimumRerankerScore": "最小リランキング・スコア",
         "retrieveCount": "ここで指定する検索結果数を取得：",
-        "maxSubqueryCount": "最大サブクエリ数",
         "resultsMergeStrategy": "結果マージ戦略",
         "resultsMergeStrategyOptions": {
             "interleaved": "インターリーブ",
@@ -144,8 +143,6 @@
             "セマンティック・リランカーから返される検索結果の最小スコアを設定します。スコアの値は0から4の範囲で変更できます。スコアの値が大きいほど、質問に対する結果の意味的な関連性が高まります。",
         "retrieveNumber":
             "Azure AI Searchの検索結果から取得する数を設定します。結果が多ければ多いほど、正しい答えを見つける可能性は高まるかもしれませんが、モデルが「途中で迷子になる」可能性もあります。",
-        "maxSubqueryCount":
-            "エージェント型リトリーバルのクエリプランで使用する最大サブクエリ数を設定します。",
         "resultsMergeStrategy":
             "複数のサブクエリからの結果をマージする戦略を設定します。インターリーブは各サブクエリから上位結果をラウンドロビン順で選択し、降順はすべての結果をリランカースコアで並び替えます。",
         "includeCategory": "検索結果に含めるカテゴリを指定します。デフォルトのデータセットはカテゴリを使用していません。",
diff --git a/app/frontend/src/pages/ask/Ask.tsx b/app/frontend/src/pages/ask/Ask.tsx
@@ -30,7 +30,6 @@ export function Component(): JSX.Element {
     const [minimumSearchScore, setMinimumSearchScore] = useState<number>(0);
     const [retrievalMode, setRetrievalMode] = useState<RetrievalMode>(RetrievalMode.Hybrid);
     const [retrieveCount, setRetrieveCount] = useState<number>(3);
-    const [maxSubqueryCount, setMaxSubqueryCount] = useState<number>(10);
     const [resultsMergeStrategy, setResultsMergeStrategy] = useState<string>("interleaved");
     const [useSemanticRanker, setUseSemanticRanker] = useState<boolean>(true);
     const [useSemanticCaptions, setUseSemanticCaptions] = useState<boolean>(false);
@@ -149,7 +148,6 @@ export function Component(): JSX.Element {
                         include_category: includeCategory.length === 0 ? undefined : includeCategory,
                         exclude_category: excludeCategory.length === 0 ? undefined : excludeCategory,
                         top: retrieveCount,
-                        max_subqueries: maxSubqueryCount,
                         results_merge_strategy: resultsMergeStrategy,
                         temperature: temperature,
                         minimum_reranker_score: minimumRerankerScore,
@@ -209,9 +207,6 @@ export function Component(): JSX.Element {
             case "retrieveCount":
                 setRetrieveCount(value);
                 break;
-            case "maxSubqueryCount":
-                setMaxSubqueryCount(value);
-                break;
             case "resultsMergeStrategy":
                 setResultsMergeStrategy(value);
                 break;
@@ -370,7 +365,6 @@ export function Component(): JSX.Element {
                     promptTemplateSuffix={promptTemplateSuffix}
                     temperature={temperature}
                     retrieveCount={retrieveCount}
-                    maxSubqueryCount={maxSubqueryCount}
                     resultsMergeStrategy={resultsMergeStrategy}
                     seed={seed}
                     minimumSearchScore={minimumSearchScore}
diff --git a/app/frontend/src/pages/chat/Chat.tsx b/app/frontend/src/pages/chat/Chat.tsx
@@ -35,7 +35,6 @@ const Chat = () => {
     const [minimumRerankerScore, setMinimumRerankerScore] = useState<number>(0);
     const [minimumSearchScore, setMinimumSearchScore] = useState<number>(0);
     const [retrieveCount, setRetrieveCount] = useState<number>(3);
-    const [maxSubqueryCount, setMaxSubqueryCount] = useState<number>(10);
     const [resultsMergeStrategy, setResultsMergeStrategy] = useState<string>("interleaved");
     const [retrievalMode, setRetrievalMode] = useState<RetrievalMode>(RetrievalMode.Hybrid);
     const [useSemanticRanker, setUseSemanticRanker] = useState<boolean>(true);
@@ -213,7 +212,6 @@ const Chat = () => {
                         include_category: includeCategory.length === 0 ? undefined : includeCategory,
                         exclude_category: excludeCategory.length === 0 ? undefined : excludeCategory,
                         top: retrieveCount,
-                        max_subqueries: maxSubqueryCount,
                         results_merge_strategy: resultsMergeStrategy,
                         temperature: temperature,
                         minimum_reranker_score: minimumRerankerScore,
@@ -310,9 +308,6 @@ const Chat = () => {
             case "retrieveCount":
                 setRetrieveCount(value);
                 break;
-            case "maxSubqueryCount":
-                setMaxSubqueryCount(value);
-                break;
             case "resultsMergeStrategy":
                 setResultsMergeStrategy(value);
                 break;
@@ -543,7 +538,6 @@ const Chat = () => {
                         promptTemplate={promptTemplate}
                         temperature={temperature}
                         retrieveCount={retrieveCount}
-                        maxSubqueryCount={maxSubqueryCount}
                         resultsMergeStrategy={resultsMergeStrategy}
                         seed={seed}
                         minimumSearchScore={minimumSearchScore}
diff --git a/docs/agentic_retrieval.md b/docs/agentic_retrieval.md
@@ -54,13 +54,7 @@ See the agentic retrieval documentation.
 
    Open the web app and start a new chat. Agentic retrieval will be used to find all sources.
 
-6. **Experiment with max subqueries:**
-
-   Select the developer options in the web app and change max subqueries to any value between 1 and 20. This controls the maximum amount of subqueries that can be created in the query plan.
-
-   ![Max subqueries screenshot](./images/max-subqueries.png)
-
-7. **Review the query plan**
+6. **Review the query plan**
 
    Agentic retrieval use additional billed tokens behind the scenes for the planning process.
    To see the token usage, select the lightbulb icon on a chat answer. This will open the "Thought process" tab, which shows the amount of tokens used by and the queries produced by the planning process
diff --git a/evals/evaluate_config.json b/evals/evaluate_config.json
@@ -6,7 +6,6 @@
     "target_parameters": {
         "overrides": {
             "top": 3,
-            "max_subqueries": 10,
             "results_merge_strategy": "interleaved",
             "temperature": 0.3,
             "minimum_reranker_score": 0,
diff --git a/evals/evaluate_config_multimodal.json b/evals/evaluate_config_multimodal.json
@@ -6,7 +6,6 @@
     "target_parameters": {
         "overrides": {
             "top": 3,
-            "max_subqueries": 10,
             "results_merge_strategy": "interleaved",
             "temperature": 0.3,
             "minimum_reranker_score": 0,
diff --git a/infra/main.bicep b/infra/main.bicep
@@ -24,7 +24,7 @@ param searchServiceLocation string = '' // Set in main.parameters.json
 @allowed(['free', 'basic', 'standard', 'standard2', 'standard3', 'storage_optimized_l1', 'storage_optimized_l2'])
 param searchServiceSkuName string // Set in main.parameters.json
 param searchIndexName string // Set in main.parameters.json
-param searchAgentName string = useAgenticRetrieval ? '${searchIndexName}-agent' : ''
+param searchAgentName string = useAgenticRetrieval ? '${searchIndexName}-agent-upgrade' : ''
 param searchQueryLanguage string // Set in main.parameters.json
 param searchQuerySpeller string // Set in main.parameters.json
 param searchServiceSemanticRankerLevel string // Set in main.parameters.json
