Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
task.py 3.29 KB
Newer Older
Timothy J. Baek's avatar
refac: title generation  
Timothy J. Baek committed
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 
import re
import math

from datetime import datetime
from typing import Optional


def prompt_template(
    template: str, user_name: str = None, current_location: str = None
) -> str:
    # Get the current date
    current_date = datetime.now()

    # Format the date to YYYY-MM-DD
    formatted_date = current_date.strftime("%Y-%m-%d")

    # Replace {{CURRENT_DATE}} in the template with the formatted date
    template = template.replace("{{CURRENT_DATE}}", formatted_date)

    if user_name:
        # Replace {{USER_NAME}} in the template with the user's name
        template = template.replace("{{USER_NAME}}", user_name)

    if current_location:
        # Replace {{CURRENT_LOCATION}} in the template with the current location
        template = template.replace("{{CURRENT_LOCATION}}", current_location)

    return template


def title_generation_template(
    template: str, prompt: str, user: Optional[dict] = None
) -> str:
    def replacement_function(match):
        full_match = match.group(0)
        start_length = match.group(1)
        end_length = match.group(2)
        middle_length = match.group(3)

        if full_match == "{{prompt}}":
            return prompt
        elif start_length is not None:
            return prompt[: int(start_length)]
        elif end_length is not None:
            return prompt[-int(end_length) :]
        elif middle_length is not None:
            middle_length = int(middle_length)
            if len(prompt) <= middle_length:
                return prompt
            start = prompt[: math.ceil(middle_length / 2)]
            end = prompt[-math.floor(middle_length / 2) :]
            return f"{start}...{end}"
        return ""

    template = re.sub(
        r"{{prompt}}|{{prompt:start:(\d+)}}|{{prompt:end:(\d+)}}|{{prompt:middletruncate:(\d+)}}",
        replacement_function,
        template,
    )

    template = prompt_template(
        template,
        **(
            {"user_name": user.get("name"), "current_location": user.get("location")}
            if user
            else {}
        ),
    )

    return template
Timothy J. Baek's avatar
refac: search query task  
Timothy J. Baek committed
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 


def search_query_generation_template(
    template: str, prompt: str, user: Optional[dict] = None
) -> str:

    def replacement_function(match):
        full_match = match.group(0)
        start_length = match.group(1)
        end_length = match.group(2)
        middle_length = match.group(3)

        if full_match == "{{prompt}}":
            return prompt
        elif start_length is not None:
            return prompt[: int(start_length)]
        elif end_length is not None:
            return prompt[-int(end_length) :]
        elif middle_length is not None:
            middle_length = int(middle_length)
            if len(prompt) <= middle_length:
                return prompt
            start = prompt[: math.ceil(middle_length / 2)]
            end = prompt[-math.floor(middle_length / 2) :]
            return f"{start}...{end}"
        return ""

    template = re.sub(
        r"{{prompt}}|{{prompt:start:(\d+)}}|{{prompt:end:(\d+)}}|{{prompt:middletruncate:(\d+)}}",
        replacement_function,
        template,
    )

    template = prompt_template(
        template,
        **(
            {"user_name": user.get("name"), "current_location": user.get("location")}
            if user
            else {}
        ),
    )
    return template