From c23d79767709b18a78d7982592d9bfc4be44078d Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Thu, 2 Nov 2023 22:46:15 -0300
Subject: [PATCH] Add LLMChain component for running queries against LLMs

---
 .../langflow/components/chains/LLMChain.py    | 24 +++++++++++++++++++
 1 file changed, 24 insertions(+)
 create mode 100644 src/backend/langflow/components/chains/LLMChain.py

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)