That's very ambitious and impactful! 👍
My recommendation, though, is that you should look into whether you can implement this outside of Vim. More specifically, if you can use the Language Server Protocol to implement this style of AI-powered code completion enhancements. (See also the Wikipedia page on LSP and also langserver.org
If you do implement this as an LSP, you'll have a wealth of LSP clients on Vim that are able to query it. Straight LSP clients, such as vim-lsp or vim-lsc, but also integration with established code completion tools such as CoC.nvim and also linting engines such as ALE.
Also, since LSP is not specific to Vim, by implementing something as an LSP means immediately you get it available in most other programmer editors, since LSP support is widespread among IDEs these days.
Now, it's quite possible that LSP will not really export the kind of methods or queries that you think would be useful or required to implement an AI-powered completion engine.
But I'd still encourage you to try to implement this as an extension of LSP, using the same protocol, only exposing new objects/methods through its JSON-RPC interface. The biggest advantage of this approach is that you can then use any of the existing clients as a base to prototype your code and you don't need to start from scratch.
Also, as LSP is evolving, features for smarter and better completion are likely to make it into the protocol. (Or perhaps some standard way of adding extensions to the protocol that a client and server would agree on.) So by building a successful prototype for AI-powered completion on top of LSP, you can influence LSP's evolution to support this kind of use case.
:help omnifunc
and:help ins-completion
for how completion works in vim, and how you can program it. There are lots of other plugins that make it “automatic,” but I’m not a fan of those. You might want to take a look at kite for python.