CoSQL-LLM/llm.py at main · emorynlp/CoSQL-LLM · 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

import openai
import time
import anthropic

openai.api_key = 'KEY'
claude_key = 'KEY'


gpt_model = "gpt-4o" # gpt 4o
claude_model = "claude-3-opus-20240229" # claude 3 opus
# client = anthropic.Anthropic(
#     api_key=claude_key,
# )


def claudecall(system: str, user: str, history: "list[dict[str, str]] | None" = None):
	while(True):
		try:
			messages = []
			if(history):
				for h in history:
					messages.append(h)
			messages.append({"role": "user", "content":
						user})
			response = client.messages.create(
				model=claude_model,
				max_tokens = 1000,
				temperature=0,
				system=system,
				messages=messages
			)
			break
		except Exception as e:
			print(e)
			time.sleep(1)
	return response.content[0].text

def gptcall(system: str, user: str, history: "list[dict[str, str]] | None" = None):
	while(True):
		try:
			messages = [
				{"role": "system", "content": system}
			]
			if(history):
				for h in history:
					messages.append(h)
			messages.append({"role": "user", "content": user})
			# print(messages)
			response = openai.ChatCompletion.create( # type: ignore
				model=gpt_model,
				temperature=0,
				messages=messages
			)
			break
		except Exception as e:
			print(e)
			time.sleep(1)

	r = response['choices'][0]['message']['content'] # type: ignore
	return str(r) # type: ignore