Skip to content

Commit e6dce77

Browse files
MinuraPunchihewaMinuraPunchihewa
committed
fixed unit tests for adding files, webpages and databases
1 parent 9c67be9 commit e6dce77

File tree

1 file changed

+163
-112
lines changed

1 file changed

+163
-112
lines changed

‎tests/test_sdk.py‎

Lines changed: 163 additions & 112 deletions
Original file line numberDiff line numberDiff line change
@@ -1606,154 +1606,186 @@ def test_delete(self, mock_delete):
16061606

16071607
@patch('requests.Session.get')
16081608
@patch('requests.Session.put')
1609-
@patch('requests.Session.post')
1610-
def test_add_file(self, mock_post, mock_put, mock_get):
1609+
def test_add_file(self, mock_put, mock_get):
16111610
server = mindsdb_sdk.connect()
16121611
responses_mock(mock_get, [
1613-
# File metadata get.
1614-
[{'name': 'tokaido_rules'}],
16151612
# Existing agent get.
16161613
{
16171614
'name': 'test_agent',
1618-
'model_name': 'test_model',
1619-
'skills': [],
1620-
'params': {},
1615+
'project_id': 1,
1616+
'model': {
1617+
'model_name': 'gpt-3.5-turbo',
1618+
'provider': 'openai',
1619+
'api_key': 'sk-...',
1620+
},
1621+
'data': {
1622+
'tables': ['test_database.test_table'],
1623+
'knowledge_bases': ['test_kb'],
1624+
},
16211625
'created_at': None,
16221626
'updated_at': None,
1623-
'provider': 'mindsdb'
16241627
},
1625-
# get KB
1626-
{
1627-
'id': 1,
1628-
'name': 'my_kb',
1629-
'project_id': 1,
1630-
'embedding_model': 'openai_emb',
1631-
'vector_database': 'pvec',
1632-
'vector_database_table': 'tbl1',
1633-
'updated_at': '2024-10-04 10:55:25.350799',
1634-
'created_at': '2024-10-04 10:55:25.350790',
1635-
'params': {}
1636-
},
1637-
# Skills get in Agent update to check if it exists.
1638-
{'name': 'new_skill', 'type': 'retrieval', 'params': {'source': 'test_agent_tokaido_rules_kb'}},
1628+
# File metadata get.
1629+
[{'name': 'tokaido_rules'}],
16391630
# Existing agent get in Agent update.
16401631
{
16411632
'name': 'test_agent',
1642-
'model_name': 'test_model',
1643-
'skills': [],
1644-
'params': {},
1633+
'project_id': 1,
1634+
'model': {
1635+
'model_name': 'gpt-3.5-turbo',
1636+
'provider': 'openai',
1637+
'api_key': 'sk-...',
1638+
},
1639+
'data': {
1640+
'tables': ['test_database.test_table'],
1641+
'knowledge_bases': ['test_kb'],
1642+
},
16451643
'created_at': None,
16461644
'updated_at': None,
1647-
'provider': 'mindsdb'
16481645
},
16491646
])
1650-
responses_mock(mock_post, [
1651-
# Skill creation.
1652-
{'name': 'new_skill', 'type': 'retrieval', 'params': {'source': 'test_agent_tokaido_rules_kb'}}
1653-
])
16541647
responses_mock(mock_put, [
1655-
# KB update.
1656-
{'name': 'test_agent_tokaido_rules_kb'},
1657-
# Agent update with new skill.
1648+
# Agent update with new data.
16581649
{
16591650
'name': 'test_agent',
1660-
'model_name': 'test_model',
1661-
'skills': [{'name': 'new_skill', 'type': 'retrieval', 'params': {'source': 'test_agent_tokaido_rules_kb'}}],
1662-
'params': {},
1651+
'project_id': 1,
1652+
'model': {
1653+
'model_name': 'gpt-3.5-turbo',
1654+
'provider': 'openai',
1655+
'api_key': 'sk-...',
1656+
},
1657+
'data': {
1658+
'tables': ['test_database.test_table', 'files.tokaido_rules'],
1659+
'knowledge_bases': ['test_kb'],
1660+
},
1661+
'prompt_template': '\nRules for the board game Tokaido',
16631662
'created_at': None,
16641663
'updated_at': None,
1665-
'provider': 'mindsdb'
16661664
},
16671665
])
1668-
server.agents.add_file('test_agent', './tokaido_rules.pdf', 'Rules for the board game Tokaido', 'existing_kb')
1666+
server.agents.add_file('test_agent', './tokaido_rules.pdf', 'Rules for the board game Tokaido')
16691667

1670-
# Check Agent was updated with a new skill.
1668+
# Check Agent was updated with a new data.
16711669
agent_update_json = mock_put.call_args[-1]['json']
16721670
expected_agent_json = {
16731671
'agent': {
16741672
'name': 'test_agent',
1675-
'model_name': 'test_model',
1676-
# Skill name is a generated UUID.
1677-
'skills_to_add': [agent_update_json['agent']['skills_to_add'][0]],
1678-
'skills_to_remove': [],
1673+
'model_name': None,
16791674
'params': {},
1680-
'provider': 'mindsdb'
1675+
'provider': None,
1676+
'prompt_template': '\nRules for the board game Tokaido',
1677+
'model': {
1678+
'model_name': 'gpt-3.5-turbo',
1679+
'provider': 'openai',
1680+
'api_key': 'sk-...',
1681+
},
1682+
'skills_to_add': [],
1683+
'skills_to_remove': [],
1684+
'data': {
1685+
'tables': ['test_database.test_table', 'files.tokaido_rules'],
1686+
'knowledge_bases': ['test_kb'],
1687+
},
16811688
}
16821689
}
16831690
assert agent_update_json == expected_agent_json
16841691

16851692
@patch('requests.Session.get')
16861693
@patch('requests.Session.put')
1687-
@patch('requests.Session.post')
1688-
def test_add_webpage(self, mock_post, mock_put, mock_get):
1694+
def test_add_webpage(self, mock_put, mock_get):
16891695
server = mindsdb_sdk.connect()
16901696
responses_mock(mock_get, [
16911697
# Existing agent get.
16921698
{
1693-
'name':'test_agent',
1694-
'model_name':'test_model',
1695-
'skills':[],
1696-
'params':{},
1697-
'created_at':None,
1698-
'updated_at':None,
1699-
'provider':'mindsdb'
1699+
'name': 'test_agent',
1700+
'project_id': 1,
1701+
'model': {
1702+
'model_name': 'gpt-3.5-turbo',
1703+
'provider': 'openai',
1704+
'api_key': 'sk-...',
1705+
},
1706+
'data': {
1707+
'tables': ['test_database.test_table'],
1708+
'knowledge_bases': ['test_kb'],
1709+
},
1710+
'created_at': None,
1711+
'updated_at': None,
17001712
},
17011713
# get KB
17021714
{
17031715
'id': 1,
1704-
'name': 'my_kb',
1716+
'name': 'existing_kb',
17051717
'project_id': 1,
1706-
'embedding_model': 'openai_emb',
1718+
'embedding_model': {
1719+
'provider': 'openai',
1720+
'model_name': 'openai_emb',
1721+
'api_key': 'sk-...'
1722+
},
17071723
'vector_database': 'pvec',
17081724
'vector_database_table': 'tbl1',
17091725
'updated_at': '2024-10-04 10:55:25.350799',
17101726
'created_at': '2024-10-04 10:55:25.350790',
17111727
'params': {}
17121728
},
1713-
# Skills get in Agent update to check if it exists.
1714-
{'name':'new_skill', 'type':'retrieval', 'params':{'source':'test_agent_docs_mdb_ai_kb'}},
17151729
# Existing agent get in Agent update.
17161730
{
1717-
'name':'test_agent',
1718-
'model_name':'test_model',
1719-
'skills':[],
1720-
'params':{},
1721-
'created_at':None,
1722-
'updated_at':None,
1723-
'provider':'mindsdb' # Added provider field
1731+
'name': 'test_agent',
1732+
'project_id': 1,
1733+
'model': {
1734+
'model_name': 'gpt-3.5-turbo',
1735+
'provider': 'openai',
1736+
'api_key': 'sk-...',
1737+
},
1738+
'data': {
1739+
'tables': ['test_database.test_table'],
1740+
'knowledge_bases': ['test_kb'],
1741+
},
1742+
'created_at': None,
1743+
'updated_at': None,
17241744
},
17251745
])
1726-
responses_mock(mock_post, [
1727-
# Skill creation.
1728-
{'name':'new_skill', 'type':'retrieval', 'params':{'source':'test_agent_docs_mdb_ai_kb'}}
1729-
])
17301746
responses_mock(mock_put, [
17311747
# KB update.
17321748
{'name':'test_agent_docs_mdb_ai_kb'},
1733-
# Agent update with new skill.
1749+
# Agent update with new data.
17341750
{
1735-
'name':'test_agent',
1736-
'model_name':'test_model',
1737-
'skills':[{'name':'new_skill', 'type':'retrieval', 'params':{'source':'test_agent_docs_mdb_ai_kb'}}],
1738-
'params':{},
1739-
'created_at':None,
1740-
'updated_at':None,
1741-
'provider':'mindsdb' # Added provider field
1751+
'name': 'test_agent',
1752+
'project_id': 1,
1753+
'model': {
1754+
'model_name': 'gpt-3.5-turbo',
1755+
'provider': 'openai',
1756+
'api_key': 'sk-...',
1757+
},
1758+
'data': {
1759+
'tables': ['test_database.test_table'],
1760+
'knowledge_bases': ['test_kb', 'existing_kb']
1761+
},
1762+
'prompt_template': 'Documentation for MindsDB',
1763+
'created_at': None,
1764+
'updated_at': None,
17421765
},
17431766
])
17441767
server.agents.add_webpage('test_agent', 'docs.mdb.ai', 'Documentation for MindsDB', 'existing_kb')
17451768

1746-
# Check Agent was updated with a new skill.
1769+
# Check Agent was updated with a new data.
17471770
agent_update_json = mock_put.call_args[-1]['json']
17481771
expected_agent_json = {
17491772
'agent':{
1750-
'name':'test_agent',
1751-
'model_name':'test_model',
1752-
# Skill name is a generated UUID.
1753-
'skills_to_add':[agent_update_json['agent']['skills_to_add'][0]],
1754-
'skills_to_remove':[],
1755-
'params':{},
1756-
'provider': 'mindsdb'
1773+
'name': 'test_agent',
1774+
'model_name': None,
1775+
'params': {},
1776+
'provider': None,
1777+
'prompt_template': '\nDocumentation for MindsDB',
1778+
'model': {
1779+
'model_name': 'gpt-3.5-turbo',
1780+
'provider': 'openai',
1781+
'api_key': 'sk-...',
1782+
},
1783+
'skills_to_add': [],
1784+
'skills_to_remove': [],
1785+
'data': {
1786+
'tables': ['test_database.test_table'],
1787+
'knowledge_bases': ['test_kb', 'existing_kb']
1788+
},
17571789
}
17581790
}
17591791
assert agent_update_json == expected_agent_json
@@ -1767,24 +1799,34 @@ def test_add_database(self, mock_post, mock_put, mock_get):
17671799
# Existing agent get.
17681800
{
17691801
'name': 'test_agent',
1770-
'model_name': 'test_model',
1771-
'skills': [],
1772-
'params': {},
1802+
'project_id': 1,
1803+
'model': {
1804+
'model_name': 'gpt-3.5-turbo',
1805+
'provider': 'openai',
1806+
'api_key': 'sk-...',
1807+
},
1808+
'data': {
1809+
'tables': ['test_database.test_table'],
1810+
'knowledge_bases': ['test_kb'],
1811+
},
17731812
'created_at': None,
17741813
'updated_at': None,
1775-
'provider': 'mindsdb'
17761814
},
1777-
# Skills get in Agent update to check if it exists.
1778-
{'name': 'new_skill', 'type': 'sql', 'params': {'database': 'existing_db', 'tables': ['existing_table']}},
17791815
# Existing agent get in Agent update.
17801816
{
17811817
'name': 'test_agent',
1782-
'model_name': 'test_model',
1783-
'skills': [],
1784-
'params': {},
1818+
'project_id': 1,
1819+
'model': {
1820+
'model_name': 'gpt-3.5-turbo',
1821+
'provider': 'openai',
1822+
'api_key': 'sk-...',
1823+
},
1824+
'data': {
1825+
'tables': ['test_database.test_table'],
1826+
'knowledge_bases': ['test_kb'],
1827+
},
17851828
'created_at': None,
17861829
'updated_at': None,
1787-
'provider': 'mindsdb'
17881830
},
17891831
])
17901832
responses_mock(
@@ -1802,39 +1844,48 @@ def test_add_database(self, mock_post, mock_put, mock_get):
18021844
),
18031845
# DB tables get (POST /sql).
18041846
pd.DataFrame([{"name": "existing_table"}]),
1805-
# Skill creation.
1806-
{
1807-
"name": "new_skill",
1808-
"type": "sql",
1809-
"params": {"database": "existing_db", "tables": ["existing_table"]},
1810-
},
18111847
],
18121848
)
18131849
responses_mock(mock_put, [
1814-
# Agent update with new skill.
1850+
# Agent update with new data.
18151851
{
18161852
'name': 'test_agent',
1817-
'model_name': 'test_model',
1818-
'skills': [{'name': 'new_skill', 'type': 'sql', 'params': {'database': 'existing_db', 'tables': ['existing_table']}}],
1819-
'params': {},
1853+
'project_id': 1,
1854+
'model': {
1855+
'model_name': 'gpt-3.5-turbo',
1856+
'provider': 'openai',
1857+
'api_key': 'sk-...',
1858+
},
1859+
'data': {
1860+
'tables': ['test_database.test_table', 'existing_db.existing_table'],
1861+
'knowledge_bases': ['test_kb'],
1862+
},
18201863
'created_at': None,
18211864
'updated_at': None,
1822-
'provider': 'mindsdb'
18231865
},
18241866
])
1825-
server.agents.add_database('test_agent', 'existing_db', ['existing_table'], 'My data')
1867+
server.agents.add_database('test_agent', 'existing_db', ['existing_table'])
18261868

1827-
# Check Agent was updated with a new skill.
1869+
# Check Agent was updated with a new data.
18281870
agent_update_json = mock_put.call_args[-1]['json']
18291871
expected_agent_json = {
18301872
'agent': {
18311873
'name': 'test_agent',
1832-
'model_name': 'test_model',
1833-
# Skill name is a generated UUID.
1834-
'skills_to_add': [agent_update_json['agent']['skills_to_add'][0]],
1874+
'model_name': None,
1875+
'params': {},
1876+
'provider': None,
1877+
'prompt_template': None,
1878+
'model': {
1879+
'model_name': 'gpt-3.5-turbo',
1880+
'provider': 'openai',
1881+
'api_key': 'sk-...',
1882+
},
1883+
'skills_to_add': [],
18351884
'skills_to_remove': [],
1836-
'params': {'prompt_template': 'using mindsdb sqltoolbox'},
1837-
'provider': 'mindsdb'
1885+
'data': {
1886+
'tables': ['test_database.test_table', 'existing_db.existing_table'],
1887+
'knowledge_bases': ['test_kb'],
1888+
},
18381889
}
18391890
}
18401891
assert agent_update_json == expected_agent_json

0 commit comments

Comments
 (0)