File size: 3,791 Bytes
ed4d993
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
import json
from datetime import datetime
from enum import Enum
from operator import itemgetter
from typing import Any, Dict, Sequence

from langchain.chains.openai_functions import convert_to_openai_function
from langchain_community.chat_models import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.pydantic_v1 import BaseModel, Field, ValidationError, conint
from langchain_core.runnables import (
    Runnable,
    RunnableBranch,
    RunnableLambda,
    RunnablePassthrough,
)


class TaskType(str, Enum):
    call = "Call"
    message = "Message"
    todo = "Todo"
    in_person_meeting = "In-Person Meeting"
    email = "Email"
    mail = "Mail"
    text = "Text"
    open_house = "Open House"


class Task(BaseModel):
    title: str = Field(..., description="The title of the tasks, reminders and alerts")
    due_date: datetime = Field(
        ..., description="Due date. Must be a valid ISO date string with timezone"
    )
    task_type: TaskType = Field(None, description="The type of task")


class Tasks(BaseModel):
    """JSON definition for creating tasks, reminders and alerts"""

    tasks: Sequence[Task]


template = """Respond to the following user query to the best of your ability:

{query}"""

generate_prompt = ChatPromptTemplate.from_template(template)

function_args = {"functions": [convert_to_openai_function(Tasks)]}

task_function_call_model = ChatOpenAI(model="gpt-3.5-turbo").bind(**function_args)

output_parser = RunnableLambda(
    lambda x: json.loads(
        x.additional_kwargs.get("function_call", {}).get("arguments", '""')
    )
)


revise_template = """
Based on the provided context, fix the incorrect result of the original prompt
and the provided errors. Only respond with an answer that satisfies the
constraints laid out in the original prompt and fixes the Pydantic errors.

Hint: Datetime fields must be valid ISO date strings.

<context>
<original_prompt>
{original_prompt}
</original_prompt>
<incorrect_result>
{completion}
</incorrect_result>
<errors>
{error}
</errors>
</context>"""

revise_prompt = ChatPromptTemplate.from_template(revise_template)

revise_chain = revise_prompt | task_function_call_model | output_parser


def output_validator(output):
    try:
        Tasks.validate(output["completion"])
    except ValidationError as e:
        return str(e)

    return None


class IntermediateType(BaseModel):
    error: str
    completion: Dict
    original_prompt: str
    max_revisions: int


validation_step = RunnablePassthrough().assign(error=RunnableLambda(output_validator))


def revise_loop(input: IntermediateType) -> IntermediateType:
    revise_step = RunnablePassthrough().assign(completion=revise_chain)

    else_step: Runnable[IntermediateType, IntermediateType] = RunnableBranch(
        (lambda x: x["error"] is None, RunnablePassthrough()),
        revise_step | validation_step,
    ).with_types(input_type=IntermediateType)

    for _ in range(max(0, input["max_revisions"] - 1)):
        else_step = RunnableBranch(
            (lambda x: x["error"] is None, RunnablePassthrough()),
            revise_step | validation_step | else_step,
        )
    return else_step


revise_lambda = RunnableLambda(revise_loop)


class InputType(BaseModel):
    query: str
    max_revisions: conint(ge=1, le=10) = 5


chain: Runnable[Any, Any] = (
    {
        "original_prompt": generate_prompt,
        "max_revisions": itemgetter("max_revisions"),
    }
    | RunnablePassthrough().assign(
        completion=(
            RunnableLambda(itemgetter("original_prompt"))
            | task_function_call_model
            | output_parser
        )
    )
    | validation_step
    | revise_lambda
    | RunnableLambda(itemgetter("completion"))
).with_types(input_type=InputType)