feat(prompt_runner.py): add PromptRunner component to run a Chain with a given PromptTemplate

This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-08-04 00:11:00 -03:00
commit d99bcbd80b

View file

@ -0,0 +1,31 @@
from langflow import CustomComponent
from langchain.llms.base import BaseLLM
from langchain import PromptTemplate
from langchain.schema import Document
class PromptRunner(CustomComponent):
display_name: str = "Prompt Runner"
description: str = "Run a Chain with the given PromptTemplate"
beta = True
field_config = {
"llm": {"display_name": "LLM"},
"prompt": {
"display_name": "Prompt Template",
"info": "Make sure the prompt has all variables filled.",
},
"code": {"show": False},
"inputs": {"field_type": "code"},
}
def build(
self,
llm: BaseLLM,
prompt: PromptTemplate,
) -> Document:
chain = prompt | llm
result = chain.invoke()
result = result[chain.output_key]
self.repr_value = result
return Document(page_content=str(result))