Repository: mindsdb/mindsdb
Pull request: Shared pgvector in cloud
## Description

**Updates:**

Using default vector storage
- if `KB_PGVECTOR_URL` env variable is set:
  - pgvector database with name `kb_pgvector_store` will be created 
  - it will use connection args from `KB_PGVECTOR_URL` (exampe: postgresql://user:password@server:5432/db)
    - database will be created if not exists
  - it will create table for every knowledge base, the name of the table is:
    - `t_{company_id}_{KB name}`
  - after deleting of KB database will be kept 
- else chromadb database will be created (the same behavior as is)

Using default embedding model 
- New behavior:
  - the name `kb_default_embedding_model` is used for model
  - if it is not exists: it will be created with engine `langchain_embedding`
  - otherwise this model will be used for KB
  - after deleting of KB embedding model will be kept 
- Old behavior:
  - a new model is create for every KB with name `{KB name}_default_model`


KB_PGVECTOR_URL var was added here:
https://github.com/mindsdb/INTERNAL-mindsdb-build-deploy-to-kubernetes/pull/45
https://github.com/mindsdb/INTERNAL-mindsdb-build-deploy-to-kubernetes/pull/46

Tested on dev:
https://www.loom.com/share/0b636961d988496f8516c9df7dcd43c4

Fixes https://linear.app/mindsdb/issue/ML-27/add-default-vector-database-for-cloud
Fixes #9340

## Type of change

- [x] ⚡ New feature (non-breaking change which adds functionality)

## Verification Process

To ensure the changes are working as expected:

 - [ ]   Test Location: Specify the URL or path for testing.
 - [ ]   Verification Steps: Outline the steps or queries needed to validate the change. Include any data, configurations, or actions required to reproduce or see the new functionality.

## Additional Media:

- [ ] I have attached a brief loom video or screenshots showcasing the new functionality or change.

## Checklist:

- [x] My code follows the style guidelines(PEP 8) of MindsDB.
- [x] I have appropriately commented on my code, especially in complex areas.
- [ ] Necessary documentation updates are either made or tracked in issues.
- [ ] Relevant unit and integration tests are updated or added.





Problem statement:
[Bug]: Unable to create knowledge base with default embedding model
### Short description of current behavior

Creating knowledge base without  embedding model defined should create langchain_embedding model under the hood and use it.
But after langchain_embedding is not permanent handler (changed in this PR) mindsdb unable create model. 
How to fix:
- create langchain_embedding ml engine automatically when model is created
  - in KB controller
  - or in `get_ml_handler` function

### Video or screenshots

_No response_

### Expected behavior

_No response_

### How to reproduce the error

Execute query:
```sql
CREATE KNOWLEDGE BASE my_kb
```
Get error in response:
```
There is no integration with name 'langchain_embedding': unknown
```


### Anything else?

_No response_

Hints from issue discussion:

Commit message: don't delete default model

Patch to review:
diff --git a/mindsdb/api/executor/datahub/datanodes/project_datanode.py b/mindsdb/api/executor/datahub/datanodes/project_datanode.py
index 3189045c9fc..7a08c94689a 100644
--- a/mindsdb/api/executor/datahub/datanodes/project_datanode.py
+++ b/mindsdb/api/executor/datahub/datanodes/project_datanode.py
@@ -181,4 +181,4 @@ def create_table(self, table_name: Identifier, result_set=None, is_replace=False
 
             df = result_set.to_df()
             return kb_table.insert(df)
-        raise NotImplementedError(f"Cant create table {table_name}")
+        raise NotImplementedError(f"Can't create table {table_name}")
diff --git a/mindsdb/integrations/handlers/pgvector_handler/pgvector_handler.py b/mindsdb/integrations/handlers/pgvector_handler/pgvector_handler.py
index c40d1fb9b65..8cf380b6aa8 100644
--- a/mindsdb/integrations/handlers/pgvector_handler/pgvector_handler.py
+++ b/mindsdb/integrations/handlers/pgvector_handler/pgvector_handler.py
@@ -1,22 +1,24 @@
+import os
 import json
 from typing import List, Union
+from urllib.parse import urlparse
 
 import pandas as pd
 import psycopg
-from mindsdb_sql import ASTNode, Parameter, Identifier, Update, BinaryOperation
+from mindsdb_sql.parser.ast import Parameter, Identifier, Update, BinaryOperation
 from pgvector.psycopg import register_vector
 
 from mindsdb.integrations.handlers.postgres_handler.postgres_handler import (
     PostgresHandler,
 )
-from mindsdb.integrations.libs.response import RESPONSE_TYPE
-from mindsdb.integrations.libs.response import HandlerResponse
+from mindsdb.integrations.libs.response import RESPONSE_TYPE, HandlerResponse as Response
 from mindsdb.integrations.libs.vectordatabase_handler import (
     FilterCondition,
     VectorStoreHandler,
 )
 from mindsdb.utilities import log
 from mindsdb.utilities.profiler import profiler
+from mindsdb.utilities.context import context as ctx
 
 logger = log.getLogger(__name__)
 
@@ -30,8 +32,42 @@ class PgVectorHandler(VectorStoreHandler, PostgresHandler):
     def __init__(self, name: str, **kwargs):
 
         super().__init__(name=name, **kwargs)
+        self._is_shared_db = False
         self.connect()
 
+    def _make_connection_args(self):
+        cloud_pgvector_url = os.environ.get('KB_PGVECTOR_URL')
+        if cloud_pgvector_url is not None:
+            result = urlparse(cloud_pgvector_url)
+            self.connection_args = {
+                'host': result.hostname,
+                'port': result.port,
+                'user': result.username,
+                'password': result.password,
+                'database': result.path[1:]
+            }
+            self._is_shared_db = True
+        return super()._make_connection_args()
+
+    def get_tables(self) -> Response:
+        # Hide list of tables from all users
+        if self._is_shared_db:
+            return Response(RESPONSE_TYPE.OK)
+        return super().get_tables()
+
+    def native_query(self, query) -> Response:
+        # Prevent execute native queries
+        if self._is_shared_db:
+            return Response(RESPONSE_TYPE.OK)
+        return super().native_query(query)
+
+    def raw_query(self, query, params=None) -> Response:
+        resp = super().native_query(query, params)
+        if resp.resp_type == RESPONSE_TYPE.ERROR:
+            raise Exception(resp.error_message)
+        if resp.resp_type == RESPONSE_TYPE.TABLE:
+            return resp.data_frame
+
     @profiler.profile()
     def connect(self) -> psycopg.connection:
         """
@@ -46,10 +82,11 @@ def connect(self) -> psycopg.connection:
                 logger.info("pg_vector extension loaded")
 
             except psycopg.Error as e:
+                self.connection.rollback()
                 logger.error(
                     f"Error loading pg_vector extension, ensure you have installed it before running, {e}!"
                 )
-                return HandlerResponse(resp_type=RESPONSE_TYPE.ERROR, error_message=str(e))
+                raise
 
         # register vector type with psycopg2 connection
         register_vector(self.connection)
@@ -157,6 +194,13 @@ def _build_select_query(
             # if no filter conditions, return all rows
             return f"SELECT {targets} FROM {table_name} {after_from_clause}"
 
+    def _check_table(self, table_name: str):
+        # Apply namespace for a user
+        if self._is_shared_db:
+            company_id = ctx.company_id or 'x'
+            return f't_{company_id}_{table_name}'
+        return table_name
+
     def select(
         self,
         table_name: str,
@@ -168,17 +212,14 @@ def select(
         """
         Retrieve the data from the SQL statement with eliminated rows that dont satisfy the WHERE condition
         """
+        table_name = self._check_table(table_name)
+
         if columns is None:
             columns = ["id", "content", "embeddings", "metadata"]
 
-        with self.connection.cursor() as cur:
-            query = self._build_select_query(table_name, columns, conditions, limit, offset)
-            cur.execute(query)
-
-            self.connection.commit()
-            result = cur.fetchall()
+        query = self._build_select_query(table_name, columns, conditions, limit, offset)
+        result = self.raw_query(query)
 
-        result = pd.DataFrame(result, columns=columns)
         # ensure embeddings are returned as string so they can be parsed by mindsdb
         if "embeddings" in columns:
             result["embeddings"] = result["embeddings"].astype(str)
@@ -189,11 +230,10 @@ def create_table(self, table_name: str, if_not_exists=True):
         """
         Run a create table query on the pgvector database.
         """
-        with self.connection.cursor() as cur:
-            cur.execute(
-                f"CREATE TABLE IF NOT EXISTS {table_name} (id text PRIMARY KEY, content text, embeddings vector, metadata jsonb)"
-            )
-            self.connection.commit()
+        table_name = self._check_table(table_name)
+
+        query = f"CREATE TABLE IF NOT EXISTS {table_name} (id text PRIMARY KEY, content text, embeddings vector, metadata jsonb)"
+        self.raw_query(query)
 
     def insert(
         self, table_name: str, data: pd.DataFrame
@@ -201,6 +241,8 @@ def insert(
         """
         Insert data into the pgvector table database.
         """
+        table_name = self._check_table(table_name)
+
         data_dict = data.to_dict(orient="list")
 
         if 'metadata' in data_dict:
@@ -212,9 +254,7 @@ def insert(
 
         insert_statement = f"INSERT INTO {table_name} ({columns}) VALUES ({values})"
 
-        with self.connection.cursor() as cur:
-            cur.executemany(insert_statement, transposed_data)
-            self.connection.commit()
+        self.raw_query(insert_statement, params=transposed_data)
 
     def update(
         self, table_name: str, data: pd.DataFrame, key_columns: List[str] = None
@@ -222,6 +262,7 @@ def update(
         """
         Udate data into the pgvector table database.
         """
+        table_name = self._check_table(table_name)
 
         where = None
         update_columns = {}
@@ -250,45 +291,36 @@ def update(
             where=where
         )
 
-        with self.connection.cursor() as cur:
-            transposed_data = []
-            for _, record in data.iterrows():
-                row = [
-                    record[col]
-                    for col in update_columns.keys()
-                ]
-                for key_column in key_columns:
-                    row.append(record[key_column])
-                transposed_data.append(row)
-
-            query_str = self.renderer.get_string(query)
-            cur.executemany(query_str, transposed_data)
-            self.connection.commit()
+        transposed_data = []
+        for _, record in data.iterrows():
+            row = [
+                record[col]
+                for col in update_columns.keys()
+            ]
+            for key_column in key_columns:
+                row.append(record[key_column])
+            transposed_data.append(row)
+
+        query_str = self.renderer.get_string(query)
+        self.raw_query(query_str, transposed_data)
 
     def delete(
         self, table_name: str, conditions: List[FilterCondition] = None
     ):
+        table_name = self._check_table(table_name)
 
         filter_conditions = self._translate_conditions(conditions)
         where_clause = self._construct_where_clause(filter_conditions)
 
-        with self.connection.cursor() as cur:
-
-            # convert search embedding to string
-
-            # we need to use the <-> operator to search for similar vectors,
-            # so we need to convert the string to a vector and also use a threshold (e.g. 0.5)
-
-            query = (
-                f"DELETE FROM {table_name} {where_clause}"
-            )
-            cur.execute(query)
-            self.connection.commit()
+        query = (
+            f"DELETE FROM {table_name} {where_clause}"
+        )
+        self.raw_query(query)
 
     def drop_table(self, table_name: str, if_exists=True):
         """
         Run a drop table query on the pgvector database.
         """
-        with self.connection.cursor() as cur:
-            cur.execute(f"DROP TABLE IF EXISTS {table_name}")
-            self.connection.commit()
+        table_name = self._check_table(table_name)
+        self.raw_query(f"DROP TABLE IF EXISTS {table_name}")
+
diff --git a/mindsdb/interfaces/knowledge_base/controller.py b/mindsdb/interfaces/knowledge_base/controller.py
index 11e7fccdba0..01ed66e5641 100644
--- a/mindsdb/interfaces/knowledge_base/controller.py
+++ b/mindsdb/interfaces/knowledge_base/controller.py
@@ -1,3 +1,4 @@
+import os
 import copy
 from typing import List
 
@@ -18,8 +19,11 @@
 import mindsdb.interfaces.storage.db as db
 from mindsdb.integrations.libs.vectordatabase_handler import TableField
 from mindsdb.interfaces.database.projects import ProjectController
+from mindsdb.interfaces.model.functions import PredictorRecordNotFound
 from mindsdb.utilities.exception import EntityExistsError, EntityNotExistsError
 
+from mindsdb.api.executor.command_executor import ExecuteCommands
+
 
 class KnowledgeBaseTable:
     """
@@ -352,10 +356,7 @@ def add(
 
         if embedding_model is None:
             # create default embedding model
-            model_name = self._create_default_embedding_model(project.name, name, params=params)
-
-            # memorize to remove it later
-            params['embedding_model'] = model_name
+            model_name = self._get_default_embedding_model(project.name, params=params)
 
         else:
             # get embedding model from input
@@ -376,14 +377,20 @@ def add(
 
         # search for the vector database table
         if storage is None:
-            # create chroma db with same name
-            vector_table_name = "default_collection"
-            vector_db_name = self._create_persistent_chroma(
-                name
-            )
 
-            # memorize to remove it later
-            params['vector_storage'] = vector_db_name
+            cloud_pg_vector = os.environ.get('KB_PGVECTOR_URL')
+            if cloud_pg_vector:
+                vector_table_name = name
+                vector_db_name = self._create_persistent_pgvector()
+            else:
+                # create chroma db with same name
+                vector_table_name = "default_collection"
+                vector_db_name = self._create_persistent_chroma(
+                    name
+                )
+
+                # memorize to remove it later
+                params['vector_storage'] = vector_db_name
         elif len(storage.parts) != 2:
             raise ValueError('Storage param has to be vector db with table')
         else:
@@ -408,6 +415,18 @@ def add(
         db.session.commit()
         return kb
 
+    def _create_persistent_pgvector(self):
+        """Create default vector database for knowledge base, if not specified"""
+
+        vector_store_name = "kb_pgvector_store"
+
+        # check if exists
+        if self.session.integration_controller.get(vector_store_name):
+            return vector_store_name
+
+        self.session.integration_controller.add(vector_store_name, 'pgvector', {})
+        return vector_store_name
+
     def _create_persistent_chroma(self, kb_name, engine="chromadb"):
         """Create default vector database for knowledge base, if not specified"""
 
@@ -423,13 +442,25 @@ def _create_persistent_chroma(self, kb_name, engine="chromadb"):
         self.session.integration_controller.add(vector_store_name, engine, connection_args)
         return vector_store_name
 
-    def _create_default_embedding_model(self, project_name, kb_name, engine="langchain_embedding", params: dict = None):
+    def _get_default_embedding_model(self, project_name, engine="langchain_embedding", params: dict = None):
         """create a default embedding model for knowledge base, if not specified"""
-        model_name = f"{kb_name}_default_model"
-        using_args = {}
+        model_name = "kb_default_embedding_model"
+
+        # check exists
+        try:
+            model = self.session.model_controller.get_model(model_name, project_name=project_name)
+            if model is not None:
+                return model_name
+        except PredictorRecordNotFound:
+            pass
+
+        using_args = {
+            'engine': engine
+        }
         if engine == 'langchain_embedding':
             # Use default embeddings.
             using_args['class'] = 'openai'
+
         # Include API key if provided.
         using_args.update({k: v for k, v in params.items() if 'api_key' in k})
         statement = CreatePredictor(
@@ -439,12 +470,9 @@ def _create_default_embedding_model(self, project_name, kb_name, engine="langcha
                 Identifier(parts=[TableField.EMBEDDINGS.value])
             ]
         )
-        ml_handler = self.session.integration_controller.get_ml_handler(engine)
 
-        _ = self.session.model_controller.create_model(
-            statement,
-            ml_handler
-        )
+        command_executor = ExecuteCommands(self.session)
+        command_executor.answer_create_predictor(statement, project_name)
 
         return model_name
 


Code review comments:
diff_hunk: @@ -30,8 +32,42 @@ class PgVectorHandler(VectorStoreHandler, PostgresHandler):
     def __init__(self, name: str, **kwargs):
 
         super().__init__(name=name, **kwargs)
+        self._is_shared_db = False
         self.connect()
 
+    def _make_connection_args(self):
+        cloud_pgvector_url = os.environ.get('KB_PGVECTOR_URL')
+        if cloud_pgvector_url is not None:
+            result = urlparse(cloud_pgvector_url)
+            self.connection_args = {
+                'host': result.hostname,
+                'port': result.port,
+                'user': result.username,
+                'password': result.password,
+                'database': result.path[1:]
+            }
+            self._is_shared_db = True
+        return super()._make_connection_args()
+
+    def get_tables(self) -> Response:
+        # Hide list of tables from all users
+        if self._is_shared_db:
+            return Response(RESPONSE_TYPE.OK)
+        return super().get_tables()
+
+    def native_query(self, query) -> Response:
+        # Prevent execute native queries
+        if self._is_shared_db:
+            return Response(RESPONSE_TYPE.OK)
+        return super().native_query(query)
+
+    def raw_query(self, query, params=None) -> Response:
+        resp = super().native_query(query, params)
+        if resp.resp_type == RESPONSE_TYPE.ERROR:
+            raise Exception(resp.error_message)
path: mindsdb/integrations/handlers/pgvector_handler/pgvector_handler.py
text: @user1:
nit: would prefer `RuntimeError` instead of using base `Exception`

@author:
Changed