LLM/DocsGPT
1
0
Fork 0
mirror of https://github.com/arc53/DocsGPT.git synced 2025-11-29 16:43:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

update: apply decorator to abstract classes

Browse Source
This commit is contained in:
Siddhant Rai
2024-04-15 18:57:28 +05:30
parent 262d160314
commit 590aa8b43f
3 changed files with 28 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
application/llm/base.py
View File

@@ -1,14 +1,28 @@
from abc import ABC, abstractmethod
from application.usage import gen_token_usage, stream_token_usage
class BaseLLM(ABC):
def __init__(self):
self.token_usage = {"prompt_tokens": 0, "generated_tokens": 0}
@abstractmethod
def gen(self, *args, **kwargs):
pass
def _apply_decorator(self, method, decorator, *args, **kwargs):
return decorator(method, *args, **kwargs)
@abstractmethod
def gen_stream(self, *args, **kwargs):
def _raw_gen(self, model, messages, stream, *args, **kwargs):
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
def _raw_gen_stream(self, model, messages, stream, *args, **kwargs):
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
)