mirror of https://github.com/arc53/DocsGPT
update: apply decorator to abstract classes
parent
262d160314
commit
590aa8b43f
@ -1,14 +1,28 @@
|
|||||||
from abc import ABC, abstractmethod
|
from abc import ABC, abstractmethod
|
||||||
|
from application.usage import gen_token_usage, stream_token_usage
|
||||||
|
|
||||||
|
|
||||||
class BaseLLM(ABC):
|
class BaseLLM(ABC):
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
self.token_usage = {"prompt_tokens": 0, "generated_tokens": 0}
|
self.token_usage = {"prompt_tokens": 0, "generated_tokens": 0}
|
||||||
|
|
||||||
|
def _apply_decorator(self, method, decorator, *args, **kwargs):
|
||||||
|
return decorator(method, *args, **kwargs)
|
||||||
|
|
||||||
@abstractmethod
|
@abstractmethod
|
||||||
def gen(self, *args, **kwargs):
|
def _raw_gen(self, model, messages, stream, *args, **kwargs):
|
||||||
pass
|
pass
|
||||||
|
|
||||||
|
def gen(self, model, messages, stream=False, *args, **kwargs):
|
||||||
|
return self._apply_decorator(self._raw_gen, gen_token_usage)(
|
||||||
|
self, model=model, messages=messages, stream=stream, *args, **kwargs
|
||||||
|
)
|
||||||
|
|
||||||
@abstractmethod
|
@abstractmethod
|
||||||
def gen_stream(self, *args, **kwargs):
|
def _raw_gen_stream(self, model, messages, stream, *args, **kwargs):
|
||||||
pass
|
pass
|
||||||
|
|
||||||
|
def gen_stream(self, model, messages, stream=True, *args, **kwargs):
|
||||||
|
return self._apply_decorator(self._raw_gen_stream, stream_token_usage)(
|
||||||
|
self, model=model, messages=messages, stream=stream, *args, **kwargs
|
||||||
|
)
|
||||||
|
Loading…
Reference in New Issue