forked from realpython/materials
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathapp.py
More file actions
99 lines (79 loc) · 3.68 KB
/
Copy pathapp.py
File metadata and controls
99 lines (79 loc) · 3.68 KB
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
import argparse
import os
import tomllib
from pathlib import Path
import openai
# Authenticate
openai.api_key = os.getenv("OPENAI_API_KEY")
class Settings(dict):
"""Handle loading and accessing application settings from file."""
@classmethod
def load(cls, path) -> "Settings":
"""Load TOML settings file and pass it to class constuctor."""
with path.open("rb") as file:
return cls(tomllib.load(file))
def __init__(self, *args, **kwargs) -> None:
"""Add general settings and prompts as instance attributes."""
super().__init__(*args, **kwargs)
# Settings
self.chat_models = self["general"]["chat_models"]
self.model = self["general"]["model"]
self.max_tokens = self["general"]["max_tokens"]
self.temperature = self["general"]["temperature"]
self.model_supports_chat_completions = self.model in self.chat_models
# Prompts
self.instruction_prompt = self["prompts"]["instruction_prompt"]
self.role_prompt = self["prompts"]["role_prompt"]
self.positive_example = self["prompts"]["positive_example"]
self.positive_reasoning = self["prompts"]["positive_reasoning"]
self.positive_output = self["prompts"]["positive_output"]
self.negative_example = self["prompts"]["negative_example"]
self.negative_reasoning = self["prompts"]["negative_reasoning"]
self.negative_output = self["prompts"]["negative_output"]
def parse_args() -> argparse.Namespace:
"""Parse command-line input."""
parser = argparse.ArgumentParser()
parser.add_argument("file_path", type=Path, help="Path to the input file")
return parser.parse_args()
def main(args: argparse.Namespace) -> None:
file_content = args.file_path.read_text("utf-8")
settings = Settings.load(Path("settings.toml"))
if settings.model_supports_chat_completions:
print(get_chat_completion(file_content, settings))
else:
print(get_completion(file_content, settings))
def get_completion(content: str, settings: Settings) -> str:
"""Send a request to the /completions endpoint."""
response = openai.Completion.create(
model=settings.model,
prompt=assemble_prompt(content, settings),
max_tokens=settings.max_tokens,
temperature=settings.temperature,
)
return response["choices"][0]["text"]
def get_chat_completion(content: str, settings: Settings) -> str:
"""Send a request to the /chat/completions endpoint."""
response = openai.ChatCompletion.create(
model=settings.model,
messages=assemble_chat_messages(content, settings),
temperature=settings.temperature,
)
return response["choices"][0]["message"]["content"]
def assemble_prompt(content: str, settings: Settings) -> str:
"""Combine all text input into a single prompt."""
return f">>>>>\n{content}\n<<<<<\n\n" + settings.instruction_prompt
def assemble_chat_messages(content: str, settings: Settings) -> list[dict]:
"""Combine all messages into a well-formatted dictionary."""
return [
{"role": "system", "content": settings.role_prompt},
{"role": "user", "content": settings.negative_example},
{"role": "system", "content": settings.negative_reasoning},
{"role": "assistant", "content": settings.negative_output},
{"role": "user", "content": settings.positive_example},
{"role": "system", "content": settings.positive_reasoning},
{"role": "assistant", "content": settings.positive_output},
{"role": "user", "content": f">>>>>\n{content}\n<<<<<"},
{"role": "user", "content": settings.instruction_prompt},
]
if __name__ == "__main__":
main(parse_args())