Skip to content

FlowStep

This module the FlowStep class, which can call a language model, record run times, and optionally invoke callbacks on the results.

FlowStep

Bases: BaseFlowStep

Represents a specific step in a Flow.

A FlowStep calls a language model using a prompt template, records the run time, and optionally invokes callback functions on the results.

Parameters:

Name Type Description Default
name str

The name of the flow step.

 required
llm BaseLLM

The language model to be used in the flow step.

 required
prompt_template PromptTemplate

Template for the prompt to be used with the language model.

 required
callbacks list[BaseCallback]

Callbacks to be invoked within the flowstep

 None

Attributes:

Name Type Description
llm BaseLLM

The language model to be used in the flow step.
prompt_template PromptTemplate

Template for the prompt to be used with the language model.
required_keys set[str]

The keys required for the flow step to run.
Source code in llmflows/flows/flowstep.py 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
class FlowStep(BaseFlowStep):
    """
    Represents a specific step in a Flow.

    A FlowStep calls a language model using a prompt template, records the run
    time, and optionally invokes callback functions on the results.

    Args:
        name (str): The name of the flow step.
        llm (BaseLLM): The language model to be used in the flow step.
        prompt_template (PromptTemplate): Template for the prompt to be used with the 
            language model.
        callbacks (list[BaseCallback]): Callbacks to be invoked within the flowstep

    Attributes:
        llm (BaseLLM): The language model to be used in the flow step.
        prompt_template (PromptTemplate): Template for the prompt to be used with the 
            language model.
        required_keys (set[str]): The keys required for the flow step to run.
    """

    def __init__(
        self,
        name: str,
        llm: BaseLLM,
        prompt_template: PromptTemplate,
        output_key: str,
        callbacks:  Union[list[BaseCallback], None] = None,
    ):
        super().__init__(name, output_key, callbacks)
        self.llm = llm
        self.prompt_template = prompt_template
        self.required_keys = prompt_template.variables

    def generate(self, inputs: dict[str, Any]) -> tuple[Any, Any, Any]:
        prompt = self.prompt_template.get_prompt(**inputs)
        text_result, call_data, model_config = self.llm.generate(prompt)
        call_data["prompt_template"] = self.prompt_template.prompt
        call_data["prompt"] = prompt

        return text_result, call_data, model_config