File size: 1,296 Bytes
a553e02 c0cd1dc a553e02 c0cd1dc a553e02 c0cd1dc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 |
from llama_index.core.bridge.langchain import BaseLanguageModel, BaseChatModel
from llama_index.llms.langchain import LangChainLLM
from llama_index.core.bridge.langchain import OpenAI, ChatOpenAI
from llama_index.core.llms import LLMMetadata
from kron.llm_predictor.openai_utils import kron_openai_modelname_to_contextsize
def is_chat_model(llm: BaseLanguageModel) -> bool:
return isinstance(llm, BaseChatModel)
class KronLangChainLLM(LangChainLLM):
"""Adapter for a LangChain LLM."""
def __init__(self, llm: BaseLanguageModel) -> None:
super().__init__(llm)
@property
def metadata(self) -> LLMMetadata:
is_chat_model_ = is_chat_model(self.llm)
if isinstance(self.llm, OpenAI):
return LLMMetadata(
context_window=kron_openai_modelname_to_contextsize(self.llm.model_name),
num_output=self.llm.max_tokens,
is_chat_model=is_chat_model_ ,
)
elif isinstance(self.llm, ChatOpenAI):
return LLMMetadata(
context_window=kron_openai_modelname_to_contextsize(self.llm.model_name),
num_output=self.llm.max_tokens or -1,
is_chat_model=is_chat_model_ ,
)
else:
return super().metadata()
|