llama-index

1
from typing import Optional
2

3
from llama_index.legacy.core.base_selector import BaseSelector
4
from llama_index.legacy.selectors.llm_selectors import (
5
    LLMMultiSelector,
6
    LLMSingleSelector,
7
)
8
from llama_index.legacy.selectors.pydantic_selectors import (
9
    PydanticMultiSelector,
10
    PydanticSingleSelector,
11
)
12
from llama_index.legacy.service_context import ServiceContext
13

14

15
def get_selector_from_context(
16
    service_context: ServiceContext, is_multi: bool = False
17
) -> BaseSelector:
18
    """Get a selector from a service context. Prefers Pydantic selectors if possible."""
19
    selector: Optional[BaseSelector] = None
20

21
    if is_multi:
22
        try:
23
            llm = service_context.llm
24
            selector = PydanticMultiSelector.from_defaults(llm=llm)  # type: ignore
25
        except ValueError:
26
            selector = LLMMultiSelector.from_defaults(service_context=service_context)
27
    else:
28
        try:
29
            llm = service_context.llm
30
            selector = PydanticSingleSelector.from_defaults(llm=llm)  # type: ignore
31
        except ValueError:
32
            selector = LLMSingleSelector.from_defaults(service_context=service_context)
33

34
    assert selector is not None
35

36
    return selector
37