summaryrefslogtreecommitdiff
path: root/py/chat.py
blob: b6dcd3512675546b1507f12767327af87aeaffb3 (plain)
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
import requests
import sys
import os

file_content = vim.eval("prompt")

config_file_path = os.path.join(os.path.expanduser("~"), ".config/openai.token")

api_key = os.getenv("OPENAI_API_KEY")

try:
    with open(config_file_path, 'r') as file:
        api_key = file.read()
except Exception:
    pass

api_key = api_key.strip()

lines = file_content.splitlines()
messages = []

for line in lines:
    if line.startswith(">>> system"):
        messages.append({"role": "system", "content": ""})
        continue
    if line.startswith(">>> user"):
        messages.append({"role": "user", "content": ""})
        continue
    if line.startswith("<<< assistant"):
        messages.append({"role": "assistant", "content": ""})
        continue
    if not messages:
        continue
    messages[-1]["content"] += "\n" + line

if not messages:
    file_content = ">>> user\n\n" + file_content
    messages.append({"role": "user", "content": file_content })


url = "https://api.openai.com/v1/chat/completions"
headers = {
    'Content-Type': 'application/json',
    'Authorization': F"Bearer {api_key}"
}
data = {
    "model": "gpt-3.5-turbo",
    "messages": messages
}
response = requests.post(url, headers=headers, json=data)
response = response.json()

answer = response['choices'][0]['message']['content']

output = f"{file_content.strip()}\n\n<<< assistant\n\n{answer.strip()}\n\n>>> user\n"