liteflow/examples/batch_email_analysis.py at main · BerriAI/liteflow · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
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
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
"""Compile per-email parallel analyses and a final batch report.

Model names are illustrative configured gateway aliases. Registration and
execution require the v2 executor. Run this file to print the compiled
definition without calling a model or reading gateway credentials.
"""

import json
from pydantic import BaseModel
from typing_extensions import TypedDict

from liteflow import ChatCompletion, ChatMessage, ForEach, Input, Parallel, Python, Sequence, Workflow, compile_workflow


class Emails(BaseModel):
    items: list[str]


class Finding(BaseModel):
    summary: str


class Analyses(TypedDict):
    intent: Finding
    sentiment: Finding


class EmailAnalysis(TypedDict):
    email: str
    intent: str
    sentiment: str


def prepare_analysis(email: str, instruction: str) -> list[ChatMessage]:
    return [{"role": "system", "content": instruction}, {"role": "user", "content": email}]


def combine_analysis(analyses: Analyses, email: str) -> EmailAnalysis:
    return {"email": email, "intent": analyses["intent"].summary, "sentiment": analyses["sentiment"].summary}


def prepare_report(analyses: list[EmailAnalysis]) -> list[ChatMessage]:
    lines = [item["email"] + "\nIntent: " + item["intent"] + "\nSentiment: " + item["sentiment"] for item in analyses]
    return [{"role": "user", "content": "Summarize these email analyses:\n\n" + "\n\n".join(lines)}]


workflow = Workflow(
    name="batch-email-analysis",
    nodes=[
        ChatCompletion(model="support-model", messages=Input.messages, response_model=Emails),
        ForEach(
            items=Input.previous.field("items"),
            max_items=100,
            max_concurrency=4,
            nodes=[
                Parallel(
                    branches={
                        "intent": [
                            Python(
                                fn=prepare_analysis,
                                inputs={"email": Input.item, "instruction": "Identify the request in this email."},
                            ),
                            ChatCompletion(model="haiku", messages=Input.previous, response_model=Finding),
                        ],
                        "sentiment": [
                            Python(
                                fn=prepare_analysis,
                                inputs={"email": Input.item, "instruction": "Assess the customer's sentiment."},
                            ),
                            ChatCompletion(model="haiku", messages=Input.previous, response_model=Finding),
                        ],
                    },
                    max_concurrency=2,
                ),
                Python(fn=combine_analysis, inputs={"analyses": Input.previous, "email": Input.item}),
            ],
        ),
        Sequence(
            nodes=[
                Python(fn=prepare_report),
                ChatCompletion(model="support-model", messages=Input.previous, timeout_seconds=35),
            ]
        ),
    ],
)


if __name__ == "__main__":
    print(json.dumps(compile_workflow(workflow), indent=2))