LLM/DocsGPT
1
0
Fork 0
mirror of https://github.com/arc53/DocsGPT.git synced 2025-11-29 08:33:20 +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
)

7
application/llm/docsgpt_provider.py
View File

@@ -1,7 +1,6 @@
from application.llm.base import BaseLLM
import json
import requests
from application.usage import gen_token_usage, stream_token_usage
class DocsGPTAPILLM(BaseLLM):
@@ -11,8 +10,7 @@ class DocsGPTAPILLM(BaseLLM):
self.api_key = api_key
self.endpoint = "https://llm.docsgpt.co.uk"
@gen_token_usage
def gen(self, model, messages, stream=False, **kwargs):
def _raw_gen(self, baseself, model, messages, stream=False, *args, **kwargs):
context = messages[0]["content"]
user_question = messages[-1]["content"]
prompt = f"### Instruction \n {user_question} \n ### Context \n {context} \n ### Answer \n"
@@ -24,8 +22,7 @@ class DocsGPTAPILLM(BaseLLM):
return response_clean
@stream_token_usage
def gen_stream(self, model, messages, stream=True, **kwargs):
def _raw_gen_stream(self, baseself, model, messages, stream=True, *args, **kwargs):
context = messages[0]["content"]
user_question = messages[-1]["content"]
prompt = f"### Instruction \n {user_question} \n ### Context \n {context} \n ### Answer \n"