-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
49 lines (41 loc) · 1.47 KB
/
main.py
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
import os
import discord
from src.discordBot import DiscordClient, Sender
from src.logger import logger
import google.generativeai as palm
from src.server import keep_alive
# Set up the PaLM API
palm.configure(api_key = "PaLM_API_KEY")
model = palm.types.Model(
name = "models/chat-bison-001",
base_model_id= "chat-bison",
version= "001",
display_name= "the Bison",
description= "the PaLM Bison",
input_token_limit= 145,
output_token_limit= 145,
supported_generation_methods= ["generateMessage"],
)
def chat_completion(message: str) -> str:
# Get the PaLM API response
response = palm.chat(messages=message, candidate_count = 0)
output = response.messages[1]
return output['content']
def run():
client = DiscordClient()
sender = Sender()
@client.tree.command(name="palm", description="Have a chat with PaLM Bison")
async def chat(interaction: discord.Interaction, *, message: str):
user_id = interaction.user.id
if interaction.user == client.user:
return
await interaction.response.defer()
receive = chat_completion(message)
# Truncate the length of the message to 2000 characters, but set this to 2000 chars would still cause this problem. so I set it to 1900.
#if len(receive) > 1900:
# receive = receive[:1900]
await sender.send_message(interaction, message, receive)
client.run("DISCORD_TOKEN")
if __name__ == '__main__':
keep_alive()
run()