-
Notifications
You must be signed in to change notification settings - Fork 4.6k
/
base_selector.py
115 lines (88 loc) · 3.34 KB
/
base_selector.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
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
85
86
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
from abc import abstractmethod
from typing import Any, List, Sequence, Union
from llama_index.core.base.query_pipeline.query import (
ChainableMixin,
QueryComponent,
)
from llama_index.core.bridge.pydantic import BaseModel
from llama_index.core.instrumentation import DispatcherSpanMixin
from llama_index.core.prompts.mixin import PromptMixin, PromptMixinType
from llama_index.core.schema import QueryBundle, QueryType
from llama_index.core.tools.types import ToolMetadata
MetadataType = Union[str, ToolMetadata]
class SingleSelection(BaseModel):
"""A single selection of a choice."""
index: int
reason: str
class MultiSelection(BaseModel):
"""A multi-selection of choices."""
selections: List[SingleSelection]
@property
def ind(self) -> int:
if len(self.selections) != 1:
raise ValueError(
f"There are {len(self.selections)} selections, " "please use .inds."
)
return self.selections[0].index
@property
def reason(self) -> str:
if len(self.reasons) != 1:
raise ValueError(
f"There are {len(self.reasons)} selections, " "please use .reasons."
)
return self.selections[0].reason
@property
def inds(self) -> List[int]:
return [x.index for x in self.selections]
@property
def reasons(self) -> List[str]:
return [x.reason for x in self.selections]
# separate name for clarity and to not confuse function calling model
SelectorResult = MultiSelection
def _wrap_choice(choice: MetadataType) -> ToolMetadata:
if isinstance(choice, ToolMetadata):
return choice
elif isinstance(choice, str):
return ToolMetadata(description=choice)
else:
raise ValueError(f"Unexpected type: {type(choice)}")
def _wrap_query(query: QueryType) -> QueryBundle:
if isinstance(query, QueryBundle):
return query
elif isinstance(query, str):
return QueryBundle(query_str=query)
else:
raise ValueError(f"Unexpected type: {type(query)}")
class BaseSelector(PromptMixin, ChainableMixin, DispatcherSpanMixin):
"""Base selector."""
def _get_prompt_modules(self) -> PromptMixinType:
"""Get prompt sub-modules."""
return {}
def select(
self, choices: Sequence[MetadataType], query: QueryType
) -> SelectorResult:
metadatas = [_wrap_choice(choice) for choice in choices]
query_bundle = _wrap_query(query)
return self._select(choices=metadatas, query=query_bundle)
async def aselect(
self, choices: Sequence[MetadataType], query: QueryType
) -> SelectorResult:
metadatas = [_wrap_choice(choice) for choice in choices]
query_bundle = _wrap_query(query)
return await self._aselect(choices=metadatas, query=query_bundle)
@abstractmethod
def _select(
self, choices: Sequence[ToolMetadata], query: QueryBundle
) -> SelectorResult:
pass
@abstractmethod
async def _aselect(
self, choices: Sequence[ToolMetadata], query: QueryBundle
) -> SelectorResult:
pass
def _as_query_component(self, **kwargs: Any) -> QueryComponent:
"""As query component."""
from llama_index.core.query_pipeline.components.router import (
SelectorComponent,
)
return SelectorComponent(selector=self)