Vanna.i + MySQL + ChromaDB : Model doesn't retrieve the data from the table

[tifDev]

Describe the bug
I ask below question and it should be able to retrieve the data from the table:

To Reproduce
Steps to reproduce the behavior:

create the main-mysql-chromadb-vanna.py

from vanna.ollama import Ollama
from vanna.chromadb import ChromaDB_VectorStore


class MyVanna(ChromaDB_VectorStore, Ollama):
    def __init__(self, config=None):
        ChromaDB_VectorStore.__init__(self, config=config)
        Ollama.__init__(self, config=config)

vn = MyVanna(config={"model": "llama3.1"})


vn.connect_to_mysql(
    host="127.0.0.1",
    dbname="test_db",
    user="root",
    password="test",
    port=3306,
)
df_information_schema = vn.run_sql("SELECT * FROM INFORMATION_SCHEMA.COLUMNS")
plan = vn.get_training_plan_generic(df_information_schema)
plan

vn.train(plan=plan)


vn.train(
    ddl="""
    CREATE DATABASE test_db
    )
"""
)


vn.train(
    ddl="""
    CREATE TABLE users (
        id INT AUTO_INCREMENT PRIMARY KEY,
        name VARCHAR(100),
        email VARCHAR(100)
    )
"""
)

vn.train(documentation="The table _users_ is within _test_db_ Database")


vn.train(sql="SELECT * FROM test_db.users WHERE name = 'John Doe'")


from vanna.flask import VannaFlaskApp

app = VannaFlaskApp(vn)
app.run()

create and run a local mysql instance

docker pull mysql:latest


sudo docker run --name db_container_mysql -e MYSQL_ROOT_PASSWORD=test -e MYSQL_DATABASE=test_db -p 3306:3306 -p 33060:33060 -d mysql


# Connect to DB to run sql commands
sudo docker exec -it db_container_mysql mysql -uroot -pmysql -h 127.0.0.1 -P 3306 -u root -p


# Generate data in DB

CREATE DATABASE test_db;
USE test_db;

CREATE USER 'test_user'@'localhost';
GRANT ALL PRIVILEGES ON *.* TO 'test_user'@'localhost' IDENTIFIED BY 'test_user_pass';

CREATE TABLE users (
    id INT AUTO_INCREMENT PRIMARY KEY,
    name VARCHAR(100),
    email VARCHAR(100)
);

INSERT INTO users (name, email) VALUES ('John Doe', 'john.doe@example.com');
INSERT INTO users (name, email) VALUES ('Jane Smith', 'jane.smith@example.com');

select * from users;
mysql> select * from users;
+----+------------+------------------------+
| id | name       | email                  |
+----+------------+------------------------+
|  1 | John Doe   | john.doe@example.com   |
|  2 | Jane Smith | jane.smith@example.com |
+----+------------+------------------------+
2 rows in set (0.00 sec)

Expected behavior
Be able to retrieve the data from the table.

Error logs/Screenshots
If applicable, add logs/screenshots to give more information about the issue.

mysql> show databases;
+--------------------+
| Database |
+--------------------+
| information_schema |
| mysql |
| performance_schema |
| sys |
| test_db |
+--------------------+
5 rows in set (0.02 sec)

mysql> use test_db
Reading table information for completion of table and column names
You can turn off this feature to get a quicker startup with -A

Database changed
mysql> show tables;
+-------------------+
| Tables_in_test_db |
+-------------------+
| users |
+-------------------+
1 row in set (0.00 sec)

mysql> select * from users;
+----+------------+------------------------+
| id | name | email |
+----+------------+------------------------+
| 1 | John Doe | john.doe@example.com |
| 2 | Jane Smith | jane.smith@example.com |
+----+------------+------------------------+
2 rows in set (0.00 sec)

Desktop (please complete the following information where):

• OS: Ubuntu
• Version: 20.04
• Python: 3.10
• Vanna: 0.6.4

[zainhoda]

This may potentially just be a weird LLM behavior where it doesn't want to respond to PII. What does it give you if you ask a question like "How many users are there?"