Skip to content

Selectors

Define selector abstractions and concrete selection strategies for retrieval results.

This module provides selector components used in the terminology-mapping pipeline to choose the best candidate from a set of retrieved expressions. It includes a simple rule-based selector that picks the first result under a distance threshold, as well as LLM-based selectors backed by OpenAI and Gemini models.

Selectors are designed to operate as pipeline stages that consume RetrieverResults and return structured SelectorResults objects.

BaseSelector

Bases: PipelineBaseClass, ABC

Define the abstract interface for selector pipeline components.

This base class establishes the contract for selectors that choose a final candidate from retrieval results. Subclasses must implement select(), while the base __call__() method provides pipeline-compatible invocation.

Source code in aatm\selectors.py 
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
class BaseSelector(PipelineBaseClass, ABC):
    """Define the abstract interface for selector pipeline components.

    This base class establishes the contract for selectors that choose a final
    candidate from retrieval results. Subclasses must implement ``select()``,
    while the base ``__call__()`` method provides pipeline-compatible
    invocation.
    """

    @abstractmethod
    def select(self, results: RetrieverResults) -> SelectorResults:
        """Select the best candidate result for each query.

        Subclasses must implement this method to examine retrieval results and
        return a structured selection output for each query.

        Args:
            results: Retrieval results containing the original queries and their
                candidate matches.

        Returns:
            A ``SelectorResults`` instance containing the selected result for
                each query.

        Raises:
            NotImplementedError: If the subclass does not override this method.
        """
        pass

    def __call__(self, results: RetrieverResults) -> SelectorResults:
        """Invoke the selector on retrieval results.

        This method makes selector instances directly callable and delegates to
        ``select()``.

        Args:
            results: Retrieval results to process.

        Returns:
            A ``SelectorResults`` instance produced by the selector.
        """
        return self.select(results)

select(results) abstractmethod

Select the best candidate result for each query.

Subclasses must implement this method to examine retrieval results and return a structured selection output for each query.

Parameters:

Name Type Description Default
results RetrieverResults

Retrieval results containing the original queries and their candidate matches.

 required

Returns:

Type Description
SelectorResults

A SelectorResults instance containing the selected result for each query.

Raises:

Type Description
NotImplementedError

If the subclass does not override this method.
Source code in aatm\selectors.py 
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
@abstractmethod
def select(self, results: RetrieverResults) -> SelectorResults:
    """Select the best candidate result for each query.

    Subclasses must implement this method to examine retrieval results and
    return a structured selection output for each query.

    Args:
        results: Retrieval results containing the original queries and their
            candidate matches.

    Returns:
        A ``SelectorResults`` instance containing the selected result for
            each query.

    Raises:
        NotImplementedError: If the subclass does not override this method.
    """
    pass

__call__(results)

Invoke the selector on retrieval results.

This method makes selector instances directly callable and delegates to select().

Parameters:

Name Type Description Default
results RetrieverResults

Retrieval results to process.

 required

Returns:

Type Description
SelectorResults

A SelectorResults instance produced by the selector.
Source code in aatm\selectors.py 
72
73
74
75
76
77
78
79
80
81
82
83
84
def __call__(self, results: RetrieverResults) -> SelectorResults:
    """Invoke the selector on retrieval results.

    This method makes selector instances directly callable and delegates to
    ``select()``.

    Args:
        results: Retrieval results to process.

    Returns:
        A ``SelectorResults`` instance produced by the selector.
    """
    return self.select(results)

FirstResultSelector

Bases: BaseSelector

Select the first retrieved result that satisfies a distance threshold.

This selector implements a simple heuristic strategy: it chooses the first candidate in each result list if its distance is below the configured threshold. If no candidate is available or the first candidate does not satisfy the threshold, an empty selection is returned.

Source code in aatm\selectors.py 
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
class FirstResultSelector(BaseSelector):
    """Select the first retrieved result that satisfies a distance threshold.

    This selector implements a simple heuristic strategy: it chooses the first
    candidate in each result list if its distance is below the configured
    threshold. If no candidate is available or the first candidate does not
    satisfy the threshold, an empty selection is returned.
    """

    def __init__(
        self, selection_threshold: float = float("inf"), *args: Any, **kwargs: Any
    ) -> None:
        """Initialize the first-result selector.

        Args:
            selection_threshold: Maximum allowed distance for accepting the
                first retrieved result. If the first result has a larger
                distance, no selection is made.
            *args: Additional positional arguments reserved for compatibility.
            **kwargs: Additional keyword arguments reserved for compatibility.

        Returns:
            None.
        """
        self.selection_threshold = selection_threshold

    def select(self, results: RetrieverResults) -> SelectorResults:
        """Select the first valid result for each query.

        This method inspects the first candidate in each query result list and
        returns it as the selected expression when its distance is below the
        configured threshold. If no candidates are available or the threshold
        is not met, an empty selection is returned.

        Args:
            results: Retrieval results containing candidate expressions for
                each query.

        Returns:
            A ``SelectorResults`` object containing either a selected
                expression or an empty selection for each query.
        """
        selected_results = []
        for idx, r in enumerate(results.results):
            try:
                if r[0].distance < self.selection_threshold:
                    selected_results.append(
                        SelectedExpressionMetadata(
                            **r[0].model_dump(), result_list_index=0
                        )
                    )
                else:
                    selected_results.append(EmptySelectionMetadata())
            except IndexError:
                selected_results.append(EmptySelectionMetadata())
                print(f"No results for query: {results.queries[idx]}. Results: {r}")
        return SelectorResults(results=selected_results, queries=results.queries)

__init__(selection_threshold=float('inf'), *args, **kwargs)

Initialize the first-result selector.

Parameters:

Name Type Description Default
selection_threshold float

Maximum allowed distance for accepting the first retrieved result. If the first result has a larger distance, no selection is made.

 float('inf')
*args Any

Additional positional arguments reserved for compatibility.

 ()
**kwargs Any

Additional keyword arguments reserved for compatibility.

 {}

Returns:

Type Description
None

None.
Source code in aatm\selectors.py 
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
def __init__(
    self, selection_threshold: float = float("inf"), *args: Any, **kwargs: Any
) -> None:
    """Initialize the first-result selector.

    Args:
        selection_threshold: Maximum allowed distance for accepting the
            first retrieved result. If the first result has a larger
            distance, no selection is made.
        *args: Additional positional arguments reserved for compatibility.
        **kwargs: Additional keyword arguments reserved for compatibility.

    Returns:
        None.
    """
    self.selection_threshold = selection_threshold

select(results)

Select the first valid result for each query.

This method inspects the first candidate in each query result list and returns it as the selected expression when its distance is below the configured threshold. If no candidates are available or the threshold is not met, an empty selection is returned.

Parameters:

Name Type Description Default
results RetrieverResults

Retrieval results containing candidate expressions for each query.

 required

Returns:

Type Description
SelectorResults

A SelectorResults object containing either a selected expression or an empty selection for each query.
Source code in aatm\selectors.py 
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
def select(self, results: RetrieverResults) -> SelectorResults:
    """Select the first valid result for each query.

    This method inspects the first candidate in each query result list and
    returns it as the selected expression when its distance is below the
    configured threshold. If no candidates are available or the threshold
    is not met, an empty selection is returned.

    Args:
        results: Retrieval results containing candidate expressions for
            each query.

    Returns:
        A ``SelectorResults`` object containing either a selected
            expression or an empty selection for each query.
    """
    selected_results = []
    for idx, r in enumerate(results.results):
        try:
            if r[0].distance < self.selection_threshold:
                selected_results.append(
                    SelectedExpressionMetadata(
                        **r[0].model_dump(), result_list_index=0
                    )
                )
            else:
                selected_results.append(EmptySelectionMetadata())
        except IndexError:
            selected_results.append(EmptySelectionMetadata())
            print(f"No results for query: {results.queries[idx]}. Results: {r}")
    return SelectorResults(results=selected_results, queries=results.queries)

OpenAILLMModels

Bases: Enum

Enumerate the supported OpenAI model identifiers for selection.

This enumeration defines the OpenAI model names that can be used by the OpenAI-based selector implementation.

Source code in aatm\selectors.py 
146
147
148
149
150
151
152
153
154
155
156
class OpenAILLMModels(Enum):
    """Enumerate the supported OpenAI model identifiers for selection.

    This enumeration defines the OpenAI model names that can be used by the
    OpenAI-based selector implementation.
    """

    GPT_52 = "gpt-5.2"
    GPT_5 = "gpt-5"
    GPT_5_NANO = "gpt-5-nano"
    GPT_5_MINI = "gpt-5-mini"

OpenAILLMSelector

Bases: BaseSelector

Select results using an OpenAI language model.

This selector formats each query and its retrieved candidate expressions into a prompt, asks an OpenAI model to choose the semantically closest expression, and validates the structured JSON response against the expected selection schema.

If the model returns no selection or selects an expression that is not present in the candidate list, the selector produces an empty selection.

Source code in aatm\selectors.py 
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
class OpenAILLMSelector(BaseSelector):
    """Select results using an OpenAI language model.

    This selector formats each query and its retrieved candidate expressions
    into a prompt, asks an OpenAI model to choose the semantically closest
    expression, and validates the structured JSON response against the expected
    selection schema.

    If the model returns no selection or selects an expression that is not
    present in the candidate list, the selector produces an empty selection.
    """

    def __init__(
        self,
        model_id: str,
        prompt_template: List[Dict[str, str]] = None,
        *args: Any,
        **kwargs: Any,
    ) -> None:
        """Initialize the OpenAI-based selector.

        This constructor resolves the configured model identifier, creates an
        OpenAI client, and sets the prompt template used to ask the model to
        choose the best candidate expression.

        Args:
            model_id: Name of the supported OpenAI model variant to use for
                selection.
            prompt_template: Optional structured prompt template. If not
                provided, a default system-and-user prompt is used.
            *args: Additional positional arguments reserved for compatibility.
            **kwargs: Additional keyword arguments reserved for compatibility.

        Returns:
            None.

        Raises:
            ValueError: If ``model_id`` is not a valid member of
                ``OpenAILLMModels``.
        """
        self.model_id = OpenAILLMModels(model_id).value
        self.prompt_template = prompt_template
        self.client = OpenAI()

        if self.prompt_template is None:
            self.prompt_template = [
                {
                    "role": "system",
                    "content": "Select the expression from the list of expressions that is semantically closest to the query. You are mapping a non-standard concept to a standard concept and the list of expressions are already mapped to standard concepts. Your response should be a JSON object and include the expression_id of the selected expression. If no expression matches the query, return '{\"expression_id\": null}'.\n\nJSON object format: {json_format}",
                },
                {
                    "role": "user",
                    "content": "Query: '{query}'\n\nExpressions: {expressions}",
                },
            ]

    def select(self, results: RetrieverResults) -> SelectorResults:
        """Select the best candidate for each query using an OpenAI model.

        This method formats a prompt for each query, sends it to the configured
        OpenAI model for structured response parsing, validates the selected
        expression identifier, and returns either the matched candidate or an
        empty selection.

        Args:
            results: Retrieval results containing candidate expressions for
                each query.

        Returns:
            A ``SelectorResults`` object containing the selected candidate or
                an empty selection for each query.

        Raises:
            AssertionError: If the parsed response is not a ``SelectedResult``
                instance.
        """
        selector_results = SelectorResults(results=[], queries=results.queries)
        for query_id, query in enumerate(results.queries):
            prompt = format_prompt(
                self.prompt_template,
                {
                    "json_format": SelectedResult.model_json_schema(),
                    "query": query.capitalize(),  # avoid case sensitivity for some queries like 'cough' and 'Cough'
                    "expressions": [
                        r.to_prompt_object() for r in results.results[query_id]
                    ],
                },
            )

            if debug_mode == DebugMode.OPENAI_LLM_SELECTOR:
                logger.debug(prompt)

            response = self.client.responses.parse(
                model=self.model_id,
                input=prompt,
                text_format=SelectedResult,
            )

            if debug_mode == DebugMode.OPENAI_LLM_SELECTOR:
                logger.debug(response)
                logger.debug(response.output_parsed)

            selected_result: SelectedResult = response.output_parsed
            assert isinstance(selected_result, SelectedResult), (
                f"Expected SelectedResult object from OpenAI, but got {type(selected_result)}"
            )

            # case where no expression is selected
            if selected_result.expression_id is None:
                selector_results.results.append(EmptySelectionMetadata())
                continue

            # case where expression is selected but it is not in the results
            results_expression_ids = set(
                [r.expression_id for r in results.results[query_id]]
            )

            if selected_result.expression_id not in results_expression_ids:
                selector_results.results.append(EmptySelectionMetadata())
                continue

            # case where expression is selected and it is in the results
            for result_idx, r in enumerate(results.results[query_id]):
                if r.expression_id == selected_result.expression_id:
                    selector_results.results.append(
                        SelectedExpressionMetadata(
                            **r.model_dump(), result_list_index=result_idx
                        )
                    )
                    break

        return selector_results

__init__(model_id, prompt_template=None, *args, **kwargs)

Initialize the OpenAI-based selector.

This constructor resolves the configured model identifier, creates an OpenAI client, and sets the prompt template used to ask the model to choose the best candidate expression.

Parameters:

Name Type Description Default
model_id str

Name of the supported OpenAI model variant to use for selection.

 required
prompt_template List[Dict[str, str]]

Optional structured prompt template. If not provided, a default system-and-user prompt is used.

 None
*args Any

Additional positional arguments reserved for compatibility.

 ()
**kwargs Any

Additional keyword arguments reserved for compatibility.

 {}

Returns:

Type Description
None

None.

Raises:

Type Description
ValueError

If model_id is not a valid member of OpenAILLMModels.
Source code in aatm\selectors.py 
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
def __init__(
    self,
    model_id: str,
    prompt_template: List[Dict[str, str]] = None,
    *args: Any,
    **kwargs: Any,
) -> None:
    """Initialize the OpenAI-based selector.

    This constructor resolves the configured model identifier, creates an
    OpenAI client, and sets the prompt template used to ask the model to
    choose the best candidate expression.

    Args:
        model_id: Name of the supported OpenAI model variant to use for
            selection.
        prompt_template: Optional structured prompt template. If not
            provided, a default system-and-user prompt is used.
        *args: Additional positional arguments reserved for compatibility.
        **kwargs: Additional keyword arguments reserved for compatibility.

    Returns:
        None.

    Raises:
        ValueError: If ``model_id`` is not a valid member of
            ``OpenAILLMModels``.
    """
    self.model_id = OpenAILLMModels(model_id).value
    self.prompt_template = prompt_template
    self.client = OpenAI()

    if self.prompt_template is None:
        self.prompt_template = [
            {
                "role": "system",
                "content": "Select the expression from the list of expressions that is semantically closest to the query. You are mapping a non-standard concept to a standard concept and the list of expressions are already mapped to standard concepts. Your response should be a JSON object and include the expression_id of the selected expression. If no expression matches the query, return '{\"expression_id\": null}'.\n\nJSON object format: {json_format}",
            },
            {
                "role": "user",
                "content": "Query: '{query}'\n\nExpressions: {expressions}",
            },
        ]

select(results)

Select the best candidate for each query using an OpenAI model.

This method formats a prompt for each query, sends it to the configured OpenAI model for structured response parsing, validates the selected expression identifier, and returns either the matched candidate or an empty selection.

Parameters:

Name Type Description Default
results RetrieverResults

Retrieval results containing candidate expressions for each query.

 required

Returns:

Type Description
SelectorResults

A SelectorResults object containing the selected candidate or an empty selection for each query.

Raises:

Type Description
AssertionError

If the parsed response is not a SelectedResult instance.
Source code in aatm\selectors.py 
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
def select(self, results: RetrieverResults) -> SelectorResults:
    """Select the best candidate for each query using an OpenAI model.

    This method formats a prompt for each query, sends it to the configured
    OpenAI model for structured response parsing, validates the selected
    expression identifier, and returns either the matched candidate or an
    empty selection.

    Args:
        results: Retrieval results containing candidate expressions for
            each query.

    Returns:
        A ``SelectorResults`` object containing the selected candidate or
            an empty selection for each query.

    Raises:
        AssertionError: If the parsed response is not a ``SelectedResult``
            instance.
    """
    selector_results = SelectorResults(results=[], queries=results.queries)
    for query_id, query in enumerate(results.queries):
        prompt = format_prompt(
            self.prompt_template,
            {
                "json_format": SelectedResult.model_json_schema(),
                "query": query.capitalize(),  # avoid case sensitivity for some queries like 'cough' and 'Cough'
                "expressions": [
                    r.to_prompt_object() for r in results.results[query_id]
                ],
            },
        )

        if debug_mode == DebugMode.OPENAI_LLM_SELECTOR:
            logger.debug(prompt)

        response = self.client.responses.parse(
            model=self.model_id,
            input=prompt,
            text_format=SelectedResult,
        )

        if debug_mode == DebugMode.OPENAI_LLM_SELECTOR:
            logger.debug(response)
            logger.debug(response.output_parsed)

        selected_result: SelectedResult = response.output_parsed
        assert isinstance(selected_result, SelectedResult), (
            f"Expected SelectedResult object from OpenAI, but got {type(selected_result)}"
        )

        # case where no expression is selected
        if selected_result.expression_id is None:
            selector_results.results.append(EmptySelectionMetadata())
            continue

        # case where expression is selected but it is not in the results
        results_expression_ids = set(
            [r.expression_id for r in results.results[query_id]]
        )

        if selected_result.expression_id not in results_expression_ids:
            selector_results.results.append(EmptySelectionMetadata())
            continue

        # case where expression is selected and it is in the results
        for result_idx, r in enumerate(results.results[query_id]):
            if r.expression_id == selected_result.expression_id:
                selector_results.results.append(
                    SelectedExpressionMetadata(
                        **r.model_dump(), result_list_index=result_idx
                    )
                )
                break

    return selector_results

GeminiLLMModels

Bases: Enum

Enumerate the supported Gemini model identifiers for selection.

This enumeration defines the Gemini model names that can be used by the Gemini-based selector implementation.

Source code in aatm\selectors.py 
293
294
295
296
297
298
299
300
301
302
303
304
class GeminiLLMModels(Enum):
    """Enumerate the supported Gemini model identifiers for selection.

    This enumeration defines the Gemini model names that can be used by the
    Gemini-based selector implementation.
    """

    GEMINI_3_PRO_PREVIEW = "gemini-3-pro-preview"
    GEMINI_3_FLASH_PREVIEW = "gemini-3-flash-preview"
    GEMINI_2_5_FLASH = "gemini-2.5-flash"
    GEMINI_2_5_FLASH_LITE = "gemini-2.5-flash-lite"
    GEMINI_2_5_PRO = "gemini-2.5-pro"

GeminiLLMSelector

Bases: BaseSelector

Select results using a Gemini language model.

This selector formats each query and its retrieved candidate expressions into a prompt, sends the prompt to a Gemini model configured for JSON output, and validates the returned structured selection against the expected schema.

If the model returns no selection or selects an expression that is not present in the candidate list, the selector produces an empty selection.

Source code in aatm\selectors.py 
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
class GeminiLLMSelector(BaseSelector):
    """Select results using a Gemini language model.

    This selector formats each query and its retrieved candidate expressions
    into a prompt, sends the prompt to a Gemini model configured for JSON
    output, and validates the returned structured selection against the
    expected schema.

    If the model returns no selection or selects an expression that is not
    present in the candidate list, the selector produces an empty selection.
    """

    def __init__(
        self,
        model_id: str,
        prompt_template: List[Dict[str, str]] = None,
        *args: Any,
        **kwargs: Any,
    ) -> None:
        """Initialize the Gemini-based selector.

        This constructor resolves the configured Gemini model identifier,
        creates a Gemini client, and sets the prompt template used to ask the
        model to choose the best candidate expression.

        Args:
            model_id: Name of the supported Gemini model variant to use for
                selection.
            prompt_template: Optional structured prompt template. If not
                provided, a default system-and-user prompt is used.
            *args: Additional positional arguments reserved for compatibility.
            **kwargs: Additional keyword arguments reserved for compatibility.

        Returns:
            None.

        Raises:
            ValueError: If ``model_id`` is not a valid member of
                ``GeminiLLMModels``.
        """
        self.model_id = GeminiLLMModels(model_id).value
        self.client = genai.Client()
        self.prompt_template = prompt_template

        if self.prompt_template is None:
            self.prompt_template = [
                {
                    "role": "system",
                    "content": "Select the expression from the list of expressions that is semantically closest to the query. You are mapping a non-standard concept to a standard concept and the list of expressions are already mapped to standard concepts. Your response should be a JSON object and include the expression_id of the selected expression. If no expression matches the query, return '{\"expression_id\": null}'.\n\nJSON object format: {json_format}",
                },
                {
                    "role": "user",
                    "content": "Query: '{query}'\n\nExpressions: {expressions}",
                },
            ]

    def select(self, results: RetrieverResults) -> SelectorResults:
        """Select the best candidate for each query using a Gemini model.

        This method formats a prompt for each query, converts it into Gemini
        content objects, requests a JSON response constrained by the selection
        schema, validates the returned expression identifier, and returns
        either the matched candidate or an empty selection.

        Args:
            results: Retrieval results containing candidate expressions for
                each query.

        Returns:
            A ``SelectorResults`` object containing the selected candidate or
                an empty selection for each query.

        Raises:
            AssertionError: If the parsed response is not a ``SelectedResult``
                instance.
        """
        selector_results = SelectorResults(results=[], queries=results.queries)
        for query_id, query in enumerate(results.queries):
            prompt = format_prompt(
                self.prompt_template,
                {
                    "json_format": SelectedResult.model_json_schema(),
                    "query": query.capitalize(),  # avoid case sensitivity for some queries like 'cough' and 'Cough'
                    "expressions": [
                        r.to_prompt_object() for r in results.results[query_id]
                    ],
                },
            )

            gemini_formatted_prompt = []
            for msg in prompt:
                gemini_formatted_prompt.append(types.UserContent(msg["content"]))

            if debug_mode == DebugMode.OPENAI_LLM_SELECTOR:
                logger.debug(prompt)

            response = self.client.models.generate_content(
                model=self.model_id,
                contents=gemini_formatted_prompt,
                config=types.GenerateContentConfig(
                    response_mime_type="application/json",
                    response_schema=SelectedResult,
                ),
            )

            if debug_mode == DebugMode.GEMINI_LLM_SELECTOR:
                logger.debug(response)
                logger.debug(response.text)
                logger.debug(SelectedResult.model_validate_json(response.text))

            selected_result: SelectedResult = SelectedResult.model_validate_json(
                response.text
            )
            assert isinstance(selected_result, SelectedResult), (
                f"Expected SelectedResult object from OpenAI, but got {type(selected_result)}"
            )

            # case where no expression is selected
            if selected_result.expression_id is None:
                selector_results.results.append(EmptySelectionMetadata())
                continue

            # case where expression is selected but it is not in the results
            results_expression_ids = set(
                [r.expression_id for r in results.results[query_id]]
            )

            if selected_result.expression_id not in results_expression_ids:
                selector_results.results.append(EmptySelectionMetadata())
                continue

            # case where expression is selected and it is in the results
            for result_idx, r in enumerate(results.results[query_id]):
                if r.expression_id == selected_result.expression_id:
                    selector_results.results.append(
                        SelectedExpressionMetadata(
                            **r.model_dump(), result_list_index=result_idx
                        )
                    )
                    break

        return selector_results

__init__(model_id, prompt_template=None, *args, **kwargs)

Initialize the Gemini-based selector.

This constructor resolves the configured Gemini model identifier, creates a Gemini client, and sets the prompt template used to ask the model to choose the best candidate expression.

Parameters:

Name Type Description Default
model_id str

Name of the supported Gemini model variant to use for selection.

 required
prompt_template List[Dict[str, str]]

Optional structured prompt template. If not provided, a default system-and-user prompt is used.

 None
*args Any

Additional positional arguments reserved for compatibility.

 ()
**kwargs Any

Additional keyword arguments reserved for compatibility.

 {}

Returns:

Type Description
None

None.

Raises:

Type Description
ValueError

If model_id is not a valid member of GeminiLLMModels.
Source code in aatm\selectors.py 
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
def __init__(
    self,
    model_id: str,
    prompt_template: List[Dict[str, str]] = None,
    *args: Any,
    **kwargs: Any,
) -> None:
    """Initialize the Gemini-based selector.

    This constructor resolves the configured Gemini model identifier,
    creates a Gemini client, and sets the prompt template used to ask the
    model to choose the best candidate expression.

    Args:
        model_id: Name of the supported Gemini model variant to use for
            selection.
        prompt_template: Optional structured prompt template. If not
            provided, a default system-and-user prompt is used.
        *args: Additional positional arguments reserved for compatibility.
        **kwargs: Additional keyword arguments reserved for compatibility.

    Returns:
        None.

    Raises:
        ValueError: If ``model_id`` is not a valid member of
            ``GeminiLLMModels``.
    """
    self.model_id = GeminiLLMModels(model_id).value
    self.client = genai.Client()
    self.prompt_template = prompt_template

    if self.prompt_template is None:
        self.prompt_template = [
            {
                "role": "system",
                "content": "Select the expression from the list of expressions that is semantically closest to the query. You are mapping a non-standard concept to a standard concept and the list of expressions are already mapped to standard concepts. Your response should be a JSON object and include the expression_id of the selected expression. If no expression matches the query, return '{\"expression_id\": null}'.\n\nJSON object format: {json_format}",
            },
            {
                "role": "user",
                "content": "Query: '{query}'\n\nExpressions: {expressions}",
            },
        ]

select(results)

Select the best candidate for each query using a Gemini model.

This method formats a prompt for each query, converts it into Gemini content objects, requests a JSON response constrained by the selection schema, validates the returned expression identifier, and returns either the matched candidate or an empty selection.

Parameters:

Name Type Description Default
results RetrieverResults

Retrieval results containing candidate expressions for each query.

 required

Returns:

Type Description
SelectorResults

A SelectorResults object containing the selected candidate or an empty selection for each query.

Raises:

Type Description
AssertionError

If the parsed response is not a SelectedResult instance.
Source code in aatm\selectors.py 
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
def select(self, results: RetrieverResults) -> SelectorResults:
    """Select the best candidate for each query using a Gemini model.

    This method formats a prompt for each query, converts it into Gemini
    content objects, requests a JSON response constrained by the selection
    schema, validates the returned expression identifier, and returns
    either the matched candidate or an empty selection.

    Args:
        results: Retrieval results containing candidate expressions for
            each query.

    Returns:
        A ``SelectorResults`` object containing the selected candidate or
            an empty selection for each query.

    Raises:
        AssertionError: If the parsed response is not a ``SelectedResult``
            instance.
    """
    selector_results = SelectorResults(results=[], queries=results.queries)
    for query_id, query in enumerate(results.queries):
        prompt = format_prompt(
            self.prompt_template,
            {
                "json_format": SelectedResult.model_json_schema(),
                "query": query.capitalize(),  # avoid case sensitivity for some queries like 'cough' and 'Cough'
                "expressions": [
                    r.to_prompt_object() for r in results.results[query_id]
                ],
            },
        )

        gemini_formatted_prompt = []
        for msg in prompt:
            gemini_formatted_prompt.append(types.UserContent(msg["content"]))

        if debug_mode == DebugMode.OPENAI_LLM_SELECTOR:
            logger.debug(prompt)

        response = self.client.models.generate_content(
            model=self.model_id,
            contents=gemini_formatted_prompt,
            config=types.GenerateContentConfig(
                response_mime_type="application/json",
                response_schema=SelectedResult,
            ),
        )

        if debug_mode == DebugMode.GEMINI_LLM_SELECTOR:
            logger.debug(response)
            logger.debug(response.text)
            logger.debug(SelectedResult.model_validate_json(response.text))

        selected_result: SelectedResult = SelectedResult.model_validate_json(
            response.text
        )
        assert isinstance(selected_result, SelectedResult), (
            f"Expected SelectedResult object from OpenAI, but got {type(selected_result)}"
        )

        # case where no expression is selected
        if selected_result.expression_id is None:
            selector_results.results.append(EmptySelectionMetadata())
            continue

        # case where expression is selected but it is not in the results
        results_expression_ids = set(
            [r.expression_id for r in results.results[query_id]]
        )

        if selected_result.expression_id not in results_expression_ids:
            selector_results.results.append(EmptySelectionMetadata())
            continue

        # case where expression is selected and it is in the results
        for result_idx, r in enumerate(results.results[query_id]):
            if r.expression_id == selected_result.expression_id:
                selector_results.results.append(
                    SelectedExpressionMetadata(
                        **r.model_dump(), result_list_index=result_idx
                    )
                )
                break

    return selector_results