Fix mypy type errors

[nikteliy]

The codebase has type errors when running mypy. These should be fixed to ensure type safety and catch potential bugs early.