Unable to access get_training_plan_postgres using VannaDefault

[coderberry]

I have the following code that I am trying to make work:

from vanna.remote import VannaDefault

vn = VannaDefault(model=os.environ.get("VANNA_MODEL"), api_key=os.environ.get("VANNA_API_KEY"))
vn.connect_to_postgres(
    host=os.environ.get("PHM_POSTGRES_HOST"),
    dbname=os.environ.get("PHM_POSTGRES_DB"),
    user=os.environ.get("PHM_POSTGRES_USER"),
    password=os.environ.get("PHM_POSTGRES_PASSWORD"),
    port=5432
)

plan = vn.get_training_plan_postgres(
    filter_databases=[os.environ.get("PHM_POSTGRES_DB")],
    filter_schemas=["public"],
    include_information_schema=False,
    use_historical_queries=True
)

When I run this (all env vars are set correctly), I get the following error:

AttributeError                            Traceback (most recent call last)
Cell In[22], [line 2](vscode-notebook-cell:?execution_count=22&line=2)
      [1](vscode-notebook-cell:?execution_count=22&line=1) # This will break up the information schema into bite-sized chunks that can be referenced by the LLM
----> [2](vscode-notebook-cell:?execution_count=22&line=2) plan = phm_vn.get_training_plan_postgres(
      [3](vscode-notebook-cell:?execution_count=22&line=3)      filter_databases=[os.environ.get("PHM_POSTGRES_DB")],
      [4](vscode-notebook-cell:?execution_count=22&line=4) 	filter_schemas=["public"],
      [5](vscode-notebook-cell:?execution_count=22&line=5) 	include_information_schema=False,
      [6](vscode-notebook-cell:?execution_count=22&line=6) 	use_historical_queries=True
      [7](vscode-notebook-cell:?execution_count=22&line=7) )
      [8](vscode-notebook-cell:?execution_count=22&line=8) plan

AttributeError: 'VannaDefault' object has no attribute 'get_training_plan_postgres'

What am I doing wrong? I want to limit the scope of what is being trained specifically to the database and schema.

[zainhoda]

I think you might have found some old documentation -- try following this notebook for training:

https://vanna.ai/docs/postgres-openai-vanna-vannadb/