Merge pull request #9069 from mindsdb/kb-columns

Kb updates 04.2024
mindsdb · Apr 15, 2024 · a30f358 · a30f358
2 parents af106f3 + eb81171
commit a30f358
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 6 deletions.
diff --git a/mindsdb/api/executor/datahub/datanodes/information_schema_datanode.py b/mindsdb/api/executor/datahub/datanodes/information_schema_datanode.py
@@ -315,7 +315,7 @@ class InformationSchemaDataNode(DataNode):
             "IS_RUNNING",
             "LAST_ERROR",
         ],
-        "KNOWLEDGE_BASES": ["NAME", "PROJECT", "MODEL", "STORAGE"],
+        "KNOWLEDGE_BASES": ["NAME", "PROJECT", "MODEL", "STORAGE", "PARAMS"],
         "SKILLS": ["NAME", "PROJECT", "TYPE", "PARAMS"],
         "AGENTS": [
             "NAME",
@@ -672,7 +672,7 @@ def _get_knowledge_bases(self, query: ASTNode = None):
 
         columns = self.information_schema['KNOWLEDGE_BASES']
 
-        # columns: NAME, PROJECT, MODEL, STORAGE
+        # columns: NAME, PROJECT, MODEL, STORAGE, PARAMS
         data = []
 
         for kb in kb_list:
@@ -684,7 +684,8 @@ def _get_knowledge_bases(self, query: ASTNode = None):
                 kb.name,
                 project_name,
                 embedding_model.name if embedding_model is not None else None,
-                vector_database_name + '.' + kb.vector_database_table
+                vector_database_name + '.' + kb.vector_database_table,
+                to_json(kb.params),
             ))
 
         return pd.DataFrame(data, columns=columns)

diff --git a/mindsdb/interfaces/knowledge_base/controller.py b/mindsdb/interfaces/knowledge_base/controller.py
@@ -251,7 +251,10 @@ def _get_vector_db(self):
         helper to get vector db handler
         """
         if self._vector_db is None:
-            database_name = db.Integration.query.get(self._kb.vector_database_id).name
+            database = db.Integration.query.get(self._kb.vector_database_id)
+            if database is None:
+                raise ValueError('Vector database not found. Is it deleted?')
+            database_name = database.name
             self._vector_db = self.session.integration_controller.get_data_handler(database_name)
         return self._vector_db
 
@@ -276,12 +279,15 @@ def _df_to_embeddings(self, df: pd.DataFrame) -> pd.DataFrame:
 
         project_datanode = self.session.datahub.get(model_project.name)
 
-        # TODO adjust input
+        # keep only content
+        df = df[[TableField.CONTENT.value]]
+
         input_col = model_rec.learn_args.get('using', {}).get('question_column')
+
         if input_col is not None and input_col != TableField.CONTENT.value:
             df = df.rename(columns={TableField.CONTENT.value: input_col})
 
-        data = df[[TableField.CONTENT.value]].to_dict('records')
+        data = df.to_dict('records')
 
         df_out = project_datanode.predict(
             model_name=model_rec.name,