mirror of
https://github.com/csunny/DB-GPT.git
synced 2026-01-25 14:54:26 +00:00
39 lines
1010 B
Python
39 lines
1010 B
Python
from abc import ABC, abstractmethod
|
|
from pydantic import BaseModel
|
|
from typing import List
|
|
from pilot.common.schema import ExampleType
|
|
|
|
|
|
class ExampleSelector(BaseModel, ABC):
|
|
examples_record: List[dict]
|
|
use_example: bool = False
|
|
type: str = ExampleType.ONE_SHOT.value
|
|
|
|
def examples(self, count: int = 2):
|
|
if ExampleType.ONE_SHOT.value == self.type:
|
|
return self.__one_show_context()
|
|
else:
|
|
return self.__few_shot_context(count)
|
|
|
|
def __few_shot_context(self, count: int = 2) -> List[dict]:
|
|
"""
|
|
Use 2 or more examples, default 2
|
|
Returns: example text
|
|
"""
|
|
if self.use_example:
|
|
need_use = self.examples_record[:count]
|
|
return need_use
|
|
return None
|
|
|
|
def __one_show_context(self) -> dict:
|
|
"""
|
|
Use one examples
|
|
Returns:
|
|
|
|
"""
|
|
if self.use_example:
|
|
need_use = self.examples_record[:1]
|
|
return need_use
|
|
|
|
return None
|