Using custom libraries while generating code using Assistants API

I have an assistant that as a part of a conversation generates code. However this code will use a custom python library. I would like to provide a list of rules for the assistant while generating code with this library (Typically what it can import, syntax for certain function calls). Any ideas how I can achieve this?
As a first step, I wanted the imports to be correct. I added a knowledge file with all the allowed imports. However, this seems to be ignored.
Any help would be much appreciated.