diff --git a/mindsdb_sdk/agents.py b/mindsdb_sdk/agents.py index 1b081cf..8c7f01b 100644 --- a/mindsdb_sdk/agents.py +++ b/mindsdb_sdk/agents.py @@ -10,9 +10,6 @@ from mindsdb_sdk.skills import Skill from mindsdb_sdk.utils.objects_collection import CollectionBase -_DEFAULT_LLM_MODEL = 'gpt-4o' -_DEFAULT_LLM_PROMPT = 'Answer the user"s question in a helpful way: {{question}}' - class AgentCompletion: """ @@ -59,17 +56,19 @@ class Agent: Create a new agent: - >>> model = models.get('my_model') # Or use models.create(...) - >>> # Connect your agent to a MindsDB table. - >>> text_to_sql_skill = skills.create('text_to_sql', 'sql', { 'tables': ['my_table'], 'database': 'my_database' }) - >>> agent = agents.create('my_agent', model, [text_to_sql_skill]) + >>> agent = agents.create( + 'my_agent', + model={ + 'model_name': 'gpt-3.5-turbo', + 'provider': 'openai', + 'api_key': 'your_openai_api_key_here' + }, + data={'tables': ['my_database.my_table'], 'knowledge_base': 'my_kb'} + ) Update an agent: - >>> new_model = models.get('new_model') - >>> agent.model_name = new_model.name - >>> new_skill = skills.create('new_skill', 'sql', { 'tables': ['new_table'], 'database': 'new_database' }) - >>> updated_agent.skills.append(new_skill) + >>> agent.data['tables'].append('my_database.my_new_table') >>> updated_agent = agents.update('my_agent', agent) Delete an agent by name: @@ -80,21 +79,25 @@ class Agent: def __init__( self, name: str, - model_name: str, - skills: List[Skill], - params: dict, created_at: datetime.datetime, updated_at: datetime.datetime, + model: Union[Model, str, dict] = None, + skills: List[Skill] = [], provider: str = None, + data: dict = {}, + prompt_template: str = None, + params: dict = {}, collection: CollectionBase = None ): self.name = name - self.model_name = model_name - self.provider = provider - self.skills = skills - self.params = params self.created_at = created_at self.updated_at = updated_at + self.model = model + self.skills = skills + self.provider = provider + self.data = data + self.prompt_template = prompt_template + self.params = params self.collection = collection def completion(self, messages: List[dict]) -> AgentCompletion: @@ -183,7 +186,7 @@ def __repr__(self): def __eq__(self, other): if self.name != other.name: return False - if self.model_name != other.model_name: + if self.model != other.model: return False if self.provider != other.provider: return False @@ -197,14 +200,22 @@ def __eq__(self, other): @classmethod def from_json(cls, json: dict, collection: CollectionBase): + skills = [] + if json.get('skills'): + skills = [Skill.from_json(skill) for skill in json['skills']] + + model = json.get('model') or json.get('model_name') + return cls( json['name'], - json['model_name'], - [Skill.from_json(skill) for skill in json['skills']], - json['params'], json['created_at'], json['updated_at'], - json['provider'], + model, + skills, + json.get('provider'), + json.get('data', {}), + json.get('prompt_template'), + json.get('params', {}), collection ) @@ -292,41 +303,32 @@ def completion_stream_v2(self, name, messages: List[dict]) -> Iterable[object]: return self.api.agent_completion_stream_v2(self.project.name, name, messages) def _create_default_knowledge_base(self, agent: Agent, name: str) -> KnowledgeBase: - # Make sure default ML engine for embeddings exists. try: - _ = self.ml_engines.get('langchain_embedding') - except AttributeError: - _ = self.ml_engines.create('langchain_embedding', 'langchain_embedding') - # Include API keys in embeddings. - if agent.provider == "mindsdb": - agent_model = self.models.get(agent.model_name) - training_options = json.loads(agent_model.data.get('training_options', '{}')) - training_options_using = training_options.get('using', {}) - api_key_params = {k: v for k, v in training_options_using.items() if 'api_key' in k} - kb = self.knowledge_bases.create(name, params=api_key_params) - else: kb = self.knowledge_bases.create(name) - # Wait for underlying embedding model to finish training. - kb.model.wait_complete() - return kb + return kb + except Exception as e: + raise ValueError( + f"Failed to automatically create knowledge base for agent {agent.name}. " + "Either provide an existing knowledge base name, " + "or set your default embedding model via server.config.set_default_embedding_model(...) or through the MindsDB UI." + ) - def add_files(self, name: str, file_paths: List[str], description: str, knowledge_base: str = None): + def add_files(self, name: str, file_paths: List[str], description: str = None): """ Add a list of files to the agent for retrieval. :param name: Name of the agent :param file_paths: List of paths or URLs to the files to be added. :param description: Description of the file. Used by agent to know when to do retrieval - :param knowledge_base: Name of an existing knowledge base to be used. Will create a default knowledge base if not given. """ if not file_paths: return + + agent = self.get(name) filename_no_extension = '' - all_filenames = [] for file_path in file_paths: filename = file_path.split('/')[-1].lower() filename_no_extension = filename.split('.')[0] - all_filenames.append(filename_no_extension) try: _ = self.api.get_file_metadata(filename_no_extension) except HTTPError as e: @@ -335,48 +337,36 @@ def add_files(self, name: str, file_paths: List[str], description: str, knowledg # upload file to mindsdb self.api.upload_file(filename, file_path) - # Insert uploaded files into new knowledge base. - agent = self.get(name) - if knowledge_base is not None: - kb = self.knowledge_bases.get(knowledge_base) - else: - kb_name = f'{name.lower()}_{filename_no_extension}_{uuid4().hex}_kb' - kb = self._create_default_knowledge_base(agent, kb_name) + # Add file to agent's data if it hasn't been added already. + if 'tables' not in agent.data or f'files.{filename_no_extension}' not in agent.data['tables']: + agent.data.setdefault('tables', []).append(f'files.{filename_no_extension}') + + # Add the description provided to the agent's prompt template. + if description: + agent.prompt_template = (agent.prompt_template or '') + f'\n{description}' - # Insert the entire file. - kb.insert_files(all_filenames) - - # Make sure skill name is unique. - skill_name = f'{filename_no_extension}_retrieval_skill_{uuid4().hex}' - retrieval_params = { - 'source': kb.name, - 'description': description, - } - file_retrieval_skill = self.skills.create(skill_name, 'retrieval', retrieval_params) - agent.skills.append(file_retrieval_skill) self.update(agent.name, agent) - def add_file(self, name: str, file_path: str, description: str, knowledge_base: str = None): + def add_file(self, name: str, file_path: str, description: str = None): """ Add a file to the agent for retrieval. :param name: Name of the agent :param file_path: Path to the file to be added, or name of existing file. :param description: Description of the file. Used by agent to know when to do retrieval - :param knowledge_base: Name of an existing knowledge base to be used. Will create a default knowledge base if not given. """ - self.add_files(name, [file_path], description, knowledge_base) + self.add_files(name, [file_path], description) def add_webpages( - self, - name: str, - urls: List[str], - description: str, - knowledge_base: str = None, - crawl_depth: int = 1, - limit: int = None, - filters: List[str] = None - ): + self, + name: str, + urls: List[str], + description: str = None, + knowledge_base: str = None, + crawl_depth: int = 1, + limit: int = None, + filters: List[str] = None + ): """ Add a list of webpages to the agent for retrieval. @@ -403,25 +393,26 @@ def add_webpages( # Insert crawled webpage. kb.insert_webpages(urls, crawl_depth=crawl_depth, filters=filters, limit=limit) - # Make sure skill name is unique. - skill_name = f'web_retrieval_skill_{uuid4().hex}' - retrieval_params = { - 'source': kb.name, - 'description': description, - } - webpage_retrieval_skill = self.skills.create(skill_name, 'retrieval', retrieval_params) - agent.skills.append(webpage_retrieval_skill) + # Add knowledge base to agent's data if it hasn't been added already. + if 'knowledge_bases' not in agent.data or kb.name not in agent.data['knowledge_bases']: + agent.data.setdefault('knowledge_bases', []).append(kb.name) + + # Add the description provided to the agent's prompt template. + if description: + agent.prompt_template = (agent.prompt_template or '') + f'\n{description}' + self.update(agent.name, agent) def add_webpage( - self, - name: str, - url: str, - description: str, - knowledge_base: str = None, - crawl_depth: int = 1, - limit: int = None, - filters: List[str] = None): + self, + name: str, + url: str, + description: str = None, + knowledge_base: str = None, + crawl_depth: int = 1, + limit: int = None, + filters: List[str] = None + ): """ Add a webpage to the agent for retrieval. @@ -436,56 +427,61 @@ def add_webpage( self.add_webpages(name, [url], description, knowledge_base=knowledge_base, crawl_depth=crawl_depth, limit=limit, filters=filters) - def add_database(self, name: str, database: str, tables: List[str], description: str): + def add_database(self, name: str, database: str, tables: List[str] = None, description: str = None): """ Add a database to the agent for retrieval. :param name: Name of the agent :param database: Name of the database to be added. - :param tables: List of tables to be added. + :param tables: List of tables to be added. If not provided, the entire database will be added. :param description: Description of the database. Used by agent to know when to do retrieval. """ # Make sure database exists. db = self.databases.get(database) - # Make sure tables exist. - all_table_names = set([t.name for t in db.tables.list()]) - for t in tables: - if t not in all_table_names: - raise ValueError(f'Table {t} does not exist in database {database}.') - - # Make sure skill name is unique. - skill_name = f'{database}_sql_skill_{uuid4().hex}' - sql_params = { - 'database': database, - 'tables': tables, - 'description': description, - } - database_sql_skill = self.skills.create(skill_name, 'sql', sql_params) + agent = self.get(name) - if not agent.params: - agent.params = {} - if 'prompt_template' not in agent.params: - # Set default prompt template. This is for langchain agent check. - agent.params['prompt_template'] = 'using mindsdb sqltoolbox' + if tables: + # Ensure the tables exist. + all_table_names = set([t.name for t in db.tables.list()]) + for t in tables: + if t not in all_table_names: + raise ValueError(f'Table {t} does not exist in database {database}.') + + # Add table to agent's data if it hasn't been added already. + if 'tables' not in agent.data or f'{database}.{t}' not in agent.data['tables']: + agent.data.setdefault('tables', []).append(f'{database}.{t}') + + else: + # If no tables are provided, add the database itself. + if 'tables' not in agent.data or f'{database}.*' not in agent.data['tables']: + agent.data.setdefault('tables', []).append(f'{database}.*') + + # Add the description provided to the agent's prompt template. + if description: + agent.prompt_template = (agent.prompt_template or '') + f'\n{description}' - agent.skills.append(database_sql_skill) self.update(agent.name, agent) def create( - self, - name: str, - model: Union[Model, dict, str] = None, - provider: str = None, - skills: List[Union[Skill, str]] = None, - params: dict = None, - **kwargs) -> Agent: + self, + name: str, + model: Union[Model, str, dict] = None, + provider: str = None, + skills: List[Union[Skill, str]] = None, + data: dict = None, + prompt_template: str = None, + params: dict = None, + **kwargs + ) -> Agent: """ Create new agent and return it :param name: Name of the agent to be created - :param model: Model to be used by the agent + :param model: Model to be used by the agent. This can be a Model object, a string with model name, or a dictionary with model parameters. :param skills: List of skills to be used by the agent. Currently only 'sql' is supported. + :param provider: Provider of the model, e.g. 'mindsdb', 'openai', etc. + :param data: Data to be used by the agent. This is usually a dictionary with 'tables' and/or 'knowledge_base' keys. :param params: Parameters for the agent :return: created agent object @@ -507,17 +503,27 @@ def create( params = {} params.update(kwargs) - if 'prompt_template' not in params: - params['prompt_template'] = _DEFAULT_LLM_PROMPT - - if model is None: - model = _DEFAULT_LLM_MODEL - elif isinstance(model, Model): - model = model.name + model_name = None + if isinstance(model_name, Model): + model_name = model_name.name provider = 'mindsdb' + model = None + elif isinstance(model, str): + model_name = model + model = None - data = self.api.create_agent(self.project.name, name, model, provider, skill_names, params) - return Agent.from_json(data, self) + agent = self.api.create_agent( + self.project.name, + name, + model_name, + provider, + skill_names, + data, + model, + prompt_template, + params + ) + return Agent.from_json(agent, self) def update(self, name: str, updated_agent: Agent): """ @@ -546,20 +552,36 @@ def update(self, name: str, updated_agent: Agent): updated_skills.add(skill.name) existing_agent = self.api.agent(self.project.name, name) - existing_skills = set([s['name'] for s in existing_agent['skills']]) + + existing_skills = set([s['name'] for s in existing_agent.get('skills', [])]) skills_to_add = updated_skills.difference(existing_skills) skills_to_remove = existing_skills.difference(updated_skills) - data = self.api.update_agent( + + updated_model_name = None + updated_provider = updated_agent.provider + updated_model = None + if isinstance(updated_agent.model, Model): + updated_model_name = updated_agent.model.name + updated_provider = 'mindsdb' + elif isinstance(updated_agent.model, str): + updated_model_name = updated_agent.model + elif isinstance(updated_agent.model, dict): + updated_model = updated_agent.model + + agent = self.api.update_agent( self.project.name, name, updated_agent.name, - updated_agent.provider, - updated_agent.model_name, + updated_provider, + updated_model_name, list(skills_to_add), list(skills_to_remove), + updated_agent.data, + updated_model, + updated_agent.prompt_template, updated_agent.params ) - return Agent.from_json(data, self) + return Agent.from_json(agent, self) def drop(self, name: str): """ diff --git a/mindsdb_sdk/connectors/rest_api.py b/mindsdb_sdk/connectors/rest_api.py index d7bf990..b31185d 100644 --- a/mindsdb_sdk/connectors/rest_api.py +++ b/mindsdb_sdk/connectors/rest_api.py @@ -306,16 +306,30 @@ def agent_completion_stream_v2(self, project: str, name: str, messages: List[dic yield e @_try_relogin - def create_agent(self, project: str, name: str, model: str = None, provider: str = None, skills: List[str] = None, params: dict = None): + def create_agent( + self, + project: str, + name: str, + model_name: str = None, + provider: str = None, + skills: List[str] = None, + data: dict = None, + model: dict = None, + prompt_template: str = None, + params: dict = None + ): url = self.url + f'/api/projects/{project}/agents' r = self.session.post( url, json={ 'agent': { 'name': name, - 'model_name': model, + 'model_name': model_name, 'provider': provider, 'skills': skills, + 'data': data, + 'model': model, + 'prompt_template': prompt_template, 'params': params } } @@ -325,26 +339,32 @@ def create_agent(self, project: str, name: str, model: str = None, provider: str @_try_relogin def update_agent( - self, - project: str, - name: str, - updated_name: str, - updated_provider: str, - updated_model: str, - skills_to_add: List[str], - skills_to_remove: List[str], - updated_params: dict - ): + self, + project: str, + name: str, + updated_name: str, + updated_provider: str, + updated_model_name: str, + skills_to_add: List[str], + skills_to_remove: List[str], + updated_data: dict, + updated_model: dict, + updated_prompt_template: str, + updated_params: dict + ): url = self.url + f'/api/projects/{project}/agents/{name}' r = self.session.put( url, json={ 'agent': { 'name': updated_name, - 'model_name': updated_model, + 'model_name': updated_model_name, 'provider': updated_provider, 'skills_to_add': skills_to_add, 'skills_to_remove': skills_to_remove, + 'data': updated_data, + 'model': updated_model, + 'prompt_template': updated_prompt_template, 'params': updated_params } } diff --git a/tests/test_sdk.py b/tests/test_sdk.py index af290e9..42dd588 100644 --- a/tests/test_sdk.py +++ b/tests/test_sdk.py @@ -1337,12 +1337,18 @@ def test_list(self, mock_get): 'id': 1, 'name': 'test_agent', 'project_id': 1, - 'model_name': 'test_model', - 'skills': [], + 'model': { + 'model_name': 'gpt-3.5-turbo', + 'provider': 'openai', + 'api_key': 'sk-...', + }, + 'data': { + 'tables': ['test_database.test_table'], + 'knowledge_bases': ['test_kb'], + }, 'params': {}, 'created_at': created_at, 'updated_at': updated_at, - 'provider': 'mindsdb' } ]) all_agents = server.agents.list() @@ -1352,12 +1358,17 @@ def test_list(self, mock_get): assert len(all_agents) == 1 expected_agent = Agent( 'test_agent', - 'test_model', - [], - {}, created_at, updated_at, - 'mindsdb' + model={ + 'model_name': 'gpt-3.5-turbo', + 'provider': 'openai', + 'api_key': 'sk-...', + }, + data={ + 'tables': ['test_database.test_table'], + 'knowledge_bases': ['test_kb'], + }, ) assert all_agents[0] == expected_agent @@ -1371,12 +1382,18 @@ def test_get(self, mock_get): 'id': 1, 'name': 'test_agent', 'project_id': 1, - 'model_name': 'test_model', - 'skills': [], + 'model': { + 'model_name': 'gpt-3.5-turbo', + 'provider': 'openai', + 'api_key': 'sk-...', + }, + 'data': { + 'tables': ['test_database.test_table'], + 'knowledge_bases': ['test_kb'], + }, 'params': {}, 'created_at': created_at, 'updated_at': updated_at, - 'provider': 'mindsdb' } ) agent = server.agents.get('test_agent') @@ -1384,53 +1401,57 @@ def test_get(self, mock_get): assert mock_get.call_args[0][0] == f'{DEFAULT_LOCAL_API_URL}/api/projects/mindsdb/agents/test_agent' expected_agent = Agent( 'test_agent', - 'test_model', - [], - {}, created_at, updated_at, - 'mindsdb' + model={ + 'model_name': 'gpt-3.5-turbo', + 'provider': 'openai', + 'api_key': 'sk-...', + }, + data={ + 'tables': ['test_database.test_table'], + 'knowledge_bases': ['test_kb'], + }, ) assert agent == expected_agent @patch('requests.Session.post') - @patch('requests.Session.get') - def test_create(self, mock_get, mock_post): + def test_create(self, mock_post): created_at = dt.datetime(2000, 3, 1, 9, 30) updated_at = dt.datetime(2001, 3, 1, 9, 30) data = { 'id': 1, 'name': 'test_agent', 'project_id': 1, - 'model_name': 'test_model', - 'skills': [{ - 'id': 0, - 'name': 'test_skill', - 'project_id': 1, - 'type': 'sql', - 'params': {'tables': ['test_table'], 'database': 'test_database', 'description': 'test_description'}, - }], - 'params': {'k1': 'v1'}, + 'model': { + 'model_name': 'gpt-3.5-turbo', + 'provider': 'openai', + 'api_key': 'sk-...', + }, + 'data': { + 'tables': ['test_database.test_table'], + 'knowledge_bases': ['test_kb'], + }, 'created_at': created_at, 'updated_at': updated_at, - 'provider': 'mindsdb', } responses_mock(mock_post, [ - # ML Engine get (SQL post for SHOW ML_ENGINES) data ]) - responses_mock(mock_get, [ - # Skill get. - {'name': 'test_skill', 'type': 'sql', 'params': {'tables': ['test_table'], 'database': 'test_database', 'description': 'test_description'}}, - ]) # Create the agent. server = mindsdb_sdk.connect() new_agent = server.agents.create( name='test_agent', - model=Model(None, {'name':'m1'}), - skills=['test_skill'], - params={'k1': 'v1'} + model={ + 'model_name': 'gpt-3.5-turbo', + 'provider': 'openai', + 'api_key': 'sk-...', + }, + data={ + 'tables': ['test_database.test_table'], + 'knowledge_bases': ['test_kb'], + } ) # Check API call. assert len(mock_post.call_args_list) == 1 @@ -1438,51 +1459,59 @@ def test_create(self, mock_get, mock_post): assert mock_post.call_args_list[-1][1]['json'] == { 'agent': { 'name': 'test_agent', - 'model_name': 'm1', - 'skills': ['test_skill'], - 'params': { - 'k1': 'v1', - 'prompt_template': 'Answer the user"s question in a helpful way: {{question}}' + 'model_name': None, + 'provider': None, + 'skills': [], + 'model': { + 'model_name': 'gpt-3.5-turbo', + 'provider': 'openai', + 'api_key': 'sk-...', + }, + 'data': { + 'tables': ['test_database.test_table'], + 'knowledge_bases': ['test_kb'], }, - 'provider': 'mindsdb' + 'prompt_template': None, + 'params': {} } } - expected_skill = SQLSkill('test_skill', ['test_table'], 'test_database', 'test_description') expected_agent = Agent( 'test_agent', - 'test_model', - [expected_skill], - {'k1': 'v1'}, - created_at, - updated_at, - 'mindsdb' + created_at=created_at, + updated_at=updated_at, + model={ + 'model_name': 'gpt-3.5-turbo', + 'provider': 'openai', + 'api_key': 'sk-...', + }, + data={ + 'tables': ['test_database.test_table'], + 'knowledge_bases': ['test_kb'], + } ) assert new_agent == expected_agent @patch('requests.Session.get') @patch('requests.Session.put') - # Mock creating new skills. - @patch('requests.Session.post') - def test_update(self, mock_get, mock_put, _): + def test_update(self, mock_put, mock_get): created_at = dt.datetime(2000, 3, 1, 9, 30) updated_at = dt.datetime(2001, 3, 1, 9, 30) data = { 'id': 1, 'name': 'test_agent', 'project_id': 1, - 'model_name': 'updated_model', - 'skills': [{ - 'id': 1, - 'name': 'updated_skill', - 'project_id': 1, - 'type': 'sql', - 'params': {'tables': ['updated_table'], 'database': 'updated_database', 'description': 'test_description'}, - }], - 'params': {'k2': 'v2'}, + 'model': { + 'model_name': 'gpt-3.5-turbo', + 'provider': 'openai', + 'api_key': 'sk-...', + }, + 'data': { + 'tables': ['test_database.test_table'], + 'knowledge_bases': ['test_kb'], + }, 'created_at': created_at, 'updated_at': updated_at, - 'provider': 'mindsdb', } response_mock(mock_put, data) @@ -1491,21 +1520,33 @@ def test_update(self, mock_get, mock_put, _): 'id': 1, 'name': 'test_agent', 'project_id': 1, - 'model_name': 'test_model', - 'skills': [], - 'params': {'k1': 'v1'}, - 'provider': 'mindsdb', + 'model': { + 'model_name': 'gpt-3.5-turbo', + 'provider': 'openai', + 'api_key': 'sk-...', + }, + 'data': { + 'tables': ['test_database.test_table'], + 'knowledge_bases': ['test_kb'], + }, + 'created_at': created_at, + 'updated_at': updated_at, }) server = mindsdb_sdk.connect() expected_agent = Agent( 'test_agent', - 'updated_model', - [SQLSkill('updated_skill', ['updated_table'], 'updated_database', 'test_description')], - {'k2': 'v2'}, created_at, updated_at, - 'mindsdb' + model={ + 'model_name': 'gpt-3.5-turbo', + 'provider': 'openai', + 'api_key': 'sk-...', + }, + data={ + 'tables': ['test_database.test_table'], + 'knowledge_bases': ['test_kb', 'test_kb2'], + }, ) updated_agent = server.agents.update('test_agent', expected_agent) @@ -1514,11 +1555,21 @@ def test_update(self, mock_get, mock_put, _): assert mock_put.call_args[1]['json'] == { 'agent': { 'name': 'test_agent', - 'model_name': 'updated_model', - 'skills_to_add': ['updated_skill'], + 'model_name': None, + 'provider': None, + 'skills_to_add': [], 'skills_to_remove': [], - 'params': {'k2': 'v2'}, - 'provider': 'mindsdb' + 'data': { + 'tables': ['test_database.test_table'], + 'knowledge_bases': ['test_kb', 'test_kb2'], + }, + 'model': { + 'model_name': 'gpt-3.5-turbo', + 'provider': 'openai', + 'api_key': 'sk-...', + }, + 'prompt_template': None, + 'params': {}, } } @@ -1555,154 +1606,186 @@ def test_delete(self, mock_delete): @patch('requests.Session.get') @patch('requests.Session.put') - @patch('requests.Session.post') - def test_add_file(self, mock_post, mock_put, mock_get): + def test_add_file(self, mock_put, mock_get): server = mindsdb_sdk.connect() responses_mock(mock_get, [ - # File metadata get. - [{'name': 'tokaido_rules'}], # Existing agent get. { 'name': 'test_agent', - 'model_name': 'test_model', - 'skills': [], - 'params': {}, + 'project_id': 1, + 'model': { + 'model_name': 'gpt-3.5-turbo', + 'provider': 'openai', + 'api_key': 'sk-...', + }, + 'data': { + 'tables': ['test_database.test_table'], + 'knowledge_bases': ['test_kb'], + }, 'created_at': None, 'updated_at': None, - 'provider': 'mindsdb' }, - # get KB - { - 'id': 1, - 'name': 'my_kb', - 'project_id': 1, - 'embedding_model': 'openai_emb', - 'vector_database': 'pvec', - 'vector_database_table': 'tbl1', - 'updated_at': '2024-10-04 10:55:25.350799', - 'created_at': '2024-10-04 10:55:25.350790', - 'params': {} - }, - # Skills get in Agent update to check if it exists. - {'name': 'new_skill', 'type': 'retrieval', 'params': {'source': 'test_agent_tokaido_rules_kb'}}, + # File metadata get. + [{'name': 'tokaido_rules'}], # Existing agent get in Agent update. { 'name': 'test_agent', - 'model_name': 'test_model', - 'skills': [], - 'params': {}, + 'project_id': 1, + 'model': { + 'model_name': 'gpt-3.5-turbo', + 'provider': 'openai', + 'api_key': 'sk-...', + }, + 'data': { + 'tables': ['test_database.test_table'], + 'knowledge_bases': ['test_kb'], + }, 'created_at': None, 'updated_at': None, - 'provider': 'mindsdb' }, ]) - responses_mock(mock_post, [ - # Skill creation. - {'name': 'new_skill', 'type': 'retrieval', 'params': {'source': 'test_agent_tokaido_rules_kb'}} - ]) responses_mock(mock_put, [ - # KB update. - {'name': 'test_agent_tokaido_rules_kb'}, - # Agent update with new skill. + # Agent update with new data. { 'name': 'test_agent', - 'model_name': 'test_model', - 'skills': [{'name': 'new_skill', 'type': 'retrieval', 'params': {'source': 'test_agent_tokaido_rules_kb'}}], - 'params': {}, + 'project_id': 1, + 'model': { + 'model_name': 'gpt-3.5-turbo', + 'provider': 'openai', + 'api_key': 'sk-...', + }, + 'data': { + 'tables': ['test_database.test_table', 'files.tokaido_rules'], + 'knowledge_bases': ['test_kb'], + }, + 'prompt_template': '\nRules for the board game Tokaido', 'created_at': None, 'updated_at': None, - 'provider': 'mindsdb' }, ]) - server.agents.add_file('test_agent', './tokaido_rules.pdf', 'Rules for the board game Tokaido', 'existing_kb') + server.agents.add_file('test_agent', './tokaido_rules.pdf', 'Rules for the board game Tokaido') - # Check Agent was updated with a new skill. + # Check Agent was updated with a new data. agent_update_json = mock_put.call_args[-1]['json'] expected_agent_json = { 'agent': { 'name': 'test_agent', - 'model_name': 'test_model', - # Skill name is a generated UUID. - 'skills_to_add': [agent_update_json['agent']['skills_to_add'][0]], - 'skills_to_remove': [], + 'model_name': None, 'params': {}, - 'provider': 'mindsdb' + 'provider': None, + 'prompt_template': '\nRules for the board game Tokaido', + 'model': { + 'model_name': 'gpt-3.5-turbo', + 'provider': 'openai', + 'api_key': 'sk-...', + }, + 'skills_to_add': [], + 'skills_to_remove': [], + 'data': { + 'tables': ['test_database.test_table', 'files.tokaido_rules'], + 'knowledge_bases': ['test_kb'], + }, } } assert agent_update_json == expected_agent_json @patch('requests.Session.get') @patch('requests.Session.put') - @patch('requests.Session.post') - def test_add_webpage(self, mock_post, mock_put, mock_get): + def test_add_webpage(self, mock_put, mock_get): server = mindsdb_sdk.connect() responses_mock(mock_get, [ # Existing agent get. { - 'name':'test_agent', - 'model_name':'test_model', - 'skills':[], - 'params':{}, - 'created_at':None, - 'updated_at':None, - 'provider':'mindsdb' + 'name': 'test_agent', + 'project_id': 1, + 'model': { + 'model_name': 'gpt-3.5-turbo', + 'provider': 'openai', + 'api_key': 'sk-...', + }, + 'data': { + 'tables': ['test_database.test_table'], + 'knowledge_bases': ['test_kb'], + }, + 'created_at': None, + 'updated_at': None, }, # get KB { 'id': 1, - 'name': 'my_kb', + 'name': 'existing_kb', 'project_id': 1, - 'embedding_model': 'openai_emb', + 'embedding_model': { + 'provider': 'openai', + 'model_name': 'openai_emb', + 'api_key': 'sk-...' + }, 'vector_database': 'pvec', 'vector_database_table': 'tbl1', 'updated_at': '2024-10-04 10:55:25.350799', 'created_at': '2024-10-04 10:55:25.350790', 'params': {} }, - # Skills get in Agent update to check if it exists. - {'name':'new_skill', 'type':'retrieval', 'params':{'source':'test_agent_docs_mdb_ai_kb'}}, # Existing agent get in Agent update. { - 'name':'test_agent', - 'model_name':'test_model', - 'skills':[], - 'params':{}, - 'created_at':None, - 'updated_at':None, - 'provider':'mindsdb' # Added provider field + 'name': 'test_agent', + 'project_id': 1, + 'model': { + 'model_name': 'gpt-3.5-turbo', + 'provider': 'openai', + 'api_key': 'sk-...', + }, + 'data': { + 'tables': ['test_database.test_table'], + 'knowledge_bases': ['test_kb'], + }, + 'created_at': None, + 'updated_at': None, }, ]) - responses_mock(mock_post, [ - # Skill creation. - {'name':'new_skill', 'type':'retrieval', 'params':{'source':'test_agent_docs_mdb_ai_kb'}} - ]) responses_mock(mock_put, [ # KB update. {'name':'test_agent_docs_mdb_ai_kb'}, - # Agent update with new skill. + # Agent update with new data. { - 'name':'test_agent', - 'model_name':'test_model', - 'skills':[{'name':'new_skill', 'type':'retrieval', 'params':{'source':'test_agent_docs_mdb_ai_kb'}}], - 'params':{}, - 'created_at':None, - 'updated_at':None, - 'provider':'mindsdb' # Added provider field + 'name': 'test_agent', + 'project_id': 1, + 'model': { + 'model_name': 'gpt-3.5-turbo', + 'provider': 'openai', + 'api_key': 'sk-...', + }, + 'data': { + 'tables': ['test_database.test_table'], + 'knowledge_bases': ['test_kb', 'existing_kb'] + }, + 'prompt_template': 'Documentation for MindsDB', + 'created_at': None, + 'updated_at': None, }, ]) server.agents.add_webpage('test_agent', 'docs.mdb.ai', 'Documentation for MindsDB', 'existing_kb') - # Check Agent was updated with a new skill. + # Check Agent was updated with a new data. agent_update_json = mock_put.call_args[-1]['json'] expected_agent_json = { 'agent':{ - 'name':'test_agent', - 'model_name':'test_model', - # Skill name is a generated UUID. - 'skills_to_add':[agent_update_json['agent']['skills_to_add'][0]], - 'skills_to_remove':[], - 'params':{}, - 'provider': 'mindsdb' + 'name': 'test_agent', + 'model_name': None, + 'params': {}, + 'provider': None, + 'prompt_template': '\nDocumentation for MindsDB', + 'model': { + 'model_name': 'gpt-3.5-turbo', + 'provider': 'openai', + 'api_key': 'sk-...', + }, + 'skills_to_add': [], + 'skills_to_remove': [], + 'data': { + 'tables': ['test_database.test_table'], + 'knowledge_bases': ['test_kb', 'existing_kb'] + }, } } assert agent_update_json == expected_agent_json @@ -1716,24 +1799,34 @@ def test_add_database(self, mock_post, mock_put, mock_get): # Existing agent get. { 'name': 'test_agent', - 'model_name': 'test_model', - 'skills': [], - 'params': {}, + 'project_id': 1, + 'model': { + 'model_name': 'gpt-3.5-turbo', + 'provider': 'openai', + 'api_key': 'sk-...', + }, + 'data': { + 'tables': ['test_database.test_table'], + 'knowledge_bases': ['test_kb'], + }, 'created_at': None, 'updated_at': None, - 'provider': 'mindsdb' }, - # Skills get in Agent update to check if it exists. - {'name': 'new_skill', 'type': 'sql', 'params': {'database': 'existing_db', 'tables': ['existing_table']}}, # Existing agent get in Agent update. { 'name': 'test_agent', - 'model_name': 'test_model', - 'skills': [], - 'params': {}, + 'project_id': 1, + 'model': { + 'model_name': 'gpt-3.5-turbo', + 'provider': 'openai', + 'api_key': 'sk-...', + }, + 'data': { + 'tables': ['test_database.test_table'], + 'knowledge_bases': ['test_kb'], + }, 'created_at': None, 'updated_at': None, - 'provider': 'mindsdb' }, ]) responses_mock( @@ -1751,39 +1844,48 @@ def test_add_database(self, mock_post, mock_put, mock_get): ), # DB tables get (POST /sql). pd.DataFrame([{"name": "existing_table"}]), - # Skill creation. - { - "name": "new_skill", - "type": "sql", - "params": {"database": "existing_db", "tables": ["existing_table"]}, - }, ], ) responses_mock(mock_put, [ - # Agent update with new skill. + # Agent update with new data. { 'name': 'test_agent', - 'model_name': 'test_model', - 'skills': [{'name': 'new_skill', 'type': 'sql', 'params': {'database': 'existing_db', 'tables': ['existing_table']}}], - 'params': {}, + 'project_id': 1, + 'model': { + 'model_name': 'gpt-3.5-turbo', + 'provider': 'openai', + 'api_key': 'sk-...', + }, + 'data': { + 'tables': ['test_database.test_table', 'existing_db.existing_table'], + 'knowledge_bases': ['test_kb'], + }, 'created_at': None, 'updated_at': None, - 'provider': 'mindsdb' }, ]) - server.agents.add_database('test_agent', 'existing_db', ['existing_table'], 'My data') + server.agents.add_database('test_agent', 'existing_db', ['existing_table']) - # Check Agent was updated with a new skill. + # Check Agent was updated with a new data. agent_update_json = mock_put.call_args[-1]['json'] expected_agent_json = { 'agent': { 'name': 'test_agent', - 'model_name': 'test_model', - # Skill name is a generated UUID. - 'skills_to_add': [agent_update_json['agent']['skills_to_add'][0]], + 'model_name': None, + 'params': {}, + 'provider': None, + 'prompt_template': None, + 'model': { + 'model_name': 'gpt-3.5-turbo', + 'provider': 'openai', + 'api_key': 'sk-...', + }, + 'skills_to_add': [], 'skills_to_remove': [], - 'params': {'prompt_template': 'using mindsdb sqltoolbox'}, - 'provider': 'mindsdb' + 'data': { + 'tables': ['test_database.test_table', 'existing_db.existing_table'], + 'knowledge_bases': ['test_kb'], + }, } } assert agent_update_json == expected_agent_json