diff --git a/examples/using_agents_with_text2sql.py b/examples/using_agents_with_text2sql.py new file mode 100644 index 0000000..fde3184 --- /dev/null +++ b/examples/using_agents_with_text2sql.py @@ -0,0 +1,41 @@ +import mindsdb_sdk +from uuid import uuid4 +import os + +con = mindsdb_sdk.connect() + +open_ai_key = os.getenv('OPENAI_API_KEY') +model_name = 'gpt-4' + +# Now create an agent that will use the model we just created. +agent = con.agents.create(name=f'mindsdb_sql_agent_{model_name}_{uuid4().hex}', + model={'model_name': model_name, 'openai_api_key': open_ai_key}, + params={'openai_api_key': open_ai_key}) + +# Set up a Postgres data source with our new agent. +data_source = 'postgres' +connection_args = { + "user": "demo_user", + "password": "demo_password", + "host": "samples.mindsdb.com", + "port": "5432", + "database": "demo", + "schema": "demo_data" +} +description = 'mindsdb demo database' +database = con.databases.create( + f'mindsdb_sql_agent_datasource_{uuid4().hex}', + data_source, + connection_args +) + +# Actually connect the agent to the datasource. +agent.add_database(database.name, [], description) + + +question = 'How many three-bedroom houses were sold in 2008?' +answer = agent.completion([{'question': question, 'answer': None}]) +print(answer.content) + +con.databases.drop(database.name) +con.agents.drop(agent.name) diff --git a/examples/using_database_mind_text2sql.py b/examples/using_database_mind_text2sql.py new file mode 100644 index 0000000..5e86830 --- /dev/null +++ b/examples/using_database_mind_text2sql.py @@ -0,0 +1,51 @@ +from uuid import uuid4 + +from openai import OpenAI +from mindsdb_sdk.utils.mind import create_mind +import os + + +# Load MindsDB API key from environment variable. or set it here. +MINDSDB_API_KEY = os.getenv('MINDSDB_API_KEY') + +# Set the model name for mind to use +model_name = 'gpt-4' + +# Set the base URL for the MindsDB LiteLLM proxy. +base_url = 'https://ai.dev.mindsdb.com' + + +# Connect to MindsDB LiteLLM proxy. +client = OpenAI( + api_key=MINDSDB_API_KEY, + base_url=base_url +) + +# create a database mind +mind = create_mind( + name = f'my_house_data_mind_{uuid4().hex}', + description= 'House Sales', + base_url= base_url, + api_key= MINDSDB_API_KEY, + model= model_name, + data_source_type='postgres', + data_source_connection_args={ + 'user': 'demo_user', + 'password': 'demo_password', + 'host': 'samples.mindsdb.com', + 'port': '5432', + 'database': 'demo', + 'schema': 'demo_data' + } +) + +# Actually pass in our tool to get a SQL completion. +completion = client.chat.completions.create( + model=mind.name, + messages=[ + {'role': 'user', 'content': 'How many 2 bedroom houses sold in 2008?'} + ], + stream=False +) + +print(completion.choices[0].message.content) diff --git a/examples/using_mindsdb_inference_with_text2sql_prompt.py b/examples/using_mindsdb_inference_with_text2sql_prompt.py deleted file mode 100644 index 8f05550..0000000 --- a/examples/using_mindsdb_inference_with_text2sql_prompt.py +++ /dev/null @@ -1,95 +0,0 @@ -from openai import OpenAI -from mindsdb_sdk.utils.openai import extract_sql_query, query_database, chat_completion_request, \ - pretty_print_conversation - -import mindsdb_sdk -import os - -from mindsdb_sdk.utils.table_schema import get_table_schemas - -# generate the key at https://llm.mdb.ai -MINDSDB_API_KEY = os.environ.get("MINDSDB_API_KEY") -OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY") - -MODEL = "gpt-3.5-turbo" - -# the prompt should be a question that can be answered by the database -SYSTEM_PROMPT = """You are a SQL expert. Given an input question, first create a syntactically correct SQL query to run, then look at the results of the query and return the answer to the input question. -Unless the user specifies in the question a specific number of examples to obtain, query for at most 5 results using the LIMIT clause as per SQL standards. You can order the results to return the most informative data in the database. -Never query for all columns from a table. You must query only the columns that are needed to answer the question. Wrap each column name in backticks (`) to denote them as identifiers. -Pay attention to use only the column names you can see in the tables below. Be careful to not query for columns that do not exist. Also, pay attention to which column is in which table. -Pay attention to use CURRENT_DATE function to get the current date, if the question involves "today". - -Use the following format: - -Question: -SQLQuery: -SQLResult: -Answer: - -Only use the following tables: - -{schema} -""" -PROMPT = "what was the average delay on arrivals?" - - -def generate_system_prompt(system_prompt: str, schema: dict) -> dict: - prompt = { - "role":"system", - "content":system_prompt.format(schema=schema) - } - return prompt - - -def generate_user_prompt(query: str) -> dict: - prompt = { - "role":"user", - "content":query - } - return prompt - - -con = mindsdb_sdk.connect() - -# given database name, returns schema and database object -# using example_db from mindsdb - -database = con.databases.get("example_db") -schema = get_table_schemas(database, included_tables=["airline_passenger_satisfaction"]) - -# client_mindsdb_serve = OpenAI( -# api_key=MINDSDB_API_KEY, -# base_url="https://llm.mdb.ai" -# ) - -client_mindsdb_serve = OpenAI( - api_key=OPENAI_API_KEY -) - -messages = [ - generate_system_prompt(SYSTEM_PROMPT, schema), - generate_user_prompt(PROMPT) -] - -chat_response = chat_completion_request(client=client_mindsdb_serve, model=MODEL, messages=messages, tools=None, - tool_choice=None) - -# extract the SQL query from the response -query = extract_sql_query(chat_response.choices[0].message.content) - -result = query_database(database, query) - -# generate the user prompt with the query result, this will be used to generate the final response -query = generate_user_prompt(f"Given this SQLResult: {str(result)} provide Answer: ") - -# add the query to the messages list -messages.append(query) - -# generate the final response -chat_completion_gpt = chat_completion_request(client=client_mindsdb_serve, model=MODEL, messages=messages, tools=None, - tool_choice=None) -response = chat_completion_gpt.choices[0].message.content - -pretty_print_conversation(messages) - diff --git a/examples/using_mindsdb_inference_with_text2sql_using_tools.py b/examples/using_mindsdb_inference_with_text2sql_using_tools.py deleted file mode 100644 index 02c8a23..0000000 --- a/examples/using_mindsdb_inference_with_text2sql_using_tools.py +++ /dev/null @@ -1,71 +0,0 @@ -from openai import OpenAI - - -from mindsdb_sdk.utils.openai import ( - make_query_tool, - execute_function_call, - chat_completion_request, - pretty_print_conversation) - -import mindsdb_sdk -import os - -from mindsdb_sdk.utils.table_schema import get_table_schemas - -# generate the key at https://llm.mdb.ai -MINDSDB_API_KEY = os.environ.get("MINDSDB_API_KEY") -OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY") - -MODEL = "gpt-3.5-turbo" - - -con = mindsdb_sdk.connect() - -# given database name, returns schema and database object -# using example_db from mindsdb - -# client_mindsdb_serve = OpenAI( -# api_key=MINDSDB_API_KEY, -# base_url="https://llm.mdb.ai" -# ) - -client_mindsdb_serve = OpenAI( - api_key=OPENAI_API_KEY -) - -database = con.databases.get("mindsdb_demo_db") -schema = get_table_schemas(database, included_tables=["airline_passenger_satisfaction"]) - -tools = [make_query_tool(schema)] - -SYSTEM_PROMPT = """You are a SQL expert. Given an input question, Answer user questions by generating SQL queries -against the database schema provided in tools -Unless the user specifies in the question a specific number of examples to obtain, query for at most 5 results using the -LIMIT clause as per SQL standards. You can order the results to return the most informative data in the database. -Never query for all columns from a table. You must query only the columns that are needed to answer the question. -Wrap each column name in backticks (`) to denote them as identifiers. -Pay attention to use only the column names you can see in the tables below. -Be careful to not query for columns that do not exist. Also, pay attention to which column is in which table. -Pay attention to use CURRENT_DATE function to get the current date, if the question involves "today".""" - -messages = [{ - "role":"system", "content":SYSTEM_PROMPT - }, {"role":"user", "content":"what was the average delay on arrivals?"}] - -chat_response = chat_completion_request(client=client_mindsdb_serve, model=MODEL, messages=messages, tools=tools, tool_choice=None) - -assistant_message = chat_response.choices[0].message - -assistant_message.content = str(assistant_message.tool_calls[0].function) - -messages.append({"role": assistant_message.role, "content": assistant_message.content}) - -if assistant_message.tool_calls: - results = execute_function_call(message=assistant_message, database=database) - messages.append({ - "role": "function", "tool_call_id": assistant_message.tool_calls[0].id, - "name": assistant_message.tool_calls[0].function.name, - "content": results - }) - -pretty_print_conversation(messages) diff --git a/mindsdb_sdk/__about__.py b/mindsdb_sdk/__about__.py index b5a66c1..6a1198a 100755 --- a/mindsdb_sdk/__about__.py +++ b/mindsdb_sdk/__about__.py @@ -1,6 +1,6 @@ __title__ = 'mindsdb_sdk' __package_name__ = 'mindsdb_sdk' -__version__ = '2.4.3' +__version__ = '2.4.4' __description__ = "MindsDB Python SDK, provides an SDK to use a remote mindsdb instance" __email__ = "jorge@mindsdb.com" __author__ = 'MindsDB Inc' diff --git a/mindsdb_sdk/agents.py b/mindsdb_sdk/agents.py index 444a7d4..7811324 100644 --- a/mindsdb_sdk/agents.py +++ b/mindsdb_sdk/agents.py @@ -58,7 +58,7 @@ class Agent: Delete an agent by name: - >>> agents.delete('my_agent') + >>> agents.drop('my_agent') """ def __init__( self, diff --git a/mindsdb_sdk/databases.py b/mindsdb_sdk/databases.py index f5e0257..9ab9821 100644 --- a/mindsdb_sdk/databases.py +++ b/mindsdb_sdk/databases.py @@ -27,9 +27,10 @@ class Database: """ - def __init__(self, server, name): + def __init__(self, server, name, engine=None): self.server = server self.name = name + self.engine = engine self.api = server.api self.tables = Tables(self, self.api) @@ -82,9 +83,9 @@ def __init__(self, api): def _list_databases(self): data = self.api.sql_query( - "select NAME from information_schema.databases where TYPE='data'" + "select NAME, ENGINE from information_schema.databases where TYPE='data'" ) - return list(data.NAME) + return dict(zip(data.NAME, data.ENGINE)) def list(self) -> List[Database]: """ @@ -92,7 +93,8 @@ def list(self) -> List[Database]: :return: list of Database objects """ - return [Database(self, name) for name in self._list_databases()] + databases = self._list_databases() + return [Database(self, name, engine=engine) for name, engine in databases.items()] def create(self, name: str, engine: Union[str, Handler], connection_args: dict) -> Database: """ @@ -112,7 +114,7 @@ def create(self, name: str, engine: Union[str, Handler], connection_args: dict) parameters=connection_args, ) self.api.sql_query(ast_query.to_string()) - return Database(self, name) + return Database(self, name, engine=engine) def drop(self, name: str): """ @@ -130,8 +132,7 @@ def get(self, name: str) -> Database: :param name: name of integration :return: Database object """ - if name not in self._list_databases(): + databases = self._list_databases() + if name not in databases: raise AttributeError("Database doesn't exist") - return Database(self, name) - - + return Database(self, name, engine=databases[name]) diff --git a/mindsdb_sdk/models.py b/mindsdb_sdk/models.py index 8b7b1ce..6d44565 100644 --- a/mindsdb_sdk/models.py +++ b/mindsdb_sdk/models.py @@ -7,7 +7,7 @@ from mindsdb_sql.parser.dialects.mindsdb import CreatePredictor, DropPredictor from mindsdb_sql.parser.dialects.mindsdb import RetrainPredictor, FinetunePredictor -from mindsdb_sql.parser.ast import Identifier, Select, Star, Join, Update, Describe, Constant +from mindsdb_sql.parser.ast import Identifier, Select, Star, Join, Describe, Set from mindsdb_sql import parse_sql from mindsdb_sql.exceptions import ParsingException @@ -388,15 +388,9 @@ def set_active(self, version: int): :param version: version to set active """ - ast_query = Update( - table=Identifier(parts=[self.project.name, 'models_versions']), - update_columns={ - 'active': Constant(1) - }, - where=dict_to_binary_op({ - 'name': self.name, - 'version': version - }) + ast_query = Set( + category='active', + value=Identifier(parts=[self.project.name, self.name, str(version)]) ) sql = ast_query.to_string() if is_saving(): @@ -609,21 +603,22 @@ def list(self, with_versions: bool = False, :return: list of Model or ModelVersion objects """ - table = 'models' model_class = Model - if with_versions: - table = 'models_versions' - model_class = ModelVersion - filters = { } + filters = {} if name is not None: filters['NAME'] = name if version is not None: filters['VERSION'] = version + if with_versions: + model_class = ModelVersion + else: + filters['ACTIVE'] = '1' + ast_query = Select( targets=[Star()], - from_table=Identifier(table), + from_table=Identifier('models'), where=dict_to_binary_op(filters) ) df = self.project.query(ast_query.to_string()).fetch() diff --git a/mindsdb_sdk/projects.py b/mindsdb_sdk/projects.py index e734ab4..b8bda07 100644 --- a/mindsdb_sdk/projects.py +++ b/mindsdb_sdk/projects.py @@ -1,10 +1,8 @@ from typing import List -from mindsdb_sql.parser.dialects.mindsdb import CreateDatabase +from mindsdb_sql.parser.dialects.mindsdb import CreateDatabase, DropPredictor from mindsdb_sql.parser.ast import DropDatabase -from mindsdb_sql.parser.ast import Identifier, Delete - -from mindsdb_sdk.utils.sql import dict_to_binary_op +from mindsdb_sql.parser.ast import Identifier from mindsdb_sdk.agents import Agents from mindsdb_sdk.databases import Databases @@ -96,7 +94,6 @@ def query(self, sql: str) -> Query: """ return Query(self.api, sql, database=self.name) - def drop_model_version(self, name: str, version: int): """ Drop version of the model @@ -104,13 +101,8 @@ def drop_model_version(self, name: str, version: int): :param name: name of the model :param version: version to drop """ - ast_query = Delete( - table=Identifier('models_versions'), - where=dict_to_binary_op({ - 'name': name, - 'version': version - }) - ) + ast_query = DropPredictor(Identifier(parts=[name, str(version)])) + self.query(ast_query.to_string()).fetch() diff --git a/mindsdb_sdk/utils/mind.py b/mindsdb_sdk/utils/mind.py index c7c298f..e7de32f 100644 --- a/mindsdb_sdk/utils/mind.py +++ b/mindsdb_sdk/utils/mind.py @@ -1,3 +1,5 @@ +from typing import Optional + import requests from logging import getLogger @@ -18,12 +20,11 @@ def __init__(self, name): def create_mind( base_url: str, api_key: str, - name: str, description: str, - model: str, data_source_type: str, data_source_connection_args: dict, + model: Optional[str] = None, ) -> Mind: """ Create a mind entity in LiteLLM proxy. @@ -54,7 +55,14 @@ def create_mind( response = requests.post(url, json=payload, headers=headers) response.raise_for_status() except requests.exceptions.HTTPError as e: - logger.error(f"Failed to create mind: {e.response.json()}") + try: + error_message = e.response.json() + except Exception: + error_message = str(e) + logger.error(f"Failed to create mind: {error_message}") + raise e + except Exception as e: + logger.error(f"Failed to create mind: {e}") raise e name = response.json()['name'] diff --git a/tests/test_sdk.py b/tests/test_sdk.py index 2d1bba8..ef8acde 100644 --- a/tests/test_sdk.py +++ b/tests/test_sdk.py @@ -183,7 +183,7 @@ def check_model(self, model, database, mock_post): # list all versions models = model.list_versions() - check_sql_call(mock_post, f"SELECT * FROM models_versions WHERE NAME = '{model.name}'", + check_sql_call(mock_post, f"SELECT * FROM models WHERE NAME = '{model.name}'", database=model.project.name) model2 = models[0] # Model object @@ -194,8 +194,7 @@ def check_model(self, model, database, mock_post): # get call before last call mock_call = mock_post.call_args_list[-2] - assert mock_call[1]['json'][ - 'query'] == f"update {model2.project.name}.models_versions set active=1 where name = '{model2.name}' AND version = 3" + assert mock_call[1]['json']['query'] == f"SET active {model2.project.name}.{model2.name}.`3`" @patch('requests.Session.post') def check_table(self, table, mock_post): @@ -238,11 +237,11 @@ def test_flow(self, mock_post, mock_put, mock_get): assert call_args[0][0] == 'https://cloud.mindsdb.com/api/status' # --------- databases ------------- - response_mock(mock_post, pd.DataFrame([{'NAME': 'db1'}])) + response_mock(mock_post, pd.DataFrame([{'NAME': 'db1','ENGINE': 'postgres'}])) databases = server.list_databases() - check_sql_call(mock_post, "select NAME from information_schema.databases where TYPE='data'") + check_sql_call(mock_post, "select NAME, ENGINE from information_schema.databases where TYPE='data'") database = databases[0] str(database) @@ -284,7 +283,7 @@ def test_flow(self, mock_post, mock_put, mock_get): check_sql_call(mock_post, 'DROP DATABASE `proj1-1`') # test upload file - response_mock(mock_post, pd.DataFrame([{'NAME': 'files'}])) + response_mock(mock_post, pd.DataFrame([{'NAME': 'files', 'ENGINE': 'file'}])) database = server.get_database('files') # create file df = pd.DataFrame([{'s': '1'}, {'s': 'a'}]) @@ -494,7 +493,7 @@ def check_project_models_versions(self, project, database, mock_post): self.check_model(model, database) project.drop_model_version('m1', 1) - check_sql_call(mock_post, f"delete from models_versions where name='m1' and version=1") + check_sql_call(mock_post, f"DROP PREDICTOR m1.`1`") @patch('requests.Session.post') @@ -609,11 +608,11 @@ def test_flow(self, mock_post, mock_put): assert call_args[1]['json']['email'] == 'a@b.com' # --------- databases ------------- - response_mock(mock_post, pd.DataFrame([{'NAME': 'db1'}])) + response_mock(mock_post, pd.DataFrame([{'NAME': 'db1', 'ENGINE': 'postgres'}])) databases = con.databases.list() - check_sql_call(mock_post, "select NAME from information_schema.databases where TYPE='data'") + check_sql_call(mock_post, "select NAME, ENGINE from information_schema.databases where TYPE='data'") database = databases[0] assert database.name == 'db1' @@ -660,7 +659,7 @@ def test_flow(self, mock_post, mock_put): check_sql_call(mock_post, 'DROP DATABASE `proj1-1`') # test upload file - response_mock(mock_post, pd.DataFrame([{'NAME': 'files'}])) + response_mock(mock_post, pd.DataFrame([{'NAME': 'files', 'ENGINE': 'file'}])) database = con.databases.files # create file df = pd.DataFrame([{'s': '1'}, {'s': 'a'}]) @@ -961,7 +960,7 @@ def check_project_models_versions(self, project, database, mock_post): self.check_model(model, database) project.models.m1.drop_version(1) - check_sql_call(mock_post, f"delete from models_versions where name='m1' and version=1") + check_sql_call(mock_post, f"DROP PREDICTOR m1.`1`") @patch('requests.Session.post') def check_database(self, database, mock_post): @@ -1614,7 +1613,7 @@ def test_add_database(self, mock_post, mock_put, mock_get): responses_mock(mock_post, [ # DB get (POST /sql). pd.DataFrame([ - {'NAME': 'existing_db'} + {'NAME': 'existing_db', 'ENGINE': 'postgres'} ]), # DB tables get (POST /sql). pd.DataFrame([