use ollama client
Michael Yang
parent
35c6f1a270
commit
5ba591fbfb
286
discollama.py
286
discollama.py
|
|
@ -1,160 +1,180 @@
|
|||
import io
|
||||
import os
|
||||
import json
|
||||
import aiohttp
|
||||
import discord
|
||||
import asyncio
|
||||
import argparse
|
||||
from redis import Redis
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
import logging
|
||||
import ollama
|
||||
import discord
|
||||
import redis
|
||||
|
||||
intents = discord.Intents.default()
|
||||
intents.message_content = True
|
||||
from logging import getLogger
|
||||
|
||||
client = discord.Client(intents=intents)
|
||||
# piggy back on the logger discord.py set up
|
||||
logging = getLogger('discord.discollama')
|
||||
|
||||
|
||||
@client.event
|
||||
async def on_ready():
|
||||
class Response:
|
||||
def __init__(self, message):
|
||||
self.message = message
|
||||
self.channel = message.channel
|
||||
|
||||
self.r = None
|
||||
self.sb = io.StringIO()
|
||||
|
||||
async def write(self, s, end=''):
|
||||
if self.sb.seek(0, io.SEEK_END) + len(s) + len(end) > 2000:
|
||||
self.r = None
|
||||
self.sb.seek(0, io.SEEK_SET)
|
||||
self.sb.truncate()
|
||||
|
||||
self.sb.write(s)
|
||||
|
||||
if self.r:
|
||||
await self.r.edit(content=self.sb.getvalue() + end)
|
||||
return
|
||||
|
||||
if self.channel.type == discord.ChannelType.text:
|
||||
self.channel = await self.channel.create_thread(name='Discollama Says', message=self.message, auto_archive_duration=60)
|
||||
|
||||
self.r = await self.channel.send(self.sb.getvalue())
|
||||
|
||||
|
||||
class Discollama:
|
||||
def __init__(self, ollama, discord, redis):
|
||||
self.ollama = ollama
|
||||
self.discord = discord
|
||||
self.redis = redis
|
||||
|
||||
# register event handlers
|
||||
self.discord.event(self.on_ready)
|
||||
self.discord.event(self.on_message)
|
||||
|
||||
async def on_ready(self):
|
||||
activity = discord.Activity(name='Discollama', state='Ask me anything!', type=discord.ActivityType.custom)
|
||||
await self.discord.change_presence(activity=activity)
|
||||
|
||||
logging.info(
|
||||
'Ready! Invite URL: %s',
|
||||
discord.utils.oauth_url(
|
||||
client.application_id,
|
||||
permissions=discord.Permissions(read_messages=True, send_messages=True),
|
||||
self.discord.application_id,
|
||||
permissions=discord.Permissions(
|
||||
read_messages=True,
|
||||
send_messages=True,
|
||||
create_public_threads=True,
|
||||
),
|
||||
scopes=['bot'],
|
||||
))
|
||||
|
||||
|
||||
async def generate_response(prompt, context=[]):
|
||||
body = {
|
||||
key: value
|
||||
for key, value in {
|
||||
'model': args.ollama_model,
|
||||
'prompt': prompt,
|
||||
'context': context,
|
||||
}.items() if value
|
||||
}
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(
|
||||
f'http://{args.ollama_host}:{args.ollama_port}/api/generate',
|
||||
json=body) as r:
|
||||
async for line in r.content:
|
||||
yield json.loads(line)
|
||||
|
||||
|
||||
async def buffered_generate_response(prompt, context=[]):
|
||||
buffer = ''
|
||||
async for part in generate_response(prompt, context):
|
||||
if error := part.get('error'):
|
||||
raise Exception(error)
|
||||
|
||||
if part['done']:
|
||||
yield buffer, part
|
||||
break
|
||||
|
||||
buffer += part['response']
|
||||
if len(buffer) >= args.buffer_size:
|
||||
yield buffer, part
|
||||
buffer = ''
|
||||
|
||||
|
||||
def save_session(response, part):
|
||||
context = part.get('context', [])
|
||||
redis.json().set(f'ollama:{response.id}', '$', {'context': context})
|
||||
|
||||
redis.expire(f'ollama:{response.id}', 60 * 60 * 24 * 7)
|
||||
logging.info('saving message=%s: len(context)=%d', response.id, len(context))
|
||||
|
||||
|
||||
def load_session(reference):
|
||||
kwargs = {}
|
||||
if reference:
|
||||
context = redis.json().get(f'ollama:{reference.message_id}', '.context')
|
||||
kwargs['context'] = context or []
|
||||
|
||||
if kwargs.get('context'):
|
||||
logging.info(
|
||||
'loading message=%s: len(context)=%d',
|
||||
reference.message_id,
|
||||
len(kwargs['context']),
|
||||
),
|
||||
)
|
||||
|
||||
return kwargs
|
||||
|
||||
|
||||
@client.event
|
||||
async def on_message(message):
|
||||
if message.author == client.user:
|
||||
async def on_message(self, message):
|
||||
if self.discord.user == message.author:
|
||||
# don't respond to ourselves
|
||||
return
|
||||
|
||||
if client.user.id in message.raw_mentions:
|
||||
raw_content = message.content.replace(f'<@{client.user.id}>', '').strip()
|
||||
if raw_content.strip() == '':
|
||||
raw_content = 'Tell me about yourself.'
|
||||
if not self.discord.user.mentioned_in(message):
|
||||
# don't respond to messages that don't mention us
|
||||
return
|
||||
|
||||
response = None
|
||||
response_content = ''
|
||||
async with message.channel.typing():
|
||||
await message.add_reaction('🤔')
|
||||
content = message.content.replace(f'<@{self.discord.user.id}>', '').strip()
|
||||
if not content:
|
||||
content = 'Hi!'
|
||||
|
||||
channel = message.channel
|
||||
|
||||
context = []
|
||||
if reference := message.reference:
|
||||
if session := load_session(message.reference):
|
||||
context = session.get('context', [])
|
||||
else:
|
||||
reference_message = await message.channel.fetch_message(
|
||||
reference.message_id)
|
||||
reference_content = reference_message.content
|
||||
raw_content = '\n'.join([
|
||||
raw_content,
|
||||
'Use it to answer the prompt:',
|
||||
reference_content,
|
||||
])
|
||||
context = await self.load(message_id=reference.message_id)
|
||||
if not context:
|
||||
reference_message = await message.channel.fetch_message(reference.message_id)
|
||||
content = '\n'.join(
|
||||
[
|
||||
content,
|
||||
'Use this to answer the question if it is relevant, otherwise ignore it:',
|
||||
reference_message.content,
|
||||
]
|
||||
)
|
||||
|
||||
async for buffer, part in buffered_generate_response(
|
||||
raw_content,
|
||||
context=context,
|
||||
):
|
||||
response_content += buffer
|
||||
if part['done']:
|
||||
save_session(response, part)
|
||||
break
|
||||
if not context:
|
||||
context = await self.load(channel_id=channel.id)
|
||||
|
||||
if not response:
|
||||
response = await message.reply(response_content)
|
||||
await message.remove_reaction('🤔', client.user)
|
||||
continue
|
||||
r = Response(message)
|
||||
task = asyncio.create_task(self.thinking(message))
|
||||
async for part in self.generate(content, context):
|
||||
task.cancel()
|
||||
|
||||
if len(response_content) + 3 >= 2000:
|
||||
response = await response.reply(buffer)
|
||||
response_content = buffer
|
||||
continue
|
||||
await r.write(part['response'], end='...')
|
||||
|
||||
await response.edit(content=response_content + '...')
|
||||
await r.write('')
|
||||
await self.save(r.channel.id, message.id, part['context'])
|
||||
|
||||
await response.edit(content=response_content)
|
||||
|
||||
|
||||
default_ollama_host = os.getenv('OLLAMA_HOST', '127.0.0.1')
|
||||
default_ollama_port = os.getenv('OLLAMA_PORT', 11434)
|
||||
default_ollama_model = os.getenv('OLLAMA_MODEL', 'llama2')
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--ollama-host', default=default_ollama_host)
|
||||
parser.add_argument('--ollama-port', default=default_ollama_port, type=int)
|
||||
parser.add_argument('--ollama-model', default=default_ollama_model, type=str)
|
||||
|
||||
parser.add_argument('--redis-host', default='localhost')
|
||||
parser.add_argument('--redis-port', default=6379)
|
||||
|
||||
parser.add_argument('--buffer-size', default=32, type=int)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
try:
|
||||
redis = Redis(host=args.redis_host, port=args.redis_port)
|
||||
client.run(os.getenv('DISCORD_TOKEN'), root_logger=True)
|
||||
except KeyboardInterrupt:
|
||||
async def thinking(self, message, timeout=999):
|
||||
try:
|
||||
await message.add_reaction('🤔')
|
||||
async with message.channel.typing():
|
||||
await asyncio.sleep(timeout)
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
await message.remove_reaction('🤔', self.discord.user)
|
||||
|
||||
redis.close()
|
||||
async def generate(self, content, context):
|
||||
sb = io.StringIO()
|
||||
|
||||
t = datetime.now()
|
||||
async for part in await self.ollama.generate(model='llama2', prompt=content, context=context, stream=True):
|
||||
sb.write(part['response'])
|
||||
|
||||
if part['done'] or datetime.now() - t > timedelta(seconds=1):
|
||||
part['response'] = sb.getvalue()
|
||||
yield part
|
||||
t = datetime.now()
|
||||
sb.seek(0, io.SEEK_SET)
|
||||
sb.truncate()
|
||||
|
||||
async def save(self, channel_id, message_id, ctx: list[int]):
|
||||
self.redis.set(f'discollama:channel:{channel_id}', message_id, ex=60 * 60 * 24 * 7)
|
||||
self.redis.set(f'discollama:message:{message_id}', json.dumps(ctx), ex=60 * 60 * 24 * 7)
|
||||
|
||||
async def load(self, channel_id=None, message_id=None) -> list[int]:
|
||||
if channel_id:
|
||||
message_id = self.redis.get(f'discollama:channel:{channel_id}')
|
||||
|
||||
ctx = self.redis.get(f'discollama:message:{message_id}')
|
||||
return json.loads(ctx) if ctx else []
|
||||
|
||||
def run(self, token):
|
||||
try:
|
||||
self.discord.run(token)
|
||||
except Exception:
|
||||
self.redis.close()
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
parser.add_argument('--ollama-scheme', default=os.getenv('OLLAMA_SCHEME', 'http'), choices=['http', 'https'])
|
||||
parser.add_argument('--ollama-host', default=os.getenv('OLLAMA_HOST', '127.0.0.1'), type=str)
|
||||
parser.add_argument('--ollama-port', default=os.getenv('OLLAMA_PORT', 11434), type=int)
|
||||
parser.add_argument('--ollama-model', default=os.getenv('OLLAMA_MODEL', 'llama2'), type=str)
|
||||
|
||||
parser.add_argument('--redis-host', default=os.getenv('REDIS_HOST', '127.0.0.1'), type=str)
|
||||
parser.add_argument('--redis-port', default=os.getenv('REDIS_PORT', 6379), type=int)
|
||||
|
||||
parser.add_argument('--buffer-size', default=32, type=int)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
intents = discord.Intents.default()
|
||||
intents.message_content = True
|
||||
|
||||
Discollama(
|
||||
ollama.AsyncClient(base_url=f'{args.ollama_scheme}://{args.ollama_host}:{args.ollama_port}'),
|
||||
discord.Client(intents=intents),
|
||||
redis.Redis(host=args.redis_host, port=args.redis_port, db=0, decode_responses=True),
|
||||
).run(os.environ['DISCORD_TOKEN'])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
|
|
|
|||
Loading…
Reference in New Issue