fix insert into kb

mindsdb · Apr 12, 2024 · eb81171 · eb81171
1 parent b15f6c7
commit eb81171
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/mindsdb/interfaces/knowledge_base/controller.py b/mindsdb/interfaces/knowledge_base/controller.py
@@ -253,7 +253,7 @@ def _get_vector_db(self):
         if self._vector_db is None:
             database = db.Integration.query.get(self._kb.vector_database_id)
             if database is None:
-                raise ValueError(f'Vector database not found. Is it deleted?')
+                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
@@ -279,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,