cdnutter
/
discollama
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

formatting

pull/3/head
Michael Yang 2023-09-20 20:15:09 -07:00
parent 924017de97
commit c564d05de5
1 changed files with 93 additions and 82 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
discollama.py
View File

@ -22,7 +22,8 @@ async def on_ready():
discord.utils.oauth_url( discord.utils.oauth_url(
client.application_id, client.application_id,
permissions=discord.Permissions(read_messages=True, send_messages=True), permissions=discord.Permissions(read_messages=True, send_messages=True),
scopes=['bot'])) scopes=['bot'],
))
async def generate_response(prompt, context=[]): async def generate_response(prompt, context=[]):
@ -61,7 +62,12 @@ def save_session(response, chunk):
redis.hset(f'ollama:{response.id}', 'context', context) redis.hset(f'ollama:{response.id}', 'context', context)
redis.expire(f'ollama:{response.id}', 60 * 60 * 24 * 7) redis.expire(f'ollama:{response.id}', 60 * 60 * 24 * 7)
logging.info('saving message=%s: len(context)=%d', response.id, len(chunk['context'])) logging.info(
'saving message=%s: len(context)=%d',
response.id,
len(chunk['context']),
)
def load_session(reference): def load_session(reference):
kwargs = {} kwargs = {}
@ -73,7 +79,8 @@ def load_session(reference):
logging.info( logging.info(
'loading message=%s: len(context)=%d', 'loading message=%s: len(context)=%d',
reference.message_id, reference.message_id,
len(kwargs['context'])) len(kwargs['context']),
)
return kwargs return kwargs
@ -98,7 +105,8 @@ async def on_message(message):
if session := load_session(message.reference): if session := load_session(message.reference):
context = session.get('context', []) context = session.get('context', [])
reference_message = await message.channel.fetch_message(reference.message_id) reference_message = await message.channel.fetch_message(
reference.message_id)
reference_content = reference_message.content reference_content = reference_message.content
raw_content = '\n'.join([ raw_content = '\n'.join([
raw_content, raw_content,
@ -106,7 +114,10 @@ async def on_message(message):
reference_content, reference_content,
]) ])
async for buffer, chunk in buffered_generate_response(raw_content, context=context): async for buffer, chunk in buffered_generate_response(
raw_content,
context=context,
):
response_content += buffer response_content += buffer
if chunk['done']: if chunk['done']:
save_session(response, chunk) save_session(response, chunk)