kotaemon/knowledgehub/llms/completions/base.py
Nguyen Trung Duc (john) c3c25db48c [AUR-385, AUR-388] Declare BaseComponent and decide LLM call interface (#2)
- Use cases related to LLM call: https://cinnamon-ai.atlassian.net/browse/AUR-388?focusedCommentId=34873
- Sample usages: `test_llms_chat_models.py` and `test_llms_completion_models.py`:

```python
from kotaemon.llms.chats.openai import AzureChatOpenAI

model = AzureChatOpenAI(
    openai_api_base="https://test.openai.azure.com/",
    openai_api_key="some-key",
    openai_api_version="2023-03-15-preview",
    deployment_name="gpt35turbo",
    temperature=0,
    request_timeout=60,
)
output = model("hello world")
```

For the LLM-call component, I decide to wrap around Langchain's LLM models and Langchain's Chat models. And set the interface as follow:

- Completion LLM component:
```python
class CompletionLLM:

    def run_raw(self, text: str) -> LLMInterface:
        # Run text completion: str in -> LLMInterface out

    def run_batch_raw(self, text: list[str]) -> list[LLMInterface]:
        # Run text completion in batch: list[str] in -> list[LLMInterface] out

# run_document and run_batch_document just reuse run_raw and run_batch_raw, due to unclear use case
```

- Chat LLM component:
```python
class ChatLLM:
    def run_raw(self, text: str) -> LLMInterface:
        # Run chat completion (no chat history): str in -> LLMInterface out

    def run_batch_raw(self, text: list[str]) -> list[LLMInterface]:
        # Run chat completion in batch mode (no chat history): list[str] in -> list[LLMInterface] out

    def run_document(self, text: list[BaseMessage]) -> LLMInterface:
        # Run chat completion (with chat history): list[langchain's BaseMessage] in -> LLMInterface out

    def run_batch_document(self, text: list[list[BaseMessage]]) -> list[LLMInterface]:
        # Run chat completion in batch mode (with chat history): list[list[langchain's BaseMessage]] in -> list[LLMInterface] out
```

- The LLMInterface is as follow:

```python
@dataclass
class LLMInterface:
    text: list[str]
    completion_tokens: int = -1
    total_tokens: int = -1
    prompt_tokens: int = -1
    logits: list[list[float]] = field(default_factory=list)
```
2023-08-29 15:47:12 +07:00

71 lines
2.0 KiB
Python

from typing import Type
from theflow.base import Param
from langchain.schema.language_model import BaseLanguageModel
from ...components import BaseComponent
from ..base import LLMInterface
class LLM(BaseComponent):
pass
class LangchainLLM(LLM):
_lc_class: Type[BaseLanguageModel]
def __init__(self, **params):
if self._lc_class is None:
raise AttributeError(
"Should set _lc_class attribute to the LLM class from Langchain "
"if using LLM from Langchain"
)
self._kwargs: dict = {}
for param in list(params.keys()):
if param in self._lc_class.__fields__:
self._kwargs[param] = params.pop(param)
super().__init__(**params)
@Param.decorate()
def agent(self):
return self._lc_class(**self._kwargs)
def run_raw(self, text: str) -> LLMInterface:
pred = self.agent.generate([text])
return LLMInterface(
text=[each.text for each in pred.generations[0]],
completion_tokens=pred.llm_output["token_usage"]["completion_tokens"],
total_tokens=pred.llm_output["token_usage"]["total_tokens"],
prompt_tokens=pred.llm_output["token_usage"]["prompt_tokens"],
logits=[],
)
def run_batch_raw(self, text: list[str]) -> list[LLMInterface]:
outputs = []
for each_text in text:
outputs.append(self.run_raw(each_text))
return outputs
def run_document(self, text: str) -> LLMInterface:
return self.run_raw(text)
def run_batch_document(self, text: list[str]) -> list[LLMInterface]:
return self.run_batch_raw(text)
def is_document(self, text) -> bool:
return False
def is_batch(self, text) -> bool:
return False if isinstance(text, str) else True
def __setattr__(self, name, value):
if name in self._lc_class.__fields__:
setattr(self.agent, name, value)
else:
super().__setattr__(name, value)
class LLMChat(BaseComponent):
pass