knowledge base completion

Author: ea-rus

mindsdb_sdk/connectors/rest_api.py

@@ -439,4 +439,11 @@ def create_knowledge_base(self, project: str, data):
 
         return r.json()
 
+    def knowledge_base_completion(self, project: str, knowledge_base_name, payload):
+        r = self.session.post(
+            self.url + f'/api/projects/{project}/knowledge_bases/{knowledge_base_name}/completions',
+            json=payload
+        )
+        _raise_for_status(r)
+        return r.json()
 

mindsdb_sdk/knowledge_bases.py

@@ -206,6 +206,11 @@ def insert_query(self, data: Query):
             data={'query': data.sql}
         )
 
+    def completion(self, query, **data):
+        data['query'] = query
+
+        return self.api.knowledge_base_completion(self.project.name, self.name, data)
+
 
 class KnowledgeBases(CollectionBase):
     """

tests/test_sdk.py

@@ -1259,6 +1259,16 @@ def check_project_kb(self, project, model, database, mock_get, mock_del, mock_po
             },
         }}}
 
+        # completion
+        kb.completion('hi', type='chat', llm_model='gpt-4')
+        args, kwargs = mock_post.call_args
+        assert args[0] == f'{DEFAULT_CLOUD_API_URL}/api/projects/{project.name}/knowledge_bases/{kb.name}/completions'
+        assert kwargs == {'json': {
+            'query': 'hi',
+            'type': 'chat',
+            'llm_model': 'gpt-4',
+        }}
+
         # drop
         project.knowledge_bases.drop('kb2')