diff --git a/src/backend/langflow/components/chains/LLMChain.py b/src/backend/langflow/components/chains/LLMChain.py new file mode 100644 index 000000000..5056c7f43 --- /dev/null +++ b/src/backend/langflow/components/chains/LLMChain.py @@ -0,0 +1,24 @@ +from langflow import CustomComponent +from langchain.chains import LLMChain +from typing import Optional, Union, Callable +from langflow.field_typing import PromptTemplate, BaseLanguageModel, BaseMemory, Chain + + +class LLMChainComponent(CustomComponent): + display_name = "LLMChain" + description = "Chain to run queries against LLMs" + + def build_config(self): + return { + "prompt": {"display_name": "Prompt"}, + "llm": {"display_name": "LLM"}, + "memory": {"display_name": "Memory"}, + } + + def build( + self, + prompt: PromptTemplate, + llm: BaseLanguageModel, + memory: Optional[BaseMemory], + ) -> Union[Chain, Callable]: + return LLMChain(prompt=prompt, llm=llm, memory=memory)