From 23122277666d4de13c32787f747e6aaa79cc7e41 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Tue, 2 Jan 2024 22:36:26 -0300 Subject: [PATCH] Add ChatAdapter class to utils/chat.py --- src/backend/langflow/utils/chat.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) create mode 100644 src/backend/langflow/utils/chat.py diff --git a/src/backend/langflow/utils/chat.py b/src/backend/langflow/utils/chat.py new file mode 100644 index 000000000..84405e38e --- /dev/null +++ b/src/backend/langflow/utils/chat.py @@ -0,0 +1,23 @@ +from typing import Any, Callable, Optional, Union + +from langchain_core.prompts import PromptTemplate as LCPromptTemplate +from langflow.utils.prompt import GenericPromptTemplate +from llama_index.prompts import PromptTemplate as LIPromptTemplate + +PromptTemplate = Union[LCPromptTemplate, LIPromptTemplate] + + +class ChatAdapter: + def __init__(self, func: Callable, inputs: list[str], output_key: str, prompt: Optional[PromptTemplate] = None): + self.func = func + self.input_keys = inputs + self.output_keys = output_key + self.prompt = prompt + + @classmethod + def from_prompt_template(cls, prompt_template: PromptTemplate, func: Callable): + prompt = GenericPromptTemplate(prompt_template) + return cls(func, prompt.input_keys, prompt_template) + + def __call__(self, inputs: dict, callbacks: Optional[Any] = None) -> dict: + return self.func(inputs, callbacks)