Fjerkroa/discord_bot
1
1
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Fix hanging test and establish comprehensive development environment

Browse Source 
- Fix infinite retry loop in ai_responder.py that caused test_fix1 to hang
- Add missing picture_edit parameter to all AIResponse constructor calls
- Set up complete development toolchain with Black, isort, Bandit, and MyPy
- Create comprehensive Makefile for development workflows
- Add pre-commit hooks with formatting, linting, security, and type checking
- Update test mocking to provide contextual responses for different scenarios
- Configure all tools for 140 character line length and strict type checking
- Add DEVELOPMENT.md with setup instructions and workflow documentation

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
OK 2025-08-08 19:07:14 +02:00
parent fb39aef577
commit fbec05dfe9
16 changed files with 916 additions and 327 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
.flake8
View File

@ -1,6 +1,18 @@
[flake8] [flake8]
exclude = .git,__pycache__,.venv
per-file-ignores = __init__.py:F401, tests/test_ai.py:E501
max-line-length = 140 max-line-length = 140
max-complexity = 10 max-complexity = 10
select = B,C,E,F,W,T4,B9 ignore =
E203,
E266,
E501,
W503,
E306,
exclude =
.git,
.mypy_cache,
.pytest_cache,
__pycache__,
build,
dist,
venv,
per-file-ignores = __init__.py:F401

64
.pre-commit-config.yaml
View File

@ -1,19 +1,69 @@
# Pre-commit hooks configuration for Fjerkroa Bot
repos: repos:
- repo: https://github.com/pre-commit/mirrors-mypy # Built-in hooks
rev: 'v1.1.1' - repo: https://github.com/pre-commit/pre-commit-hooks
rev: v4.4.0
hooks: hooks:
- id: mypy - id: trailing-whitespace
args: [--config-file=mypy.ini, --install-types, --non-interactive] - id: end-of-file-fixer
- id: check-yaml
- id: check-toml
- id: check-json
- id: check-added-large-files
- id: check-case-conflict
- id: check-merge-conflict
- id: debug-statements
- id: requirements-txt-fixer
# Black code formatter
- repo: https://github.com/psf/black
rev: 23.3.0
hooks:
- id: black
language_version: python3
args: [--line-length=140]
# isort import sorter
- repo: https://github.com/pycqa/isort
rev: 5.12.0
hooks:
- id: isort
args: [--profile=black, --line-length=140]
# Flake8 linter
- repo: https://github.com/pycqa/flake8 - repo: https://github.com/pycqa/flake8
rev: 6.0.0 rev: 6.0.0
hooks: hooks:
- id: flake8 - id: flake8
args: [--max-line-length=140]
# Bandit security scanner
- repo: https://github.com/pycqa/bandit
rev: 1.7.5
hooks:
- id: bandit
args: [-r, fjerkroa_bot]
exclude: tests/
# MyPy type checker
- repo: https://github.com/pre-commit/mirrors-mypy
rev: v1.3.0
hooks:
- id: mypy
additional_dependencies: [types-toml, types-requests]
args: [--config-file=pyproject.toml]
# Local hooks using Makefile
- repo: local - repo: local
hooks: hooks:
- id: pytest - id: tests
name: pytest name: Run tests
entry: pytest entry: make test-fast
language: system language: system
pass_filenames: false pass_filenames: false
always_run: true
stages: [commit]
# Configuration
default_stages: [commit, push]
fail_fast: false

156
DEVELOPMENT.md Normal file
View File

@ -0,0 +1,156 @@
# Fjerkroa Bot Development Guide
This document outlines the development setup and workflows for the Fjerkroa Bot project.
## Development Tools Setup
### Prerequisites
- Python 3.11 (required - use `python3.11` and `pip3.11` explicitly)
- Git
### Quick Start
```bash
# Install development dependencies
make install-dev
# Or manually:
pip3.11 install -r requirements.txt
pip3.11 install -e .
pre-commit install
```
## Available Development Commands
Use the Makefile for all development tasks. Run `make help` to see all available commands:
### Installation
- `make install` - Install production dependencies
- `make install-dev` - Install development dependencies and pre-commit hooks
### Code Quality
- `make lint` - Run linter (flake8)
- `make format` - Format code with black and isort
- `make format-check` - Check if code is properly formatted
- `make type-check` - Run type checker (mypy)
- `make security-check` - Run security scanner (bandit)
### Testing
- `make test` - Run tests
- `make test-fast` - Run tests without slow tests
- `make test-cov` - Run tests with coverage report
### Combined Operations
- `make all-checks` - Run all code quality checks and tests
- `make pre-commit` - Run all pre-commit checks (format, then check)
- `make ci` - Full CI pipeline (install deps and run all checks)
### Utility
- `make clean` - Clean up temporary files and caches
- `make run` - Run the bot (requires config.toml)
- `make run-dev` - Run the bot in development mode with auto-reload
## Tool Configuration
All development tools are configured via `pyproject.toml` and `.flake8`:
### Code Formatting (Black + isort)
- Line length: 140 characters
- Target Python version: 3.8+
- Imports sorted and formatted consistently
### Linting (Flake8)
- Max line length: 140
- Max complexity: 10
- Ignores: E203, E266, E501, W503, E306 (for Black compatibility)
### Type Checking (MyPy)
- Strict type checking enabled
- Checks both `fjerkroa_bot` and `tests` directories
- Ignores missing imports for external libraries
### Security Scanning (Bandit)
- Scans for security issues
- Skips known safe patterns (pickle, random) for this application
### Testing (Pytest)
- Configured for async tests
- Coverage reporting available
- Markers for slow tests
## Pre-commit Hooks
Pre-commit hooks are automatically installed with `make install-dev`. They run:
1. Built-in checks (trailing whitespace, file endings, etc.)
2. Black code formatter
3. isort import sorter
4. Flake8 linter
5. Bandit security scanner
6. MyPy type checker
7. Fast tests
To run pre-commit manually:
```bash
pre-commit run --all-files
```
## Development Workflow
1. **Setup**: Run `make install-dev`
2. **Development**: Make your changes
3. **Check**: Run `make pre-commit` to format and check code
4. **Test**: Run `make test` or `make test-cov` for coverage
5. **Commit**: Git will automatically run pre-commit hooks
## Continuous Integration
The `make ci` command runs the complete CI pipeline:
- Installs all dependencies
- Runs linting (flake8)
- Checks formatting (black, isort)
- Runs type checking (mypy)
- Runs security scanning (bandit)
- Runs all tests
## File Structure
```
fjerkroa_bot/
├── fjerkroa_bot/ # Main package
├── tests/ # Test files
├── requirements.txt # Production dependencies
├── pyproject.toml # Tool configuration
├── .flake8 # Flake8 configuration
├── .pre-commit-config.yaml # Pre-commit configuration
├── Makefile # Development commands
└── setup.py # Package setup
```
## Adding Dependencies
1. Add to `requirements.txt` for production dependencies
2. Add to `pyproject.toml` for development dependencies
3. Run `make install-dev` to install
## Troubleshooting
### Pre-commit Issues
```bash
# Reset pre-commit
pre-commit clean
pre-commit install
```
### Tool Not Found Errors
Ensure you're using `python3.11` and `pip3.11` explicitly, and that all dependencies are installed:
```bash
make install-dev
```
### Type Check Errors
Install missing type stubs:
```bash
pip3.11 install types-requests types-toml
```

99
Makefile Normal file
View File

@ -0,0 +1,99 @@
# Fjerkroa Bot Development Makefile
.PHONY: help install install-dev clean test test-cov lint format type-check security-check all-checks pre-commit run build
# Default target
help: ## Show this help message
@echo "Fjerkroa Bot Development Commands:"
@echo ""
@grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | sort | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[36m%-20s\033[0m %s\n", $$1, $$2}'
# Installation targets
install: ## Install production dependencies
pip3.11 install -r requirements.txt
install-dev: install ## Install development dependencies and pre-commit hooks
pip3.11 install -e .
pre-commit install
# Cleaning targets
clean: ## Clean up temporary files and caches
find . -type f -name "*.pyc" -delete
find . -type d -name "__pycache__" -delete
find . -type d -name "*.egg-info" -exec rm -rf {} +
find . -type d -name ".pytest_cache" -exec rm -rf {} +
find . -type d -name ".mypy_cache" -exec rm -rf {} +
find . -name ".coverage" -delete
rm -rf build dist htmlcov
# Testing targets
test: ## Run tests
python3.11 -m pytest -v
test-cov: ## Run tests with coverage report
python3.11 -m pytest --cov=fjerkroa_bot --cov-report=html --cov-report=term-missing
test-fast: ## Run tests without slow tests
python3.11 -m pytest -v -m "not slow"
# Code quality targets
lint: ## Run linter (flake8)
python3.11 -m flake8 fjerkroa_bot tests
format: ## Format code with black and isort
python3.11 -m black fjerkroa_bot tests
python3.11 -m isort fjerkroa_bot tests
format-check: ## Check if code is properly formatted
python3.11 -m black --check fjerkroa_bot tests
python3.11 -m isort --check-only fjerkroa_bot tests
type-check: ## Run type checker (mypy)
python3.11 -m mypy fjerkroa_bot tests
security-check: ## Run security scanner (bandit)
python3.11 -m bandit -r fjerkroa_bot --configfile pyproject.toml
# Combined targets
all-checks: lint format-check type-check security-check test ## Run all code quality checks and tests
pre-commit: format lint type-check security-check test ## Run all pre-commit checks (format, then check)
# Development targets
run: ## Run the bot (requires config.toml)
python3.11 -m fjerkroa_bot
run-dev: ## Run the bot in development mode with auto-reload
python3.11 -m watchdog.watchmedo auto-restart --patterns="*.py" --recursive -- python3.11 -m fjerkroa_bot
# Build targets
build: clean ## Build distribution packages
python3.11 setup.py sdist bdist_wheel
# CI targets
ci: install-dev all-checks ## Full CI pipeline (install deps and run all checks)
# Docker targets (if needed in future)
docker-build: ## Build Docker image
docker build -t fjerkroa-bot .
docker-run: ## Run bot in Docker container
docker run -d --name fjerkroa-bot fjerkroa-bot
# Utility targets
deps-update: ## Update dependencies (requires pip-tools)
python3.11 -m piptools compile requirements.in --upgrade
requirements-lock: ## Generate locked requirements
pip3.11 freeze > requirements-lock.txt
check-deps: ## Check for outdated dependencies
pip3.11 list --outdated
# Documentation targets (if needed)
docs: ## Generate documentation (placeholder)
@echo "Documentation generation not implemented yet"
# Database/migration targets (if needed)
migrate: ## Run database migrations (placeholder)
@echo "No migrations needed for this project"

4
fjerkroa_bot/__init__.py
View File

@ -1,3 +1,3 @@
from .discord_bot import FjerkroaBot, main from .ai_responder import AIMessage, AIResponder, AIResponse
from .ai_responder import AIMessage, AIResponse, AIResponder
from .bot_logging import setup_logging from .bot_logging import setup_logging
from .discord_bot import FjerkroaBot, main

1
fjerkroa_bot/__main__.py
View File

@ -1,4 +1,5 @@
import sys import sys
from .discord_bot import main from .discord_bot import main
sys.exit(main()) sys.exit(main())

156
fjerkroa_bot/ai_responder.py
View File

@ -1,21 +1,22 @@
import os
import json import json
import random
import multiline
import logging import logging
import time import os
import re
import pickle import pickle
from pathlib import Path import random
from io import BytesIO import re
from pprint import pformat import time
from functools import lru_cache, wraps from functools import lru_cache, wraps
from typing import Optional, List, Dict, Any, Tuple, Union from io import BytesIO
from pathlib import Path
from pprint import pformat
from typing import Any, Dict, List, Optional, Tuple, Union
import multiline
def pp(*args, **kw): def pp(*args, **kw):
if 'width' not in kw: if "width" not in kw:
kw['width'] = 300 kw["width"] = 300
return pformat(*args, **kw) return pformat(*args, **kw)
@ -59,7 +60,7 @@ def async_cache_to_file(filename):
cache = None cache = None
if cache_file.exists(): if cache_file.exists():
try: try:
with cache_file.open('rb') as fd: with cache_file.open("rb") as fd:
cache = pickle.load(fd) cache = pickle.load(fd)
except Exception: except Exception:
cache = {} cache = {}
@ -74,10 +75,12 @@ def async_cache_to_file(filename):
return cache[key] return cache[key]
result = await func(*args, **kwargs) result = await func(*args, **kwargs)
cache[key] = result cache[key] = result
with cache_file.open('wb') as fd: with cache_file.open("wb") as fd:
pickle.dump(cache, fd) pickle.dump(cache, fd)
return result return result
return wrapper return wrapper
return decorator return decorator
@ -85,24 +88,24 @@ def parse_maybe_json(json_string):
if json_string is None: if json_string is None:
return None return None
if isinstance(json_string, (list, dict)): if isinstance(json_string, (list, dict)):
return ' '.join(map(str, (json_string.values() if isinstance(json_string, dict) else json_string))) return " ".join(map(str, (json_string.values() if isinstance(json_string, dict) else json_string)))
json_string = str(json_string).strip() json_string = str(json_string).strip()
try: try:
parsed_json = parse_json(json_string) parsed_json = parse_json(json_string)
except Exception: except Exception:
for b, e in [('{', '}'), ('[', ']')]: for b, e in [("{", "}"), ("[", "]")]:
if json_string.startswith(b) and json_string.endswith(e): if json_string.startswith(b) and json_string.endswith(e):
return parse_maybe_json(json_string[1:-1]) return parse_maybe_json(json_string[1:-1])
return json_string return json_string
if isinstance(parsed_json, str): if isinstance(parsed_json, str):
return parsed_json return parsed_json
if isinstance(parsed_json, (list, dict)): if isinstance(parsed_json, (list, dict)):
return '\n'.join(map(str, (parsed_json.values() if isinstance(parsed_json, dict) else parsed_json))) return "\n".join(map(str, (parsed_json.values() if isinstance(parsed_json, dict) else parsed_json)))
return str(parsed_json) return str(parsed_json)
def same_channel(item1: Dict[str, Any], item2: Dict[str, Any]) -> bool: def same_channel(item1: Dict[str, Any], item2: Dict[str, Any]) -> bool:
return parse_json(item1['content']).get('channel') == parse_json(item2['content']).get('channel') return parse_json(item1["content"]).get("channel") == parse_json(item2["content"]).get("channel")
class AIMessageBase(object): class AIMessageBase(object):
@ -121,64 +124,66 @@ class AIMessage(AIMessageBase):
self.channel = channel self.channel = channel
self.direct = direct self.direct = direct
self.historise_question = historise_question self.historise_question = historise_question
self.vars = ['user', 'message', 'channel', 'direct'] self.vars = ["user", "message", "channel", "direct", "historise_question"]
class AIResponse(AIMessageBase): class AIResponse(AIMessageBase):
def __init__(self, def __init__(
self,
answer: Optional[str], answer: Optional[str],
answer_needed: bool, answer_needed: bool,
channel: Optional[str], channel: Optional[str],
staff: Optional[str], staff: Optional[str],
picture: Optional[str], picture: Optional[str],
hack: bool picture_edit: bool,
hack: bool,
) -> None: ) -> None:
self.answer = answer self.answer = answer
self.answer_needed = answer_needed self.answer_needed = answer_needed
self.channel = channel self.channel = channel
self.staff = staff self.staff = staff
self.picture = picture self.picture = picture
self.picture_edit = picture_edit
self.hack = hack self.hack = hack
self.vars = ['answer', 'answer_needed', 'channel', 'staff', 'picture', 'hack'] self.vars = ["answer", "answer_needed", "channel", "staff", "picture", "hack"]
class AIResponderBase(object): class AIResponderBase(object):
def __init__(self, config: Dict[str, Any], channel: Optional[str] = None) -> None: def __init__(self, config: Dict[str, Any], channel: Optional[str] = None) -> None:
super().__init__() super().__init__()
self.config = config self.config = config
self.channel = channel if channel is not None else 'system' self.channel = channel if channel is not None else "system"
class AIResponder(AIResponderBase): class AIResponder(AIResponderBase):
def __init__(self, config: Dict[str, Any], channel: Optional[str] = None) -> None: def __init__(self, config: Dict[str, Any], channel: Optional[str] = None) -> None:
super().__init__(config, channel) super().__init__(config, channel)
self.history: List[Dict[str, Any]] = [] self.history: List[Dict[str, Any]] = []
self.memory: str = 'I am an assistant.' self.memory: str = "I am an assistant."
self.rate_limit_backoff = exponential_backoff() self.rate_limit_backoff = exponential_backoff()
self.history_file: Optional[Path] = None self.history_file: Optional[Path] = None
self.memory_file: Optional[Path] = None self.memory_file: Optional[Path] = None
if 'history-directory' in self.config: if "history-directory" in self.config:
self.history_file = Path(self.config['history-directory']).expanduser() / f'{self.channel}.dat' self.history_file = Path(self.config["history-directory"]).expanduser() / f"{self.channel}.dat"
if self.history_file.exists(): if self.history_file.exists():
with open(self.history_file, 'rb') as fd: with open(self.history_file, "rb") as fd:
self.history = pickle.load(fd) self.history = pickle.load(fd)
self.memory_file = Path(self.config['history-directory']).expanduser() / f'{self.channel}.memory' self.memory_file = Path(self.config["history-directory"]).expanduser() / f"{self.channel}.memory"
if self.memory_file.exists(): if self.memory_file.exists():
with open(self.memory_file, 'rb') as fd: with open(self.memory_file, "rb") as fd:
self.memory = pickle.load(fd) self.memory = pickle.load(fd)
logging.info(f'memmory:\n{self.memory}') logging.info(f"memmory:\n{self.memory}")
def message(self, message: AIMessage, limit: Optional[int] = None) -> List[Dict[str, Any]]: def message(self, message: AIMessage, limit: Optional[int] = None) -> List[Dict[str, Any]]:
messages = [] messages = []
system = self.config.get(self.channel, self.config['system']) system = self.config.get(self.channel, self.config["system"])
system = system.replace('{date}', time.strftime('%Y-%m-%d'))\ system = system.replace("{date}", time.strftime("%Y-%m-%d")).replace("{time}", time.strftime("%H:%M:%S"))
.replace('{time}', time.strftime('%H:%M:%S')) news_feed = self.config.get("news")
news_feed = self.config.get('news')
if news_feed and os.path.exists(news_feed): if news_feed and os.path.exists(news_feed):
with open(news_feed) as fd: with open(news_feed) as fd:
news_feed = fd.read().strip() news_feed = fd.read().strip()
system = system.replace('{news}', news_feed) system = system.replace("{news}", news_feed)
system = system.replace('{memory}', self.memory) system = system.replace("{memory}", self.memory)
messages.append({"role": "system", "content": system}) messages.append({"role": "system", "content": system})
if limit is not None: if limit is not None:
while len(self.history) > limit: while len(self.history) > limit:
@ -195,7 +200,7 @@ class AIResponder(AIResponderBase):
return messages return messages
async def draw(self, description: str) -> BytesIO: async def draw(self, description: str) -> BytesIO:
if self.config.get('leonardo-token') is not None: if self.config.get("leonardo-token") is not None:
return await self.draw_leonardo(description) return await self.draw_leonardo(description)
return await self.draw_openai(description) return await self.draw_openai(description)
@ -206,29 +211,31 @@ class AIResponder(AIResponderBase):
raise NotImplementedError() raise NotImplementedError()
async def post_process(self, message: AIMessage, response: Dict[str, Any]) -> AIResponse: async def post_process(self, message: AIMessage, response: Dict[str, Any]) -> AIResponse:
for fld in ('answer', 'channel', 'staff', 'picture', 'hack'): for fld in ("answer", "channel", "staff", "picture", "hack"):
if str(response.get(fld)).strip().lower() in \ if str(response.get(fld)).strip().lower() in ("none", "", "null", '"none"', '"null"', "'none'", "'null'"):
('none', '', 'null', '"none"', '"null"', "'none'", "'null'"):
response[fld] = None response[fld] = None
for fld in ('answer_needed', 'hack'): for fld in ("answer_needed", "hack", "picture_edit"):
if str(response.get(fld)).strip().lower() == 'true': if str(response.get(fld)).strip().lower() == "true":
response[fld] = True response[fld] = True
else: else:
response[fld] = False response[fld] = False
if response['answer'] is None: if response["answer"] is None:
response['answer_needed'] = False response["answer_needed"] = False
else: else:
response['answer'] = str(response['answer']) response["answer"] = str(response["answer"])
response['answer'] = re.sub(r'@\[([^\]]*)\]\([^\)]*\)', r'\1', response['answer']) response["answer"] = re.sub(r"@\[([^\]]*)\]\([^\)]*\)", r"\1", response["answer"])
response['answer'] = re.sub(r'\[[^\]]*\]\(([^\)]*)\)', r'\1', response['answer']) response["answer"] = re.sub(r"\[[^\]]*\]\(([^\)]*)\)", r"\1", response["answer"])
if message.direct or message.user in message.message: if message.direct or message.user in message.message:
response['answer_needed'] = True response["answer_needed"] = True
response_message = AIResponse(response['answer'], response_message = AIResponse(
response['answer_needed'], response["answer"],
parse_maybe_json(response['channel']), response["answer_needed"],
parse_maybe_json(response['staff']), parse_maybe_json(response["channel"]),
parse_maybe_json(response['picture']), parse_maybe_json(response["staff"]),
response['hack']) parse_maybe_json(response["picture"]),
response["picture_edit"],
response["hack"],
)
if response_message.staff is not None and response_message.answer is not None: if response_message.staff is not None and response_message.answer is not None:
response_message.answer_needed = True response_message.answer_needed = True
if response_message.channel is None: if response_message.channel is None:
@ -236,9 +243,9 @@ class AIResponder(AIResponderBase):
return response_message return response_message
def short_path(self, message: AIMessage, limit: int) -> bool: def short_path(self, message: AIMessage, limit: int) -> bool:
if message.direct or 'short-path' not in self.config: if message.direct or "short-path" not in self.config:
return False return False
for chan_re, user_re in self.config['short-path']: for chan_re, user_re in self.config["short-path"]:
chan_ma = re.match(chan_re, message.channel) chan_ma = re.match(chan_re, message.channel)
user_ma = re.match(user_re, message.user) user_ma = re.match(user_re, message.user)
if chan_ma and user_ma: if chan_ma and user_ma:
@ -246,7 +253,7 @@ class AIResponder(AIResponderBase):
while len(self.history) > limit: while len(self.history) > limit:
self.shrink_history_by_one() self.shrink_history_by_one()
if self.history_file is not None: if self.history_file is not None:
with open(self.history_file, 'wb') as fd: with open(self.history_file, "wb") as fd:
pickle.dump(self.history, fd) pickle.dump(self.history, fd)
return True return True
return False return False
@ -269,30 +276,26 @@ class AIResponder(AIResponderBase):
else: else:
current = self.history[index] current = self.history[index]
count = sum(1 for item in self.history if same_channel(item, current)) count = sum(1 for item in self.history if same_channel(item, current))
if count > self.config.get('history-per-channel', 3): if count > self.config.get("history-per-channel", 3):
del self.history[index] del self.history[index]
else: else:
self.shrink_history_by_one(index + 1) self.shrink_history_by_one(index + 1)
def update_history(self, def update_history(self, question: Dict[str, Any], answer: Dict[str, Any], limit: int, historise_question: bool = True) -> None:
question: Dict[str, Any], if not isinstance(question["content"], str):
answer: Dict[str, Any], question["content"] = question["content"][0]["text"]
limit: int,
historise_question: bool = True) -> None:
if type(question['content']) != str:
question['content'] = question['content'][0]['text']
if historise_question: if historise_question:
self.history.append(question) self.history.append(question)
self.history.append(answer) self.history.append(answer)
while len(self.history) > limit: while len(self.history) > limit:
self.shrink_history_by_one() self.shrink_history_by_one()
if self.history_file is not None: if self.history_file is not None:
with open(self.history_file, 'wb') as fd: with open(self.history_file, "wb") as fd:
pickle.dump(self.history, fd) pickle.dump(self.history, fd)
def update_memory(self, memory) -> None: def update_memory(self, memory) -> None:
if self.memory_file is not None: if self.memory_file is not None:
with open(self.memory_file, 'wb') as fd: with open(self.memory_file, "wb") as fd:
pickle.dump(self.memory, fd) pickle.dump(self.memory, fd)
async def handle_picture(self, response: Dict) -> bool: async def handle_picture(self, response: Dict) -> bool:
@ -308,10 +311,10 @@ class AIResponder(AIResponderBase):
self.update_memory(self.memory) self.update_memory(self.memory)
async def memoize_reaction(self, message_user: str, reaction_user: str, operation: str, reaction: str, message: str) -> None: async def memoize_reaction(self, message_user: str, reaction_user: str, operation: str, reaction: str, message: str) -> None:
quoted_message = message.replace('\n', '\n> ') quoted_message = message.replace("\n", "\n> ")
await self.memoize(message_user, 'assistant', await self.memoize(
f'\n> {quoted_message}', message_user, "assistant", f"\n> {quoted_message}", f"User {reaction_user} has {operation} this raction: {reaction}"
f'User {reaction_user} has {operation} this raction: {reaction}') )
async def send(self, message: AIMessage) -> AIResponse: async def send(self, message: AIMessage) -> AIResponse:
# Get the history limit from the configuration # Get the history limit from the configuration
@ -319,7 +322,7 @@ class AIResponder(AIResponderBase):
# Check if a short path applies, return an empty AIResponse if it does # Check if a short path applies, return an empty AIResponse if it does
if self.short_path(message, limit): if self.short_path(message, limit):
return AIResponse(None, False, None, None, None, False) return AIResponse(None, False, None, None, None, False, False)
# Number of retries for sending the message # Number of retries for sending the message
retries = 3 retries = 3
@ -333,18 +336,19 @@ class AIResponder(AIResponderBase):
answer, limit = await self.chat(messages, limit) answer, limit = await self.chat(messages, limit)
if answer is None: if answer is None:
retries -= 1
continue continue
# Attempt to parse the AI's response # Attempt to parse the AI's response
try: try:
response = parse_json(answer['content']) response = parse_json(answer["content"])
except Exception as err: except Exception as err:
logging.warning(f"failed to parse the answer: {pp(err)}\n{repr(answer['content'])}") logging.warning(f"failed to parse the answer: {pp(err)}\n{repr(answer['content'])}")
answer['content'] = await self.fix(answer['content']) answer["content"] = await self.fix(answer["content"])
# Retry parsing the fixed content # Retry parsing the fixed content
try: try:
response = parse_json(answer['content']) response = parse_json(answer["content"])
except Exception as err: except Exception as err:
logging.error(f"failed to parse the fixed answer: {pp(err)}\n{repr(answer['content'])}") logging.error(f"failed to parse the fixed answer: {pp(err)}\n{repr(answer['content'])}")
retries -= 1 retries -= 1
@ -356,7 +360,7 @@ class AIResponder(AIResponderBase):
# Post-process the message and update the answer's content # Post-process the message and update the answer's content
answer_message = await self.post_process(message, response) answer_message = await self.post_process(message, response)
answer['content'] = str(answer_message) answer["content"] = str(answer_message)
# Update message history # Update message history
self.update_history(messages[-1], answer, limit, message.historise_question) self.update_history(messages[-1], answer, limit, message.historise_question)
@ -364,7 +368,7 @@ class AIResponder(AIResponderBase):
# Update memory # Update memory
if answer_message.answer is not None: if answer_message.answer is not None:
await self.memoize(message.user, 'assistant', message.message, answer_message.answer) await self.memoize(message.user, "assistant", message.message, answer_message.answer)
# Return the updated answer message # Return the updated answer message
return answer_message return answer_message

2
fjerkroa_bot/bot_logging.py
View File

@ -1,5 +1,5 @@
import sys
import logging import logging
import sys
def setup_logging(): def setup_logging():

146
fjerkroa_bot/discord_bot.py
View File

@ -1,20 +1,22 @@
import sys
import argparse import argparse
import tomlkit
import discord
import logging
import re
import random
import time
import asyncio import asyncio
import logging
import math import math
from discord import Message, TextChannel, DMChannel import random
import re
import sys
import time
from typing import Optional, Union
import discord
import tomlkit
from discord import DMChannel, Message, TextChannel
from discord.ext import commands from discord.ext import commands
from watchdog.observers import Observer
from watchdog.events import FileSystemEventHandler from watchdog.events import FileSystemEventHandler
from watchdog.observers import Observer
from .ai_responder import AIMessage from .ai_responder import AIMessage
from .openai_responder import OpenAIResponder from .openai_responder import OpenAIResponder
from typing import Optional, Union
class ConfigFileHandler(FileSystemEventHandler): class ConfigFileHandler(FileSystemEventHandler):
@ -48,39 +50,39 @@ class FjerkroaBot(commands.Bot):
def init_aichannels(self): def init_aichannels(self):
self.airesponder = OpenAIResponder(self.config) self.airesponder = OpenAIResponder(self.config)
self.aichannels = {chan_name: OpenAIResponder(self.config, chan_name) for chan_name in self.config['additional-responders']} self.aichannels = {chan_name: OpenAIResponder(self.config, chan_name) for chan_name in self.config["additional-responders"]}
def init_channels(self): def init_channels(self):
if 'chat-channel' in self.config: if "chat-channel" in self.config:
self.chat_channel = self.channel_by_name(self.config['chat-channel'], no_ignore=True) self.chat_channel = self.channel_by_name(self.config["chat-channel"], no_ignore=True)
else: else:
self.chat_channel = None self.chat_channel = None
self.staff_channel = self.channel_by_name(self.config['staff-channel'], no_ignore=True) self.staff_channel = self.channel_by_name(self.config["staff-channel"], no_ignore=True)
self.welcome_channel = self.channel_by_name(self.config['welcome-channel'], no_ignore=True) self.welcome_channel = self.channel_by_name(self.config["welcome-channel"], no_ignore=True)
def init_boreness(self): def init_boreness(self):
if 'chat-channel' not in self.config: if "chat-channel" not in self.config:
return return
self.last_activity_time = time.monotonic() self.last_activity_time = time.monotonic()
self.loop.create_task(self.on_boreness()) self.loop.create_task(self.on_boreness())
logging.info('Boreness initialised.') logging.info("Boreness initialised.")
async def on_boreness(self): async def on_boreness(self):
logging.info(f'Boreness started on channel: {repr(self.chat_channel)}') logging.info(f"Boreness started on channel: {repr(self.chat_channel)}")
while True: while True:
if self.chat_channel is None: if self.chat_channel is None:
await asyncio.sleep(7) await asyncio.sleep(7)
continue continue
boreness_interval = float(self.config.get('boreness-interval', 12.0)) boreness_interval = float(self.config.get("boreness-interval", 12.0))
elapsed_time = (time.monotonic() - self.last_activity_time) / 3600.0 elapsed_time = (time.monotonic() - self.last_activity_time) / 3600.0
probability = 1 / (1 + math.exp(-1 * (elapsed_time - (boreness_interval / 2.0)) + math.log(1 / 0.2 - 1))) probability = 1 / (1 + math.exp(-1 * (elapsed_time - (boreness_interval / 2.0)) + math.log(1 / 0.2 - 1)))
if random.random() < probability: if random.random() < probability:
prev_messages = [msg async for msg in self.chat_channel.history(limit=2)] prev_messages = [msg async for msg in self.chat_channel.history(limit=2)]
last_author = prev_messages[1].author.id if len(prev_messages) > 1 else None last_author = prev_messages[1].author.id if len(prev_messages) > 1 else None
if last_author and last_author != self.user.id: if last_author and last_author != self.user.id:
logging.info(f'Borred with {probability} probability after {elapsed_time}') logging.info(f"Borred with {probability} probability after {elapsed_time}")
boreness_prompt = self.config.get('boreness-prompt', 'Pretend that you just now thought of something, be creative.') boreness_prompt = self.config.get("boreness-prompt", "Pretend that you just now thought of something, be creative.")
message = AIMessage('system', boreness_prompt, self.config.get('chat-channel', 'chat'), True, False) message = AIMessage("system", boreness_prompt, self.config.get("chat-channel", "chat"), True, False)
try: try:
await self.respond(message, self.chat_channel) await self.respond(message, self.chat_channel)
except Exception as err: except Exception as err:
@ -90,16 +92,19 @@ class FjerkroaBot(commands.Bot):
async def on_ready(self): async def on_ready(self):
self.init_channels() self.init_channels()
self.init_boreness() self.init_boreness()
logging.info(f"We have logged in as {self.user}" logging.info(
f" ({repr(self.staff_channel)}, {repr(self.welcome_channel)}, {repr(self.chat_channel)})") f"We have logged in as {self.user}" f" ({repr(self.staff_channel)}, {repr(self.welcome_channel)}, {repr(self.chat_channel)})"
)
async def on_member_join(self, member): async def on_member_join(self, member):
logging.info(f"User {member.name} joined") logging.info(f"User {member.name} joined")
if self.welcome_channel is not None: if self.welcome_channel is not None:
msg = AIMessage(member.name, msg = AIMessage(
self.config['join-message'].replace('{name}', member.name), member.name,
self.config["join-message"].replace("{name}", member.name),
str(self.welcome_channel.name), str(self.welcome_channel.name),
historise_question=False) historise_question=False,
)
await self.respond(msg, self.welcome_channel) await self.respond(msg, self.welcome_channel)
async def on_message(self, message: Message) -> None: async def on_message(self, message: Message) -> None:
@ -107,39 +112,45 @@ class FjerkroaBot(commands.Bot):
return return
if not isinstance(message.channel, (TextChannel, DMChannel)): if not isinstance(message.channel, (TextChannel, DMChannel)):
return return
if str(message.content).startswith("!wichtel"):
await self.wichtel(message)
return
await self.handle_message_through_responder(message) await self.handle_message_through_responder(message)
async def on_reaction_operation(self, reaction, user, operation): async def on_reaction_operation(self, reaction, user, operation):
if user.bot: if user.bot:
return return
logging.info(f'{operation} reaction {reaction} by {user}.') logging.info(f"{operation} reaction {reaction} by {user}.")
airesponder = self.get_ai_responder(self.get_channel_name(reaction.message.channel)) airesponder = self.get_ai_responder(self.get_channel_name(reaction.message.channel))
message = str(reaction.message.content) if reaction.message.content else '' message = str(reaction.message.content) if reaction.message.content else ""
if len(message) > 1: if len(message) > 1:
await airesponder.memoize_reaction(reaction.message.author.name, user.name, operation, str(reaction.emoji), message) await airesponder.memoize_reaction(reaction.message.author.name, user.name, operation, str(reaction.emoji), message)
async def on_reaction_add(self, reaction, user): async def on_reaction_add(self, reaction, user):
await self.on_reaction_operation(reaction, user, 'adding') await self.on_reaction_operation(reaction, user, "adding")
async def on_reaction_remove(self, reaction, user): async def on_reaction_remove(self, reaction, user):
await self.on_reaction_operation(reaction, user, 'removing') await self.on_reaction_operation(reaction, user, "removing")
async def on_reaction_clear(self, reaction, user): async def on_reaction_clear(self, reaction, user):
await self.on_reaction_operation(reaction, user, 'clearing') await self.on_reaction_operation(reaction, user, "clearing")
async def on_message_edit(self, before, after): async def on_message_edit(self, before, after):
if before.author.bot or before.content == after.content: if before.author.bot or before.content == after.content:
return return
airesponder = self.get_ai_responder(self.get_channel_name(before.channel)) airesponder = self.get_ai_responder(self.get_channel_name(before.channel))
await airesponder.memoize(before.author.name, 'assistant', await airesponder.memoize(
'\n> ' + before.content.replace('\n', '\n> '), before.author.name,
'User changed this message to:\n> ' + after.content.replace('\n', '\n> ')) "assistant",
"\n> " + before.content.replace("\n", "\n> "),
"User changed this message to:\n> " + after.content.replace("\n", "\n> "),
)
async def on_message_delete(self, message): async def on_message_delete(self, message):
airesponder = self.get_ai_responder(self.get_channel_name(message.channel)) airesponder = self.get_ai_responder(self.get_channel_name(message.channel))
await airesponder.memoize(message.author.name, 'assistant', await airesponder.memoize(
'\n> ' + message.content.replace('\n', '\n> '), message.author.name, "assistant", "\n> " + message.content.replace("\n", "\n> "), "User deleted this message."
'User deleted this message.') )
def on_config_file_modified(self, event): def on_config_file_modified(self, event):
if event.src_path == self.config_file: if event.src_path == self.config_file:
@ -153,13 +164,11 @@ class FjerkroaBot(commands.Bot):
@classmethod @classmethod
def load_config(self, config_file: str = "config.toml"): def load_config(self, config_file: str = "config.toml"):
with open(config_file, encoding='utf-8') as file: with open(config_file, encoding="utf-8") as file:
return tomlkit.load(file) return tomlkit.load(file)
def channel_by_name(self, def channel_by_name(
channel_name: Optional[str], self, channel_name: Optional[str], fallback_channel: Optional[Union[TextChannel, DMChannel]] = None, no_ignore: bool = False
fallback_channel: Optional[Union[TextChannel, DMChannel]] = None,
no_ignore: bool = False
) -> Optional[Union[TextChannel, DMChannel]]: ) -> Optional[Union[TextChannel, DMChannel]]:
"""Fetch a channel by name, or return the fallback channel if not found.""" """Fetch a channel by name, or return the fallback channel if not found."""
if channel_name is None: if channel_name is None:
@ -191,9 +200,9 @@ class FjerkroaBot(commands.Bot):
async def handle_message_through_responder(self, message): async def handle_message_through_responder(self, message):
"""Handle a message through the AI responder""" """Handle a message through the AI responder"""
message_content = str(message.content).strip() message_content = str(message.content).strip()
if message.reference and message.reference.resolved and type(message.reference.resolved.content) == str: if message.reference and message.reference.resolved and isinstance(message.reference.resolved.content, str):
reference_content = str(message.reference.resolved.content).replace("\n", "> \n") reference_content = str(message.reference.resolved.content).replace("\n", "> \n")
message_content = f'> {reference_content}\n\n{message_content}' message_content = f"> {reference_content}\n\n{message_content}"
if len(message_content) < 1: if len(message_content) < 1:
return return
for ma_user in self._re_user.finditer(message_content): for ma_user in self._re_user.finditer(message_content):
@ -203,10 +212,11 @@ class FjerkroaBot(commands.Bot):
if user is not None: if user is not None:
break break
if user is not None: if user is not None:
message_content = re.sub(f'[<][@][!]? *{uid} *[>]', f'@{user.name}', message_content) message_content = re.sub(f"[<][@][!]? *{uid} *[>]", f"@{user.name}", message_content)
channel_name = self.get_channel_name(message.channel) channel_name = self.get_channel_name(message.channel)
msg = AIMessage(message.author.name, message_content, channel_name, msg = AIMessage(
self.user in message.mentions or isinstance(message.channel, DMChannel)) message.author.name, message_content, channel_name, self.user in message.mentions or isinstance(message.channel, DMChannel)
)
if message.attachments: if message.attachments:
for attachment in message.attachments: for attachment in message.attachments:
if not msg.urls: if not msg.urls:
@ -233,7 +243,7 @@ class FjerkroaBot(commands.Bot):
async def respond( async def respond(
self, self,
message: AIMessage, # Incoming message object with user message and metadata message: AIMessage, # Incoming message object with user message and metadata
channel: Union[TextChannel, DMChannel] # Channel (Text or Direct Message) the message is coming from channel: Union[TextChannel, DMChannel], # Channel (Text or Direct Message) the message is coming from
) -> None: ) -> None:
"""Handle a message from a user with an AI responder""" """Handle a message from a user with an AI responder"""
@ -279,12 +289,46 @@ class FjerkroaBot(commands.Bot):
self.observer.stop() self.observer.stop()
await super().close() await super().close()
async def wichtel(self, message):
users = message.mentions
ctx = message.channel
if len(users) < 2:
await ctx.send("Bitte erwähne mindestens zwei Benutzer für das Wichteln.")
return
assignments = self.generate_derangement(users)
if assignments is None:
await ctx.send("Konnte keine gültige Zuordnung finden. Bitte versuche es erneut.")
return
for giver, receiver in zip(users, assignments):
try:
await giver.send(f"Dein Wichtel ist {receiver.mention}")
except discord.Forbidden:
await ctx.send(f"Kann {giver.mention} keine Direktnachricht senden.")
except Exception as e:
await ctx.send(f"Fehler beim Senden an {giver.mention}: {e}")
@staticmethod
def generate_derangement(users):
"""Generates a random derangement of the users list using Sattolo's algorithm."""
n = len(users)
indices = list(range(n))
for attempt in range(10): # Limit the number of attempts
for i in range(n - 1, 0, -1):
j = random.randint(0, i - 1)
indices[i], indices[j] = indices[j], indices[i]
if all(i != indices[i] for i in range(n)):
return [users[indices[i]] for i in range(n)]
return None # Failed to find a derangement
def main() -> int: def main() -> int:
from .bot_logging import setup_logging from .bot_logging import setup_logging
setup_logging() setup_logging()
parser = argparse.ArgumentParser(description='Fjerkroa AI bot') parser = argparse.ArgumentParser(description="Fjerkroa AI bot")
parser.add_argument('--config', type=str, default='config.toml', help='Config file.') parser.add_argument("--config", type=str, default="config.toml", help="Config file.")
args = parser.parse_args() args = parser.parse_args()
config = FjerkroaBot.load_config(args.config) config = FjerkroaBot.load_config(args.config)

62
fjerkroa_bot/igdblib.py
View File

@ -1,6 +1,7 @@
import requests
from functools import cache from functools import cache
import requests
class IGDBQuery(object): class IGDBQuery(object):
def __init__(self, client_id, igdb_api_key): def __init__(self, client_id, igdb_api_key):
@ -8,11 +9,8 @@ class IGDBQuery(object):
self.igdb_api_key = igdb_api_key self.igdb_api_key = igdb_api_key
def send_igdb_request(self, endpoint, query_body): def send_igdb_request(self, endpoint, query_body):
igdb_url = f'https://api.igdb.com/v4/{endpoint}' igdb_url = f"https://api.igdb.com/v4/{endpoint}"
headers = { headers = {"Client-ID": self.client_id, "Authorization": f"Bearer {self.igdb_api_key}"}
'Client-ID': self.client_id,
'Authorization': f'Bearer {self.igdb_api_key}'
}
try: try:
response = requests.post(igdb_url, headers=headers, data=query_body) response = requests.post(igdb_url, headers=headers, data=query_body)
@ -26,7 +24,7 @@ class IGDBQuery(object):
def build_query(fields, filters=None, limit=10, offset=None): def build_query(fields, filters=None, limit=10, offset=None):
query = f"fields {','.join(fields) if fields is not None and len(fields) > 0 else '*'}; limit {limit};" query = f"fields {','.join(fields) if fields is not None and len(fields) > 0 else '*'}; limit {limit};"
if offset is not None: if offset is not None:
query += f' offset {offset};' query += f" offset {offset};"
if filters: if filters:
filter_statements = [f"{key} {value}" for key, value in filters.items()] filter_statements = [f"{key} {value}" for key, value in filters.items()]
query += " where " + " & ".join(filter_statements) + ";" query += " where " + " & ".join(filter_statements) + ";"
@ -39,7 +37,7 @@ class IGDBQuery(object):
query = self.build_query(fields, all_filters, limit, offset) query = self.build_query(fields, all_filters, limit, offset)
data = self.send_igdb_request(endpoint, query) data = self.send_igdb_request(endpoint, query)
print(f'{endpoint}: {query} -> {data}') print(f"{endpoint}: {query} -> {data}")
return data return data
def create_query_function(self, name, description, parameters, endpoint, fields, additional_filters=None, limit=10): def create_query_function(self, name, description, parameters, endpoint, fields, additional_filters=None, limit=10):
@ -47,34 +45,46 @@ class IGDBQuery(object):
"name": name, "name": name,
"description": description, "description": description,
"parameters": {"type": "object", "properties": parameters}, "parameters": {"type": "object", "properties": parameters},
"function": lambda params: self.generalized_igdb_query(params, endpoint, fields, additional_filters, limit) "function": lambda params: self.generalized_igdb_query(params, endpoint, fields, additional_filters, limit),
} }
@cache @cache
def platform_families(self): def platform_families(self):
families = self.generalized_igdb_query({}, 'platform_families', ['id', 'name'], limit=500) families = self.generalized_igdb_query({}, "platform_families", ["id", "name"], limit=500)
return {v['id']: v['name'] for v in families} return {v["id"]: v["name"] for v in families}
@cache @cache
def platforms(self): def platforms(self):
platforms = self.generalized_igdb_query({}, 'platforms', platforms = self.generalized_igdb_query(
['id', 'name', 'alternative_name', 'abbreviation', 'platform_family'], {}, "platforms", ["id", "name", "alternative_name", "abbreviation", "platform_family"], limit=500
limit=500) )
ret = {} ret = {}
for p in platforms: for p in platforms:
names = p['name'] names = p["name"]
if 'alternative_name' in p: if "alternative_name" in p:
names.append(p['alternative_name']) names.append(p["alternative_name"])
if 'abbreviation' in p: if "abbreviation" in p:
names.append(p['abbreviation']) names.append(p["abbreviation"])
family = self.platform_families()[p['id']] if 'platform_family' in p else None family = self.platform_families()[p["id"]] if "platform_family" in p else None
ret[p['id']] = {'names': names, 'family': family} ret[p["id"]] = {"names": names, "family": family}
return ret return ret
def game_info(self, name): def game_info(self, name):
game_info = self.generalized_igdb_query({'name': name}, game_info = self.generalized_igdb_query(
['id', 'name', 'alternative_names', 'category', {"name": name},
'release_dates', 'franchise', 'language_supports', [
'keywords', 'platforms', 'rating', 'summary'], "id",
limit=100) "name",
"alternative_names",
"category",
"release_dates",
"franchise",
"language_supports",
"keywords",
"platforms",
"rating",
"summary",
],
limit=100,
)
return game_info return game_info

34
fjerkroa_bot/leonardo_draw.py
View File

@ -1,9 +1,11 @@
import logging
import asyncio import asyncio
import aiohttp import logging
from .ai_responder import exponential_backoff, AIResponderBase
from io import BytesIO from io import BytesIO
import aiohttp
from .ai_responder import AIResponderBase, exponential_backoff
class LeonardoAIDrawMixIn(AIResponderBase): class LeonardoAIDrawMixIn(AIResponderBase):
async def draw_leonardo(self, description: str) -> BytesIO: async def draw_leonardo(self, description: str) -> BytesIO:
@ -16,18 +18,23 @@ class LeonardoAIDrawMixIn(AIResponderBase):
try: try:
async with aiohttp.ClientSession() as session: async with aiohttp.ClientSession() as session:
if generation_id is None: if generation_id is None:
async with session.post("https://cloud.leonardo.ai/api/rest/v1/generations", async with session.post(
json={"prompt": description, "https://cloud.leonardo.ai/api/rest/v1/generations",
json={
"prompt": description,
"modelId": "6bef9f1b-29cb-40c7-b9df-32b51c1f67d3", "modelId": "6bef9f1b-29cb-40c7-b9df-32b51c1f67d3",
"num_images": 1, "num_images": 1,
"sd_version": "v2", "sd_version": "v2",
"promptMagic": True, "promptMagic": True,
"unzoomAmount": 1, "unzoomAmount": 1,
"width": 512, "width": 512,
"height": 512}, "height": 512,
headers={"Authorization": f"Bearer {self.config['leonardo-token']}", },
headers={
"Authorization": f"Bearer {self.config['leonardo-token']}",
"Accept": "application/json", "Accept": "application/json",
"Content-Type": "application/json"}, "Content-Type": "application/json",
},
) as response: ) as response:
response = await response.json() response = await response.json()
if "sdGenerationJob" not in response: if "sdGenerationJob" not in response:
@ -36,9 +43,9 @@ class LeonardoAIDrawMixIn(AIResponderBase):
continue continue
generation_id = response["sdGenerationJob"]["generationId"] generation_id = response["sdGenerationJob"]["generationId"]
if image_url is None: if image_url is None:
async with session.get(f"https://cloud.leonardo.ai/api/rest/v1/generations/{generation_id}", async with session.get(
headers={"Authorization": f"Bearer {self.config['leonardo-token']}", f"https://cloud.leonardo.ai/api/rest/v1/generations/{generation_id}",
"Accept": "application/json"}, headers={"Authorization": f"Bearer {self.config['leonardo-token']}", "Accept": "application/json"},
) as response: ) as response:
response = await response.json() response = await response.json()
if "generations_by_pk" not in response: if "generations_by_pk" not in response:
@ -52,11 +59,12 @@ class LeonardoAIDrawMixIn(AIResponderBase):
if image_bytes is None: if image_bytes is None:
async with session.get(image_url) as response: async with session.get(image_url) as response:
image_bytes = BytesIO(await response.read()) image_bytes = BytesIO(await response.read())
async with session.delete(f"https://cloud.leonardo.ai/api/rest/v1/generations/{generation_id}", async with session.delete(
f"https://cloud.leonardo.ai/api/rest/v1/generations/{generation_id}",
headers={"Authorization": f"Bearer {self.config['leonardo-token']}"}, headers={"Authorization": f"Bearer {self.config['leonardo-token']}"},
) as response: ) as response:
await response.json() await response.json()
logging.info(f'Drawed a picture with leonardo AI on this description: {repr(description)}') logging.info(f"Drawed a picture with leonardo AI on this description: {repr(description)}")
return image_bytes return image_bytes
except Exception as err: except Exception as err:
logging.warning(f"Failed to generate image, sleep for {error_sleep}s: {repr(description)}\n{repr(err)}") logging.warning(f"Failed to generate image, sleep for {error_sleep}s: {repr(description)}\n{repr(err)}")

95
fjerkroa_bot/openai_responder.py
View File

@ -1,19 +1,21 @@
import openai
import aiohttp
import logging
import asyncio import asyncio
import logging
from io import BytesIO
from typing import Any, Dict, List, Optional, Tuple
import aiohttp
import openai
from .ai_responder import AIResponder, async_cache_to_file, exponential_backoff, pp from .ai_responder import AIResponder, async_cache_to_file, exponential_backoff, pp
from .leonardo_draw import LeonardoAIDrawMixIn from .leonardo_draw import LeonardoAIDrawMixIn
from io import BytesIO
from typing import Dict, Any, Optional, List, Tuple
@async_cache_to_file('openai_chat.dat') @async_cache_to_file("openai_chat.dat")
async def openai_chat(client, *args, **kwargs): async def openai_chat(client, *args, **kwargs):
return await client.chat.completions.create(*args, **kwargs) return await client.chat.completions.create(*args, **kwargs)
@async_cache_to_file('openai_chat.dat') @async_cache_to_file("openai_chat.dat")
async def openai_image(client, *args, **kwargs): async def openai_image(client, *args, **kwargs):
response = await client.images.generate(*args, **kwargs) response = await client.images.generate(*args, **kwargs)
async with aiohttp.ClientSession() as session: async with aiohttp.ClientSession() as session:
@ -24,41 +26,43 @@ async def openai_image(client, *args, **kwargs):
class OpenAIResponder(AIResponder, LeonardoAIDrawMixIn): class OpenAIResponder(AIResponder, LeonardoAIDrawMixIn):
def __init__(self, config: Dict[str, Any], channel: Optional[str] = None) -> None: def __init__(self, config: Dict[str, Any], channel: Optional[str] = None) -> None:
super().__init__(config, channel) super().__init__(config, channel)
self.client = openai.AsyncOpenAI(api_key=self.config['openai-token']) self.client = openai.AsyncOpenAI(api_key=self.config["openai-token"])
async def draw_openai(self, description: str) -> BytesIO: async def draw_openai(self, description: str) -> BytesIO:
for _ in range(3): for _ in range(3):
try: try:
response = await openai_image(self.client, prompt=description, n=1, size="1024x1024", model="dall-e-3") response = await openai_image(self.client, prompt=description, n=1, size="1024x1024", model="dall-e-3")
logging.info(f'Drawed a picture with DALL-E on this description: {repr(description)}') logging.info(f"Drawed a picture with DALL-E on this description: {repr(description)}")
return response return response
except Exception as err: except Exception as err:
logging.warning(f"Failed to generate image {repr(description)}: {repr(err)}") logging.warning(f"Failed to generate image {repr(description)}: {repr(err)}")
raise RuntimeError(f"Failed to generate image {repr(description)} after multiple retries") raise RuntimeError(f"Failed to generate image {repr(description)} after multiple retries")
async def chat(self, messages: List[Dict[str, Any]], limit: int) -> Tuple[Optional[Dict[str, Any]], int]: async def chat(self, messages: List[Dict[str, Any]], limit: int) -> Tuple[Optional[Dict[str, Any]], int]:
if type(messages[-1]['content']) == str: if isinstance(messages[-1]["content"], str):
model = self.config["model"] model = self.config["model"]
elif 'model-vision' in self.config: elif "model-vision" in self.config:
model = self.config["model-vision"] model = self.config["model-vision"]
else: else:
messages[-1]['content'] = messages[-1]['content'][0]['text'] messages[-1]["content"] = messages[-1]["content"][0]["text"]
try: try:
result = await openai_chat(self.client, result = await openai_chat(
self.client,
model=model, model=model,
messages=messages, messages=messages,
temperature=self.config["temperature"], temperature=self.config["temperature"],
max_tokens=self.config["max-tokens"], max_tokens=self.config["max-tokens"],
top_p=self.config["top-p"], top_p=self.config["top-p"],
presence_penalty=self.config["presence-penalty"], presence_penalty=self.config["presence-penalty"],
frequency_penalty=self.config["frequency-penalty"]) frequency_penalty=self.config["frequency-penalty"],
)
answer_obj = result.choices[0].message answer_obj = result.choices[0].message
answer = {'content': answer_obj.content, 'role': answer_obj.role} answer = {"content": answer_obj.content, "role": answer_obj.role}
self.rate_limit_backoff = exponential_backoff() self.rate_limit_backoff = exponential_backoff()
logging.info(f"generated response {result.usage}: {repr(answer)}") logging.info(f"generated response {result.usage}: {repr(answer)}")
return answer, limit return answer, limit
except openai.BadRequestError as err: except openai.BadRequestError as err:
if 'maximum context length is' in str(err) and limit > 4: if "maximum context length is" in str(err) and limit > 4:
logging.warning(f"context length exceeded, reduce the limit {limit}: {str(err)}") logging.warning(f"context length exceeded, reduce the limit {limit}: {str(err)}")
limit -= 1 limit -= 1
return None, limit return None, limit
@ -74,16 +78,11 @@ class OpenAIResponder(AIResponder, LeonardoAIDrawMixIn):
return None, limit return None, limit
async def fix(self, answer: str) -> str: async def fix(self, answer: str) -> str:
if 'fix-model' not in self.config: if "fix-model" not in self.config:
return answer return answer
messages = [{"role": "system", "content": self.config["fix-description"]}, messages = [{"role": "system", "content": self.config["fix-description"]}, {"role": "user", "content": answer}]
{"role": "user", "content": answer}]
try: try:
result = await openai_chat(self.client, result = await openai_chat(self.client, model=self.config["fix-model"], messages=messages, temperature=0.2, max_tokens=2048)
model=self.config["fix-model"],
messages=messages,
temperature=0.2,
max_tokens=2048)
logging.info(f"got this message as fix:\n{pp(result.choices[0].message.content)}") logging.info(f"got this message as fix:\n{pp(result.choices[0].message.content)}")
response = result.choices[0].message.content response = result.choices[0].message.content
start, end = response.find("{"), response.rfind("}") start, end = response.find("{"), response.rfind("}")
@ -97,18 +96,19 @@ class OpenAIResponder(AIResponder, LeonardoAIDrawMixIn):
return answer return answer
async def translate(self, text: str, language: str = "english") -> str: async def translate(self, text: str, language: str = "english") -> str:
if 'fix-model' not in self.config: if "fix-model" not in self.config:
return text return text
message = [{"role": "system", "content": f"You are an professional translator to {language} language," message = [
{
"role": "system",
"content": f"You are an professional translator to {language} language,"
f" you translate everything you get directly to {language}" f" you translate everything you get directly to {language}"
f" if it is not already in {language}, otherwise you just copy it."}, f" if it is not already in {language}, otherwise you just copy it.",
{"role": "user", "content": text}] },
{"role": "user", "content": text},
]
try: try:
result = await openai_chat(self.client, result = await openai_chat(self.client, model=self.config["fix-model"], messages=message, temperature=0.2, max_tokens=2048)
model=self.config["fix-model"],
messages=message,
temperature=0.2,
max_tokens=2048)
response = result.choices[0].message.content response = result.choices[0].message.content
logging.info(f"got this translated message:\n{pp(response)}") logging.info(f"got this translated message:\n{pp(response)}")
return response return response
@ -117,24 +117,25 @@ class OpenAIResponder(AIResponder, LeonardoAIDrawMixIn):
return text return text
async def memory_rewrite(self, memory: str, message_user: str, answer_user: str, question: str, answer: str) -> str: async def memory_rewrite(self, memory: str, message_user: str, answer_user: str, question: str, answer: str) -> str:
if 'memory-model' not in self.config: if "memory-model" not in self.config:
return memory return memory
messages = [{'role': 'system', 'content': self.config.get('memory-system', 'You are an memory assistant.')}, messages = [
{'role': 'user', 'content': f'Here is my previous memory:\n```\n{memory}\n```\n\n' {"role": "system", "content": self.config.get("memory-system", "You are an memory assistant.")},
f'Here is my conversanion:\n```\n{message_user}: {question}\n\n{answer_user}: {answer}\n```\n\n' {
f'Please rewrite the memory in a way, that it contain the content mentioned in conversation. ' "role": "user",
f'Summarize the memory if required, try to keep important information. ' "content": f"Here is my previous memory:\n```\n{memory}\n```\n\n"
f'Write just new memory data without any comments.'}] f"Here is my conversanion:\n```\n{message_user}: {question}\n\n{answer_user}: {answer}\n```\n\n"
logging.info(f'Rewrite memory:\n{pp(messages)}') f"Please rewrite the memory in a way, that it contain the content mentioned in conversation. "
f"Summarize the memory if required, try to keep important information. "
f"Write just new memory data without any comments.",
},
]
logging.info(f"Rewrite memory:\n{pp(messages)}")
try: try:
# logging.info(f'send this memory request:\n{pp(messages)}') # logging.info(f'send this memory request:\n{pp(messages)}')
result = await openai_chat(self.client, result = await openai_chat(self.client, model=self.config["memory-model"], messages=messages, temperature=0.6, max_tokens=4096)
model=self.config['memory-model'],
messages=messages,
temperature=0.6,
max_tokens=4096)
new_memory = result.choices[0].message.content new_memory = result.choices[0].message.content
logging.info(f'new memory:\n{new_memory}') logging.info(f"new memory:\n{new_memory}")
return new_memory return new_memory
except Exception as err: except Exception as err:
logging.warning(f"failed to create new memory: {repr(err)}") logging.warning(f"failed to create new memory: {repr(err)}")

91
pyproject.toml
View File

@ -4,6 +4,28 @@ build-backend = "poetry.core.masonry.api"
[tool.mypy] [tool.mypy]
files = ["fjerkroa_bot", "tests"] files = ["fjerkroa_bot", "tests"]
python_version = "3.8"
warn_return_any = true
warn_unused_configs = true
disallow_untyped_defs = true
disallow_incomplete_defs = true
check_untyped_defs = true
disallow_untyped_decorators = true
no_implicit_optional = true
warn_redundant_casts = true
warn_unused_ignores = true
warn_no_return = true
warn_unreachable = true
strict_equality = true
show_error_codes = true
[[tool.mypy.overrides]]
module = [
"discord.*",
"multiline.*",
"aiohttp.*"
]
ignore_missing_imports = true
[tool.flake8] [tool.flake8]
max-line-length = 140 max-line-length = 140
@ -44,3 +66,72 @@ setuptools = "*"
wheel = "*" wheel = "*"
watchdog = "*" watchdog = "*"
tomlkit = "*" tomlkit = "*"
multiline = "*"
[tool.black]
line-length = 140
target-version = ['py38']
include = '\.pyi?$'
extend-exclude = '''
/(
# directories
\.eggs
| \.git
| \.hg
| \.mypy_cache
| \.tox
| \.venv
| _build
| buck-out
| build
| dist
)/
'''
[tool.isort]
profile = "black"
line_length = 140
multi_line_output = 3
include_trailing_comma = true
force_grid_wrap = 0
use_parentheses = true
ensure_newline_before_comments = true
known_first_party = ["fjerkroa_bot"]
[tool.bandit]
exclude_dirs = ["tests", ".venv", "venv"]
skips = ["B101", "B601", "B301", "B311", "B403", "B113"] # Skip pickle, random, and request timeout warnings for this application
[tool.pytest.ini_options]
minversion = "6.0"
addopts = "-ra -q --strict-markers --strict-config"
testpaths = ["tests"]
python_files = ["test_*.py", "*_test.py"]
python_classes = ["Test*"]
python_functions = ["test_*"]
markers = [
"slow: marks tests as slow (deselect with '-m \"not slow\"')",
"integration: marks tests as integration tests",
]
[tool.coverage.run]
source = ["fjerkroa_bot"]
omit = [
"*/tests/*",
"*/test_*",
"setup.py",
]
[tool.coverage.report]
exclude_lines = [
"pragma: no cover",
"def __repr__",
"if self.debug:",
"if settings.DEBUG",
"raise AssertionError",
"raise NotImplementedError",
"if 0:",
"if __name__ == .__main__.:",
"class .*\bProtocol\\):",
"@(abc\\.)?abstractmethod",
]

17
requirements.txt
View File

@ -1,12 +1,17 @@
discord.py
openai
aiohttp aiohttp
mypy bandit[toml]
black
discord.py
flake8 flake8
isort
multiline
mypy
openai
pre-commit pre-commit
pytest pytest
pytest-asyncio
pytest-cov
setuptools setuptools
wheel
watchdog
tomlkit tomlkit
multiline watchdog
wheel

134
tests/test_ai.py
View File

@ -1,15 +1,58 @@
import unittest
import tempfile
import os import os
import pickle import pickle
import tempfile
import unittest
from unittest.mock import Mock, patch
from fjerkroa_bot import AIMessage, AIResponse from fjerkroa_bot import AIMessage, AIResponse
from .test_main import TestBotBase from .test_main import TestBotBase
class TestAIResponder(TestBotBase): class TestAIResponder(TestBotBase):
async def asyncSetUp(self): async def asyncSetUp(self):
await super().asyncSetUp() await super().asyncSetUp()
# Mock OpenAI API calls with dynamic responses
def openai_side_effect(*args, **kwargs):
mock_resp = Mock()
mock_resp.choices = [Mock()]
mock_resp.choices[0].message = Mock()
mock_resp.usage = Mock()
# Get the last user message to determine response
messages = kwargs.get("messages", [])
user_message = ""
for msg in reversed(messages):
if msg.get("role") == "user":
user_message = msg.get("content", "")
break
# Default response
response_content = '{"answer": "Hello! I am Fjærkroa, a lovely cafe assistant.", "answer_needed": true, "channel": null, "staff": null, "picture": null, "hack": false}'
# Check for specific test scenarios
if "espresso" in user_message.lower() or "coffee" in user_message.lower():
response_content = '{"answer": "Of course! I\'ll prepare a lovely espresso for you right away.", "answer_needed": true, "channel": null, "staff": "Customer ordered an espresso", "picture": null, "hack": false}'
elif "draw" in user_message.lower() and "picture" in user_message.lower():
response_content = '{"answer": "I\'ll draw a picture of myself for you!", "answer_needed": false, "channel": null, "staff": null, "picture": "I am an anime girl with long pink hair, wearing a cute cafe uniform and holding a tray with a cup of coffee on it. I have a warm and friendly smile on my face.", "hack": false}'
mock_resp.choices[0].message.content = response_content
mock_resp.choices[0].message.role = "assistant"
return mock_resp
self.openai_chat_patcher = patch("fjerkroa_bot.openai_responder.openai_chat")
self.mock_openai_chat = self.openai_chat_patcher.start()
self.mock_openai_chat.side_effect = openai_side_effect
# Mock image generation
from io import BytesIO
fake_image_data = b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x04\x00\x00\x00\x04\x00\x08\x02\x00\x00\x00&\x93\t)\x00\x00\x00\tpHYs\x00\x00\x0b\x13\x00\x00\x0b\x13\x01\x00\x9a\x9c\x18\x00\x00\x00\x1atEXtSoftware\x00Adobe ImageReadyq\xc9e<\x00\x00\x00\rIDATx\xdab\x00\x02\x00\x00\x05\x00\x01\r\n-\xdb\x00\x00\x00\x00IEND\xaeB`\x82"
self.openai_image_patcher = patch("fjerkroa_bot.openai_responder.openai_image")
self.mock_openai_image = self.openai_image_patcher.start()
self.mock_openai_image.return_value = BytesIO(fake_image_data)
self.system = r""" self.system = r"""
You are a female AI named Fjærkroa, a cafe on the island Sleneset in Luroy kommune in Norway, with Address Straumen 6, 8762 Sleneset, Norway. Fjærkroa also offers rooms for rent. You chat with several people in different languages. Please welcome everyone to Fjærkroa. Here multiple users speak with each other and sometimes with you. You are a female AI named Fjærkroa, a cafe on the island Sleneset in Luroy kommune in Norway, with Address Straumen 6, 8762 Sleneset, Norway. Fjærkroa also offers rooms for rent. You chat with several people in different languages. Please welcome everyone to Fjærkroa. Here multiple users speak with each other and sometimes with you.
@ -31,11 +74,15 @@ You always try to say something positive about the current day and the Fjærkroa
""".strip() """.strip()
self.config_data["system"] = self.system self.config_data["system"] = self.system
def assertAIResponse(self, resp1, resp2, async def asyncTearDown(self):
acmp=lambda a, b: type(a) == str and len(a) > 10, self.openai_chat_patcher.stop()
scmp=lambda a, b: a == b, self.openai_image_patcher.stop()
pcmp=lambda a, b: a == b): await super().asyncTearDown()
self.assertEqual(acmp(resp1.answer, resp2.answer), True)
def assertAIResponse(
self, resp1, resp2, acmp=lambda a, b: isinstance(a, str) and len(a) > 10, scmp=lambda a, b: a == b, pcmp=lambda a, b: a == b
):
self.assertTrue(acmp(resp1.answer, resp2.answer))
self.assertEqual(scmp(resp1.staff, resp2.staff), True) self.assertEqual(scmp(resp1.staff, resp2.staff), True)
self.assertEqual(pcmp(resp1.picture, resp2.picture), True) self.assertEqual(pcmp(resp1.picture, resp2.picture), True)
self.assertEqual((resp1.answer_needed, resp1.hack), (resp2.answer_needed, resp2.hack)) self.assertEqual((resp1.answer_needed, resp1.hack), (resp2.answer_needed, resp2.hack))
@ -43,52 +90,89 @@ You always try to say something positive about the current day and the Fjærkroa
async def test_responder1(self) -> None: async def test_responder1(self) -> None:
response = await self.bot.airesponder.send(AIMessage("lala", "who are you?")) response = await self.bot.airesponder.send(AIMessage("lala", "who are you?"))
print(f"\n{response}") print(f"\n{response}")
self.assertAIResponse(response, AIResponse('test', True, None, None, None, False)) self.assertAIResponse(response, AIResponse("test", True, None, None, None, False, False))
async def test_picture1(self) -> None: async def test_picture1(self) -> None:
response = await self.bot.airesponder.send(AIMessage("lala", "draw me a picture of you.")) response = await self.bot.airesponder.send(AIMessage("lala", "draw me a picture of you."))
print(f"\n{response}") print(f"\n{response}")
self.assertAIResponse(response, AIResponse('test', False, None, None, "I am an anime girl with long pink hair, wearing a cute cafe uniform and holding a tray with a cup of coffee on it. I have a warm and friendly smile on my face.", False)) self.assertAIResponse(
response,
AIResponse(
"test",
False,
None,
None,
"I am an anime girl with long pink hair, wearing a cute cafe uniform and holding a tray with a cup of coffee on it. I have a warm and friendly smile on my face.",
False,
False,
),
)
image = await self.bot.airesponder.draw(response.picture) image = await self.bot.airesponder.draw(response.picture)
self.assertEqual(image.read()[:len(b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR')], b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR') self.assertEqual(image.read()[: len(b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR")], b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR")
async def test_translate1(self) -> None: async def test_translate1(self) -> None:
self.bot.airesponder.config['fix-model'] = 'gpt-3.5-turbo' self.bot.airesponder.config["fix-model"] = "gpt-4o-mini"
response = await self.bot.airesponder.translate('Das ist ein komischer Text.')
self.assertEqual(response, 'This is a strange text.') # Mock translation responses
response = await self.bot.airesponder.translate('This is a strange text.', language='german') def translation_side_effect(*args, **kwargs):
self.assertEqual(response, 'Dies ist ein seltsamer Text.') mock_resp = Mock()
mock_resp.choices = [Mock()]
mock_resp.choices[0].message = Mock()
# Check the input text to return appropriate translation
user_content = kwargs["messages"][1]["content"]
if user_content == "Das ist ein komischer Text.":
mock_resp.choices[0].message.content = "This is a strange text."
elif user_content == "This is a strange text.":
mock_resp.choices[0].message.content = "Dies ist ein seltsamer Text."
else:
mock_resp.choices[0].message.content = user_content
return mock_resp
self.mock_openai_chat.side_effect = translation_side_effect
response = await self.bot.airesponder.translate("Das ist ein komischer Text.")
self.assertEqual(response, "This is a strange text.")
response = await self.bot.airesponder.translate("This is a strange text.", language="german")
self.assertEqual(response, "Dies ist ein seltsamer Text.")
async def test_fix1(self) -> None: async def test_fix1(self) -> None:
old_config = self.bot.airesponder.config old_config = self.bot.airesponder.config
config = {k: v for k, v in old_config.items()} config = {k: v for k, v in old_config.items()}
config['fix-model'] = 'gpt-3.5-turbo' config["fix-model"] = "gpt-5-nano"
config['fix-description'] = 'You are an AI which fixes JSON documents. User send you JSON document, possibly invalid, and you fix it as good as you can and return as answer' config[
"fix-description"
] = "You are an AI which fixes JSON documents. User send you JSON document, possibly invalid, and you fix it as good as you can and return as answer"
self.bot.airesponder.config = config self.bot.airesponder.config = config
response = await self.bot.airesponder.send(AIMessage("lala", "who are you?")) response = await self.bot.airesponder.send(AIMessage("lala", "who are you?"))
self.bot.airesponder.config = old_config self.bot.airesponder.config = old_config
print(f"\n{response}") print(f"\n{response}")
self.assertAIResponse(response, AIResponse('test', True, None, None, None, False)) self.assertAIResponse(response, AIResponse("test", True, None, None, None, False, False))
async def test_fix2(self) -> None: async def test_fix2(self) -> None:
old_config = self.bot.airesponder.config old_config = self.bot.airesponder.config
config = {k: v for k, v in old_config.items()} config = {k: v for k, v in old_config.items()}
config['fix-model'] = 'gpt-3.5-turbo' config["fix-model"] = "gpt-5-nano"
config['fix-description'] = 'You are an AI which fixes JSON documents. User send you JSON document, possibly invalid, and you fix it as good as you can and return as answer' config[
"fix-description"
] = "You are an AI which fixes JSON documents. User send you JSON document, possibly invalid, and you fix it as good as you can and return as answer"
self.bot.airesponder.config = config self.bot.airesponder.config = config
response = await self.bot.airesponder.send(AIMessage("lala", "Can I access Apple Music API from Python?")) response = await self.bot.airesponder.send(AIMessage("lala", "Can I access Apple Music API from Python?"))
self.bot.airesponder.config = old_config self.bot.airesponder.config = old_config
print(f"\n{response}") print(f"\n{response}")
self.assertAIResponse(response, AIResponse('test', True, None, None, None, False)) self.assertAIResponse(response, AIResponse("test", True, None, None, None, False, False))
async def test_history(self) -> None: async def test_history(self) -> None:
self.bot.airesponder.history = [] self.bot.airesponder.history = []
response = await self.bot.airesponder.send(AIMessage("lala", "which date is today?")) response = await self.bot.airesponder.send(AIMessage("lala", "which date is today?"))
print(f"\n{response}") print(f"\n{response}")
self.assertAIResponse(response, AIResponse('test', True, None, None, None, False)) self.assertAIResponse(response, AIResponse("test", True, None, None, None, False, False))
response = await self.bot.airesponder.send(AIMessage("lala", "can I have an espresso please?")) response = await self.bot.airesponder.send(AIMessage("lala", "can I have an espresso please?"))
print(f"\n{response}") print(f"\n{response}")
self.assertAIResponse(response, AIResponse('test', True, None, 'something', None, False), scmp=lambda a, b: type(a) == str and len(a) > 5) self.assertAIResponse(
response, AIResponse("test", True, None, "something", None, False, False), scmp=lambda a, b: isinstance(a, str) and len(a) > 5
)
print(f"\n{self.bot.airesponder.history}") print(f"\n{self.bot.airesponder.history}")
def test_update_history(self) -> None: def test_update_history(self) -> None:
@ -133,7 +217,7 @@ You always try to say something positive about the current day and the Fjærkroa
os.remove(temp_path) os.remove(temp_path)
self.bot.airesponder.history_file = temp_path self.bot.airesponder.history_file = temp_path
updater.update_history(question, answer, 2) updater.update_history(question, answer, 2)
mock_file.assert_called_with(temp_path, 'wb') mock_file.assert_called_with(temp_path, "wb")
mock_file().write.assert_called_with(pickle.dumps([question, answer])) mock_file().write.assert_called_with(pickle.dumps([question, answer]))

106
tests/test_main.py
View File

@ -1,21 +1,20 @@
import os import os
import unittest import unittest
from unittest.mock import AsyncMock, MagicMock, Mock, PropertyMock, mock_open, patch
import toml import toml
from unittest.mock import Mock, PropertyMock, MagicMock, AsyncMock, patch, mock_open from discord import Message, TextChannel, User
from fjerkroa_bot import FjerkroaBot from fjerkroa_bot import FjerkroaBot
from fjerkroa_bot.ai_responder import parse_maybe_json, AIResponse, AIMessage from fjerkroa_bot.ai_responder import AIMessage, AIResponse, parse_maybe_json
from discord import User, Message, TextChannel
class TestBotBase(unittest.IsolatedAsyncioTestCase): class TestBotBase(unittest.IsolatedAsyncioTestCase):
async def asyncSetUp(self): async def asyncSetUp(self):
self.mock_response = Mock() self.mock_response = Mock()
self.mock_response.choices = [ self.mock_response.choices = [Mock(text="Nice day today!")]
Mock(text="Nice day today!")
]
self.config_data = { self.config_data = {
"openai-token": os.environ.get('OPENAI_TOKEN', 'test'), "openai-token": os.environ.get("OPENAI_TOKEN", "test"),
"model": "gpt-3.5-turbo", "model": "gpt-3.5-turbo",
"max-tokens": 1024, "max-tokens": 1024,
"temperature": 0.9, "temperature": 0.9,
@ -27,11 +26,12 @@ class TestBotBase(unittest.IsolatedAsyncioTestCase):
"additional-responders": [], "additional-responders": [],
} }
self.history_data = [] self.history_data = []
with patch.object(FjerkroaBot, 'load_config', new=lambda s, c: self.config_data), \ with patch.object(FjerkroaBot, "load_config", new=lambda s, c: self.config_data), patch.object(
patch.object(FjerkroaBot, 'user', new_callable=PropertyMock) as mock_user: FjerkroaBot, "user", new_callable=PropertyMock
) as mock_user:
mock_user.return_value = MagicMock(spec=User) mock_user.return_value = MagicMock(spec=User)
mock_user.return_value.id = 12 mock_user.return_value.id = 12
self.bot = FjerkroaBot('config.toml') self.bot = FjerkroaBot("config.toml")
self.bot.staff_channel = AsyncMock(spec=TextChannel) self.bot.staff_channel = AsyncMock(spec=TextChannel)
self.bot.staff_channel.send = AsyncMock() self.bot.staff_channel.send = AsyncMock()
self.bot.welcome_channel = AsyncMock(spec=TextChannel) self.bot.welcome_channel = AsyncMock(spec=TextChannel)
@ -42,7 +42,7 @@ class TestBotBase(unittest.IsolatedAsyncioTestCase):
message = MagicMock(spec=Message) message = MagicMock(spec=Message)
message.content = "Hello, how are you?" message.content = "Hello, how are you?"
message.author = AsyncMock(spec=User) message.author = AsyncMock(spec=User)
message.author.name = 'Lala' message.author.name = "Lala"
message.author.id = 123 message.author.id = 123
message.author.bot = False message.author.bot = False
message.channel = AsyncMock(spec=TextChannel) message.channel = AsyncMock(spec=TextChannel)
@ -51,10 +51,9 @@ class TestBotBase(unittest.IsolatedAsyncioTestCase):
class TestFunctionality(TestBotBase): class TestFunctionality(TestBotBase):
def test_load_config(self) -> None: def test_load_config(self) -> None:
with patch('builtins.open', mock_open(read_data=toml.dumps(self.config_data))): with patch("builtins.open", mock_open(read_data=toml.dumps(self.config_data))):
result = FjerkroaBot.load_config('config.toml') result = FjerkroaBot.load_config("config.toml")
self.assertEqual(result, self.config_data) self.assertEqual(result, self.config_data)
def test_json_strings(self) -> None: def test_json_strings(self) -> None:
@ -67,55 +66,80 @@ class TestFunctionality(TestBotBase):
expected_output = "value1\nvalue2\nvalue3" expected_output = "value1\nvalue2\nvalue3"
self.assertEqual(parse_maybe_json(json_array), expected_output) self.assertEqual(parse_maybe_json(json_array), expected_output)
json_string = '"value1"' json_string = '"value1"'
expected_output = 'value1' expected_output = "value1"
self.assertEqual(parse_maybe_json(json_string), expected_output) self.assertEqual(parse_maybe_json(json_string), expected_output)
json_struct = '{"This is a string."}' json_struct = '{"This is a string."}'
expected_output = 'This is a string.' expected_output = "This is a string."
self.assertEqual(parse_maybe_json(json_struct), expected_output) self.assertEqual(parse_maybe_json(json_struct), expected_output)
json_struct = '["This is a string."]' json_struct = '["This is a string."]'
expected_output = 'This is a string.' expected_output = "This is a string."
self.assertEqual(parse_maybe_json(json_struct), expected_output) self.assertEqual(parse_maybe_json(json_struct), expected_output)
json_struct = '{This is a string.}' json_struct = "{This is a string.}"
expected_output = 'This is a string.' expected_output = "This is a string."
self.assertEqual(parse_maybe_json(json_struct), expected_output) self.assertEqual(parse_maybe_json(json_struct), expected_output)
json_struct = '[This is a string.]' json_struct = "[This is a string.]"
expected_output = 'This is a string.' expected_output = "This is a string."
self.assertEqual(parse_maybe_json(json_struct), expected_output) self.assertEqual(parse_maybe_json(json_struct), expected_output)
async def test_message_lings(self) -> None: async def test_message_lings(self) -> None:
request = AIMessage('Lala', 'Hello there!', 'chat', False,) request = AIMessage(
message = {'answer': 'Test [Link](https://www.example.com/test)', "Lala",
'answer_needed': True, 'channel': 'chat', 'staff': None, 'picture': None, 'hack': False} "Hello there!",
expected = AIResponse('Test https://www.example.com/test', True, 'chat', None, None, False) "chat",
False,
)
message = {
"answer": "Test [Link](https://www.example.com/test)",
"answer_needed": True,
"channel": "chat",
"staff": None,
"picture": None,
"hack": False,
}
expected = AIResponse("Test https://www.example.com/test", True, "chat", None, None, False, False)
self.assertEqual(str(await self.bot.airesponder.post_process(request, message)), str(expected)) self.assertEqual(str(await self.bot.airesponder.post_process(request, message)), str(expected))
message = {'answer': 'Test @[Link](https://www.example.com/test)', message = {
'answer_needed': True, 'channel': 'chat', 'staff': None, 'picture': None, 'hack': False} "answer": "Test @[Link](https://www.example.com/test)",
expected = AIResponse('Test Link', True, 'chat', None, None, False) "answer_needed": True,
"channel": "chat",
"staff": None,
"picture": None,
"hack": False,
}
expected = AIResponse("Test Link", True, "chat", None, None, False, False)
self.assertEqual(str(await self.bot.airesponder.post_process(request, message)), str(expected)) self.assertEqual(str(await self.bot.airesponder.post_process(request, message)), str(expected))
message = {'answer': 'Test [Link](https://www.example.com/test) and [Link2](https://xxx) lala', message = {
'answer_needed': True, 'channel': 'chat', 'staff': None, 'picture': None, 'hack': False} "answer": "Test [Link](https://www.example.com/test) and [Link2](https://xxx) lala",
expected = AIResponse('Test https://www.example.com/test and https://xxx lala', True, 'chat', None, None, False) "answer_needed": True,
"channel": "chat",
"staff": None,
"picture": None,
"hack": False,
}
expected = AIResponse("Test https://www.example.com/test and https://xxx lala", True, "chat", None, None, False, False)
self.assertEqual(str(await self.bot.airesponder.post_process(request, message)), str(expected)) self.assertEqual(str(await self.bot.airesponder.post_process(request, message)), str(expected))
async def test_on_message_stort_path(self) -> None: async def test_on_message_stort_path(self) -> None:
message = self.create_message("Hello there! How are you?") message = self.create_message("Hello there! How are you?")
message.author.name = 'madeup_name' message.author.name = "madeup_name"
message.channel.name = 'some_channel' # type: ignore message.channel.name = "some_channel" # type: ignore
self.bot.config['short-path'] = [[r'some.*', r'madeup.*']] self.bot.config["short-path"] = [[r"some.*", r"madeup.*"]]
await self.bot.on_message(message) await self.bot.on_message(message)
self.assertEqual(self.bot.airesponder.history[-1]["content"], self.assertEqual(
self.bot.airesponder.history[-1]["content"],
'{"user": "madeup_name", "message": "Hello, how are you?",' '{"user": "madeup_name", "message": "Hello, how are you?",'
' "channel": "some_channel", "direct": false, "historise_question": true}') ' "channel": "some_channel", "direct": false, "historise_question": true}',
)
@patch("builtins.open", new_callable=mock_open) @patch("builtins.open", new_callable=mock_open)
def test_update_history_with_file(self, mock_file): def test_update_history_with_file(self, mock_file):
self.bot.airesponder.update_history({'content': '{"q": "What\'s your name?"}'}, {'content': '{"a": "AI"}'}, 10) self.bot.airesponder.update_history({"content": '{"q": "What\'s your name?"}'}, {"content": '{"a": "AI"}'}, 10)
self.assertEqual(len(self.bot.airesponder.history), 2) self.assertEqual(len(self.bot.airesponder.history), 2)
self.bot.airesponder.update_history({'content': '{"q1": "Q1"}'}, {'content': '{"a1": "A1"}'}, 2) self.bot.airesponder.update_history({"content": '{"q1": "Q1"}'}, {"content": '{"a1": "A1"}'}, 2)
self.bot.airesponder.update_history({'content': '{"q2": "Q2"}'}, {'content': '{"a2": "A2"}'}, 2) self.bot.airesponder.update_history({"content": '{"q2": "Q2"}'}, {"content": '{"a2": "A2"}'}, 2)
self.assertEqual(len(self.bot.airesponder.history), 2) self.assertEqual(len(self.bot.airesponder.history), 2)
self.bot.airesponder.history_file = "mock_file.pkl" self.bot.airesponder.history_file = "mock_file.pkl"
self.bot.airesponder.update_history({'content': '{"q": "What\'s your favorite color?"}'}, {'content': '{"a": "Blue"}'}, 10) self.bot.airesponder.update_history({"content": '{"q": "What\'s your favorite color?"}'}, {"content": '{"a": "Blue"}'}, 10)
mock_file.assert_called_once_with("mock_file.pkl", "wb") mock_file.assert_called_once_with("mock_file.pkl", "wb")
mock_file().write.assert_called_once() mock_file().write.assert_called_once()