mirror of
https://github.com/arc53/DocsGPT.git
synced 2025-11-29 16:43:16 +00:00
67 lines
2.0 KiB
Python
67 lines
2.0 KiB
Python
from abc import ABC, abstractmethod
|
|
|
|
from application.cache import gen_cache, stream_cache
|
|
from application.usage import gen_token_usage, stream_token_usage
|
|
|
|
|
|
class BaseLLM(ABC):
|
|
def __init__(self, decoded_token=None):
|
|
self.decoded_token = decoded_token
|
|
self.token_usage = {"prompt_tokens": 0, "generated_tokens": 0}
|
|
|
|
def _apply_decorator(self, method, decorators, *args, **kwargs):
|
|
for decorator in decorators:
|
|
method = decorator(method)
|
|
return method(self, *args, **kwargs)
|
|
|
|
@abstractmethod
|
|
def _raw_gen(self, model, messages, stream, tools, *args, **kwargs):
|
|
pass
|
|
|
|
def gen(self, model, messages, stream=False, tools=None, *args, **kwargs):
|
|
decorators = [gen_token_usage, gen_cache]
|
|
return self._apply_decorator(
|
|
self._raw_gen,
|
|
decorators=decorators,
|
|
model=model,
|
|
messages=messages,
|
|
stream=stream,
|
|
tools=tools,
|
|
*args,
|
|
**kwargs
|
|
)
|
|
|
|
@abstractmethod
|
|
def _raw_gen_stream(self, model, messages, stream, *args, **kwargs):
|
|
pass
|
|
|
|
def gen_stream(self, model, messages, stream=True, tools=None, *args, **kwargs):
|
|
decorators = [stream_cache, stream_token_usage]
|
|
return self._apply_decorator(
|
|
self._raw_gen_stream,
|
|
decorators=decorators,
|
|
model=model,
|
|
messages=messages,
|
|
stream=stream,
|
|
tools=tools,
|
|
*args,
|
|
**kwargs
|
|
)
|
|
|
|
def supports_tools(self):
|
|
return hasattr(self, "_supports_tools") and callable(
|
|
getattr(self, "_supports_tools")
|
|
)
|
|
|
|
def _supports_tools(self):
|
|
raise NotImplementedError("Subclass must implement _supports_tools method")
|
|
|
|
def get_supported_attachment_types(self):
|
|
"""
|
|
Return a list of MIME types supported by this LLM for file uploads.
|
|
|
|
Returns:
|
|
list: List of supported MIME types
|
|
"""
|
|
return [] # Default: no attachments supported
|