diff --git a/README.md b/README.md index 4d630a7..62787c4 100644 --- a/README.md +++ b/README.md @@ -110,12 +110,4 @@ We value your feedback and strive to improve MindSQL. Here's how you can share y Thank you for your interest in contributing to our project! We appreciate your support and look forward to working with you. 🚀 -## 🌟 Contributors - -| GitHub Profile | Link + Image | Name | -|---------------------|-------------------------------------------------------------------------------------------------|-----------------| -| siddhant-mi | [![](https://github.com/siddhant-mi.png?size=50)](https://github.com/siddhant-mi) | Siddhant Pandey | -| ishika-mi | [![](https://github.com/ishika-mi.png?size=50)](https://github.com/ishika-mi) | Ishika Shah | -| Hasmukhsuthar05 | [![](https://github.com/Hasmukhsuthar05.png?size=50)](https://github.com/Hasmukhsuthar05) | Hasmukh Suthar | -| krishna-thakkar-mi | [![](https://github.com/krishna-thakkar-mi.png?size=50)](https://github.com/krishna-thakkar-mi) | Krishna Thakkar | -| UjjawalKRoy | [![](https://github.com/UjjawalKRoy.png?size=50)](https://github.com/UjjawalKRoy) | Ujjawal Roy | + diff --git a/mindsql/_utils/constants.py b/mindsql/_utils/constants.py index 9f5dea1..66a7408 100644 --- a/mindsql/_utils/constants.py +++ b/mindsql/_utils/constants.py @@ -17,7 +17,7 @@ ERROR_WHILE_RUNNING_QUERY = "Error while running query: {}" MYSQL_SHOW_DATABASE_QUERY = "SHOW DATABASES;" MYSQL_DB_TABLES_INFO_SCHEMA_QUERY = "SELECT table_name FROM information_schema.tables WHERE table_schema = '{}';" -MYSQL_SHOW_CREATE_TABLE_QUERY = "SHOW CREATE TABLE {};" +MYSQL_SHOW_CREATE_TABLE_QUERY = "SHOW CREATE TABLE `{}`;" POSTGRESQL_SHOW_DATABASE_QUERY = "SELECT datname as DATABASE_NAME FROM pg_database WHERE datistemplate = false;" POSTGRESQL_DB_TABLES_INFO_SCHEMA_QUERY = "SELECT table_name FROM information_schema.tables WHERE table_schema = 'public' AND table_catalog = '{db}';" ERROR_DOWNLOADING_SQLITE_DB_CONSTANT = "Error downloading sqlite db: {}" @@ -30,5 +30,10 @@ CONFIG_REQUIRED_ERROR = "Configuration is required." LLAMA_PROMPT_EXCEPTION = "Prompt cannot be empty." OPENAI_VALUE_ERROR = "OpenAI API key is required" -OPENAI_PROMPT_EMPTY_EXCEPTION = "Prompt cannot be empty." +PROMPT_EMPTY_EXCEPTION = "Prompt cannot be empty." POSTGRESQL_SHOW_CREATE_TABLE_QUERY = """SELECT 'CREATE TABLE "' || table_name || '" (' || array_to_string(array_agg(column_name || ' ' || data_type), ', ') || ');' AS create_statement FROM information_schema.columns WHERE table_name = '{table}' GROUP BY table_name;""" +ANTHROPIC_VALUE_ERROR = "Anthropic API key is required" +SQLSERVER_SHOW_DATABASE_QUERY= "SELECT name FROM sys.databases;" +SQLSERVER_DB_TABLES_INFO_SCHEMA_QUERY = "SELECT CONCAT(TABLE_SCHEMA,'.',TABLE_NAME) FROM [{db}].INFORMATION_SCHEMA.TABLES WHERE TABLE_TYPE = 'BASE TABLE'" +SQLSERVER_SHOW_CREATE_TABLE_QUERY = "DECLARE @TableName NVARCHAR(MAX) = '{table}'; DECLARE @SchemaName NVARCHAR(MAX) = '{schema}'; DECLARE @SQL NVARCHAR(MAX); SELECT @SQL = 'CREATE TABLE ' + @SchemaName + '.' + t.name + ' (' + CHAR(13) + ( SELECT ' ' + c.name + ' ' + UPPER(tp.name) + CASE WHEN tp.name IN ('char', 'varchar', 'nchar', 'nvarchar') THEN '(' + CASE WHEN c.max_length = -1 THEN 'MAX' ELSE CAST(c.max_length AS VARCHAR(10)) END + ')' WHEN tp.name IN ('decimal', 'numeric') THEN '(' + CAST(c.precision AS VARCHAR(10)) + ',' + CAST(c.scale AS VARCHAR(10)) + ')' ELSE '' END + ',' + CHAR(13) FROM sys.columns c JOIN sys.types tp ON c.user_type_id = tp.user_type_id WHERE c.object_id = t.object_id ORDER BY c.column_id FOR XML PATH(''), TYPE ).value('.', 'NVARCHAR(MAX)') + CHAR(13) + ')' FROM sys.tables t JOIN sys.schemas s ON t.schema_id = s.schema_id WHERE t.name = @TableName AND s.name = @SchemaName; SELECT @SQL AS SQLQuery;" +OLLAMA_CONFIG_REQUIRED = "{type} configuration is required." diff --git a/mindsql/databases/__init__.py b/mindsql/databases/__init__.py index 109c0ef..0034303 100644 --- a/mindsql/databases/__init__.py +++ b/mindsql/databases/__init__.py @@ -2,3 +2,4 @@ from .mysql import MySql from .postgres import Postgres from .sqlite import Sqlite +from .sqlserver import SQLServer diff --git a/mindsql/databases/sqlserver.py b/mindsql/databases/sqlserver.py new file mode 100644 index 0000000..97de59a --- /dev/null +++ b/mindsql/databases/sqlserver.py @@ -0,0 +1,147 @@ +from typing import List, Optional +from urllib.parse import urlparse + +import pandas as pd +import pyodbc + +from . import IDatabase +from .._utils import logger +from .._utils.constants import ERROR_WHILE_RUNNING_QUERY, ERROR_CONNECTING_TO_DB_CONSTANT, INVALID_DB_CONNECTION_OBJECT, \ + CONNECTION_ESTABLISH_ERROR_CONSTANT, SQLSERVER_SHOW_DATABASE_QUERY, SQLSERVER_DB_TABLES_INFO_SCHEMA_QUERY, \ + SQLSERVER_SHOW_CREATE_TABLE_QUERY + +log = logger.init_loggers("SQL Server") + + +class SQLServer(IDatabase): + @staticmethod + def create_connection(url: str, **kwargs) -> any: + """ + Connects to a SQL Server database using the provided URL. + + Parameters: + - url (str): The connection string to the SQL Server database in the format: + 'DRIVER={ODBC Driver 17 for SQL Server};SERVER=server_name;DATABASE=database_name;UID=user;PWD=password' + - **kwargs: Additional keyword arguments for the connection + + Returns: + - connection: A connection to the SQL Server database + """ + + try: + connection = pyodbc.connect(url, **kwargs) + return connection + except pyodbc.Error as e: + log.error(ERROR_CONNECTING_TO_DB_CONSTANT.format("SQL Server", e)) + + def execute_sql(self, connection, sql:str) -> Optional[pd.DataFrame]: + """ + A function that runs an SQL query using the provided connection and returns the results as a pandas DataFrame. + + Parameters: + connection: The connection object for the database. + sql (str): The SQL query to be executed + + Returns: + pd.DataFrame: A DataFrame containing the results of the SQL query. + """ + try: + self.validate_connection(connection) + cursor = connection.cursor() + cursor.execute(sql) + columns = [column[0] for column in cursor.description] + data = cursor.fetchall() + data = [list(row) for row in data] + cursor.close() + return pd.DataFrame(data, columns=columns) + except pyodbc.Error as e: + log.error(ERROR_WHILE_RUNNING_QUERY.format(e)) + return None + + def get_databases(self, connection) -> List[str]: + """ + Get a list of databases from the given connection and SQL query. + + Parameters: + connection: The connection object for the database. + + Returns: + List[str]: A list of unique database names. + """ + try: + self.validate_connection(connection) + cursor = connection.cursor() + cursor.execute(SQLSERVER_SHOW_DATABASE_QUERY) + databases = [row[0] for row in cursor.fetchall()] + cursor.close() + return databases + except pyodbc.Error as e: + log.error(ERROR_WHILE_RUNNING_QUERY.format(e)) + return [] + + def get_table_names(self, connection, database: str) -> pd.DataFrame: + """ + Retrieves the tables along with their schema (schema.table_name) from the information schema for the specified + database. + + Parameters: + connection: The database connection object. + database (str): The name of the database. + + Returns: + DataFrame: A pandas DataFrame containing the table names from the information schema. + """ + self.validate_connection(connection) + query = SQLSERVER_DB_TABLES_INFO_SCHEMA_QUERY.format(db=database) + return self.execute_sql(connection, query) + + + + + def get_all_ddls(self, connection: any, database: str) -> pd.DataFrame: + """ + A method to get the DDLs for all the tables in the database. + + Parameters: + connection (any): The connection object. + database (str): The name of the database. + + Returns: + DataFrame: A pandas DataFrame containing the DDLs for all the tables in the database. + """ + df_tables = self.get_table_names(connection, database) + ddl_df = pd.DataFrame(columns=['Table', 'DDL']) + for index, row in df_tables.iterrows(): + ddl = self.get_ddl(connection, row.iloc[0]) + ddl_df = ddl_df._append({'Table': row.iloc[0], 'DDL': ddl}, ignore_index=True) + + return ddl_df + + + + def validate_connection(self, connection: any) -> None: + """ + A function that validates if the provided connection is a SQL Server connection. + + Parameters: + connection: The connection object for accessing the database. + + Raises: + ValueError: If the provided connection is not a SQL Server connection. + + Returns: + None + """ + if connection is None: + raise ValueError(CONNECTION_ESTABLISH_ERROR_CONSTANT) + if not isinstance(connection, pyodbc.Connection): + raise ValueError(INVALID_DB_CONNECTION_OBJECT.format("SQL Server")) + + def get_ddl(self, connection: any, table_name: str, **kwargs) -> str: + schema_name, table_name = table_name.split('.') + query = SQLSERVER_SHOW_CREATE_TABLE_QUERY.format(table=table_name, schema=schema_name) + df_ddl = self.execute_sql(connection, query) + return df_ddl['SQLQuery'][0] + + def get_dialect(self) -> str: + return 'tsql' diff --git a/mindsql/llms/__init__.py b/mindsql/llms/__init__.py index 61ec7b8..9aff339 100644 --- a/mindsql/llms/__init__.py +++ b/mindsql/llms/__init__.py @@ -1,5 +1,6 @@ from .illm import ILlm +from .anthropic import AnthropicAi from .googlegenai import GoogleGenAi +from .huggingface import HuggingFace from .llama import LlamaCpp from .open_ai import OpenAi -from .huggingface import HuggingFace diff --git a/mindsql/llms/anthropic.py b/mindsql/llms/anthropic.py new file mode 100644 index 0000000..4f37984 --- /dev/null +++ b/mindsql/llms/anthropic.py @@ -0,0 +1,91 @@ +from anthropic import Anthropic + +from .illm import ILlm +from .._utils.constants import ANTHROPIC_VALUE_ERROR, PROMPT_EMPTY_EXCEPTION + + +class AnthropicAi(ILlm): + def __init__(self, config=None, client=None): + """ + Initialize the class with an optional config parameter. + + Parameters: + config (any): The configuration parameter. + client (any): The client parameter. + + Returns: + None + """ + self.config = config + self.client = client + + if client is not None: + self.client = client + return + + if 'api_key' not in config: + raise ValueError(ANTHROPIC_VALUE_ERROR) + api_key = config.pop('api_key') + self.client = Anthropic(api_key=api_key, **config) + + def system_message(self, message: str) -> any: + """ + Create a system message. + + Parameters: + message (str): The message parameter. + + Returns: + any + """ + return {"role": "system", "content": message} + + def user_message(self, message: str) -> any: + """ + Create a user message. + + Parameters: + message (str): The message parameter. + + Returns: + any + """ + return {"role": "user", "content": message} + + def assistant_message(self, message: str) -> any: + """ + Create an assistant message. + + Parameters: + message (str): The message parameter. + + Returns: + any + """ + return {"role": "assistant", "content": message} + + def invoke(self, prompt, **kwargs) -> str: + """ + Submit a prompt to the model for generating a response. + + Parameters: + prompt (str): The prompt parameter. + **kwargs: Additional keyword arguments (optional). + - temperature (float): The temperature parameter for controlling randomness in generation. + - max_tokens (int): Maximum number of tokens to be generated. + Returns: + str: The generated response from the model. + """ + if prompt is None or len(prompt) == 0: + raise Exception(PROMPT_EMPTY_EXCEPTION) + + model = self.config.get("model", "claude-3-opus-20240229") + temperature = kwargs.get("temperature", 0.1) + max_tokens = kwargs.get("max_tokens", 1024) + response = self.client.messages.create(model=model, messages=[{"role": "user", "content": prompt}], + max_tokens=max_tokens, temperature=temperature) + for content in response.content: + if isinstance(content, dict) and content.get("type") == "text": + return content["text"] + elif hasattr(content, "text"): + return content.text diff --git a/mindsql/llms/googlegenai.py b/mindsql/llms/googlegenai.py index 9e80582..80723c9 100644 --- a/mindsql/llms/googlegenai.py +++ b/mindsql/llms/googlegenai.py @@ -1,7 +1,7 @@ import google.generativeai as genai from .._utils.constants import GOOGLE_GEN_AI_VALUE_ERROR, GOOGLE_GEN_AI_APIKEY_ERROR -from . import ILlm +from .illm import ILlm class GoogleGenAi(ILlm): diff --git a/mindsql/llms/ollama.py b/mindsql/llms/ollama.py new file mode 100644 index 0000000..647bdd9 --- /dev/null +++ b/mindsql/llms/ollama.py @@ -0,0 +1,105 @@ +from ollama import Client, Options + +from .illm import ILlm +from .._utils.constants import PROMPT_EMPTY_EXCEPTION, OLLAMA_CONFIG_REQUIRED +from .._utils import logger + +log = logger.init_loggers("Ollama Client") + + +class Ollama(ILlm): + def __init__(self, model_config: dict, client_config=None, client: Client = None): + """ + Initialize the class with an optional config parameter. + + Parameters: + model_config (dict): The model configuration parameter. + config (dict): The configuration parameter. + client (Client): The client parameter. + + Returns: + None + """ + self.client = client + self.client_config = client_config + self.model_config = model_config + + if self.client is not None: + if self.client_config is not None: + log.warning("Client object provided. Ignoring client_config parameter.") + return + + if client_config is None: + raise ValueError(OLLAMA_CONFIG_REQUIRED.format(type="Client")) + + if model_config is None: + raise ValueError(OLLAMA_CONFIG_REQUIRED.format(type="Model")) + + if 'model' not in model_config: + raise ValueError(OLLAMA_CONFIG_REQUIRED.format(type="Model name")) + + self.client = Client(**client_config) + + def system_message(self, message: str) -> any: + """ + Create a system message. + + Parameters: + message (str): The message parameter. + + Returns: + any + """ + return {"role": "system", "content": message} + + def user_message(self, message: str) -> any: + """ + Create a user message. + + Parameters: + message (str): The message parameter. + + Returns: + any + """ + return {"role": "user", "content": message} + + def assistant_message(self, message: str) -> any: + """ + Create an assistant message. + + Parameters: + message (str): The message parameter. + + Returns: + any + """ + return {"role": "assistant", "content": message} + + def invoke(self, prompt, **kwargs) -> str: + """ + Submit a prompt to the model for generating a response. + + Parameters: + prompt (str): The prompt parameter. + **kwargs: Additional keyword arguments (optional). + - temperature (float): The temperature parameter for controlling randomness in generation. + + Returns: + str + """ + if not prompt: + raise ValueError(PROMPT_EMPTY_EXCEPTION) + + model = self.model_config.get('model') + temperature = kwargs.get('temperature', 0.1) + + response = self.client.chat( + model=model, + messages=[self.user_message(prompt)], + options=Options( + temperature=temperature + ) + ) + + return response['message']['content'] diff --git a/mindsql/llms/open_ai.py b/mindsql/llms/open_ai.py index 5cf63a9..b9bd4f9 100644 --- a/mindsql/llms/open_ai.py +++ b/mindsql/llms/open_ai.py @@ -1,7 +1,7 @@ from openai import OpenAI -from . import ILlm -from .._utils.constants import OPENAI_VALUE_ERROR, OPENAI_PROMPT_EMPTY_EXCEPTION +from .illm import ILlm +from .._utils.constants import OPENAI_VALUE_ERROR, PROMPT_EMPTY_EXCEPTION class OpenAi(ILlm): @@ -77,7 +77,7 @@ def invoke(self, prompt, **kwargs) -> str: str: The generated response from the model. """ if prompt is None or len(prompt) == 0: - raise Exception(OPENAI_PROMPT_EMPTY_EXCEPTION) + raise Exception(PROMPT_EMPTY_EXCEPTION) model = self.config.get("model", "gpt-3.5-turbo") temperature = kwargs.get("temperature", 0.1) diff --git a/mindsql/vectorstores/__init__.py b/mindsql/vectorstores/__init__.py index c8e0797..ad17496 100644 --- a/mindsql/vectorstores/__init__.py +++ b/mindsql/vectorstores/__init__.py @@ -1,3 +1,4 @@ from .ivectorstore import IVectorstore from .chromadb import ChromaDB from .faiss_db import Faiss +from .qdrant import Qdrant diff --git a/mindsql/vectorstores/qdrant.py b/mindsql/vectorstores/qdrant.py new file mode 100644 index 0000000..fc74b6c --- /dev/null +++ b/mindsql/vectorstores/qdrant.py @@ -0,0 +1,158 @@ +import json +import os +import uuid +from typing import List + +import pandas as pd +from qdrant_client import QdrantClient +from qdrant_client.http.models import Distance, VectorParams, PointStruct +from sentence_transformers import SentenceTransformer + +from . import IVectorstore + +sentence_transformer_ef = SentenceTransformer("WhereIsAI/UAE-Large-V1") + + +class Qdrant(IVectorstore): + def __init__(self, config=None): + if config is not None: + self.embedding_function = config.get( + "embedding_function", sentence_transformer_ef + ) + self.dimension = config.get("dimension", 1024) + qdrant_client_options = config.get("qdrant_client_options", {}) + else: + self.embedding_function = sentence_transformer_ef + self.dimension = 1024 + qdrant_client_options = {} + self.client = QdrantClient(**qdrant_client_options) + self._init_collections() + + def _init_collections(self): + for name in ["sql", "ddl", "documentation"]: + if not self.client.collection_exists(collection_name=name): + self.client.create_collection( + collection_name=name, + vectors_config=VectorParams( + size=self.dimension, distance=Distance.COSINE + ), + ) + + def index_question_sql(self, question: str, sql: str, **kwargs) -> str: + question_sql_json = json.dumps( + {"question": question, "sql": sql}, ensure_ascii=False + ) + chunk_id = str(uuid.uuid4()) + vector = self.embedding_function.encode([question_sql_json])[0] + self.client.upsert( + collection_name="sql", + points=[ + PointStruct( + id=chunk_id, vector=vector, payload={"data": question_sql_json} + ) + ], + ) + return chunk_id + "-sql" + + def index_ddl(self, ddl: str, **kwargs) -> str: + chunk_id = str(uuid.uuid4()) + table = kwargs.get("table", None) + vector = self.embedding_function.encode([ddl])[0] + payload = {"data": ddl} + if table: + payload["table_name"] = table + self.client.upsert( + collection_name="ddl", + points=[PointStruct(id=chunk_id, vector=vector, payload=payload)], + ) + return chunk_id + "-ddl" + + def index_documentation(self, documentation: str, **kwargs) -> str: + chunk_id = str(uuid.uuid4()) + vector = self.embedding_function.encode([documentation])[0] + self.client.upsert( + collection_name="documentation", + points=[ + PointStruct(id=chunk_id, vector=vector, payload={"data": documentation}) + ], + ) + return chunk_id + "-doc" + + def fetch_all_vectorstore_data(self, **kwargs) -> pd.DataFrame: + data = [] + for name in ["sql", "ddl", "documentation"]: + points = self.client.scroll(collection_name=name, limit=10000)[0] + for point in points: + payload = point.payload or {} + if name == "sql": + doc = json.loads(payload.get("data", "{}")) + question = doc.get("question") + content = doc.get("sql") + else: + question = None + content = payload.get("data") + data.append( + { + "id": point.id, + "question": question, + "content": content, + "training_data_type": name, + } + ) + return pd.DataFrame(data) + + def delete_vectorstore_data(self, item_id: str, **kwargs) -> bool: + uuid_str = item_id[:-4] + if item_id.endswith("-sql"): + self.client.delete(collection_name="sql", points_selector=[uuid_str]) + return True + elif item_id.endswith("-ddl"): + self.client.delete(collection_name="ddl", points_selector=[uuid_str]) + return True + elif item_id.endswith("-doc"): + self.client.delete( + collection_name="documentation", points_selector=[uuid_str] + ) + return True + else: + return False + + def remove_collection(self, collection_name: str) -> bool: + if self.client.collection_exists(collection_name=collection_name): + self.client.delete_collection(collection_name=collection_name) + self.client.create_collection( + collection_name=collection_name, + vectors_config=VectorParams( + size=self.dimension, distance=Distance.COSINE + ), + ) + return True + return False + + def retrieve_relevant_question_sql(self, question: str, **kwargs) -> list: + n = kwargs.get("n_results", 2) + vector = self.embedding_function.encode([question])[0] + hits = self.client.query_points( + collection_name="sql", query=vector, limit=n + ).points + results = [] + for hit in hits: + doc = json.loads(hit.payload.get("data", "{}")) + results.append(doc) + return results + + def retrieve_relevant_ddl(self, question: str, **kwargs) -> list: + n = kwargs.get("n_results", 2) + vector = self.embedding_function.encode([question])[0] + hits = self.client.query_points( + collection_name="ddl", query=vector, limit=n + ).points + return [hit.payload.get("data") for hit in hits] + + def retrieve_relevant_documentation(self, question: str, **kwargs) -> list: + n = kwargs.get("n_results", 2) + vector = self.embedding_function.encode([question])[0] + hits = self.client.query_points( + collection_name="documentation", query=vector, limit=n + ).points + return [hit.payload.get("data") for hit in hits] diff --git a/pyproject.toml b/pyproject.toml index 8251b96..dfa1649 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,21 +16,21 @@ classifiers = [ [tool.poetry.dependencies] -python = "^3.10" -chromadb = "^0.4.22" -pandas = "2.2.0" +python = "^3.11" +chromadb = "^1.0.15" +pandas = "2.3.1" plotly = "5.19.0" mysql-connector-python = "^8.3.0" google-generativeai="0.3.2" llama-cpp-python = "0.2.47" openai = "^1.12.0" sqlparse = "^0.4.4" -numpy = "^1.26.4" +numpy = "2.3.1" sentence-transformers = "^2.3.1" psycopg2-binary = "^2.9.9" -faiss-cpu = "^1.8.0" -pysqlite3-binary = "^0.5.2.post3" +faiss-cpu = "^1.11.0.post1" transformers = "^4.38.2" +qdrant-client = "^1.14.3" [build-system] diff --git a/tests/ollama_test.py b/tests/ollama_test.py new file mode 100644 index 0000000..385424f --- /dev/null +++ b/tests/ollama_test.py @@ -0,0 +1,83 @@ +import unittest +from unittest.mock import MagicMock, patch +from ollama import Client, Options + +from mindsql.llms import ILlm +from mindsql.llms import Ollama +from mindsql._utils.constants import PROMPT_EMPTY_EXCEPTION, OLLAMA_CONFIG_REQUIRED + + +class TestOllama(unittest.TestCase): + + def setUp(self): + # Common setup for each test case + self.model_config = {'model': 'sqlcoder'} + self.client_config = {'host': 'http://localhost:11434/'} + self.client_mock = MagicMock(spec=Client) + + def test_initialization_with_client(self): + ollama = Ollama(model_config=self.model_config, client=self.client_mock) + self.assertEqual(ollama.client, self.client_mock) + self.assertIsNone(ollama.client_config) + self.assertEqual(ollama.model_config, self.model_config) + + def test_initialization_with_client_config(self): + ollama = Ollama(model_config=self.model_config, client_config=self.client_config) + self.assertIsNotNone(ollama.client) + self.assertEqual(ollama.client_config, self.client_config) + self.assertEqual(ollama.model_config, self.model_config) + + def test_initialization_missing_client_and_client_config(self): + with self.assertRaises(ValueError) as context: + Ollama(model_config=self.model_config) + self.assertEqual(str(context.exception), OLLAMA_CONFIG_REQUIRED.format(type="Client")) + + def test_initialization_missing_model_config(self): + with self.assertRaises(ValueError) as context: + Ollama(model_config=None, client_config=self.client_config) + self.assertEqual(str(context.exception), OLLAMA_CONFIG_REQUIRED.format(type="Model")) + + def test_initialization_missing_model_name(self): + with self.assertRaises(ValueError) as context: + Ollama(model_config={}, client_config=self.client_config) + self.assertEqual(str(context.exception), OLLAMA_CONFIG_REQUIRED.format(type="Model name")) + + def test_system_message(self): + ollama = Ollama(model_config=self.model_config, client=self.client_mock) + message = ollama.system_message("Test system message") + self.assertEqual(message, {"role": "system", "content": "Test system message"}) + + def test_user_message(self): + ollama = Ollama(model_config=self.model_config, client=self.client_mock) + message = ollama.user_message("Test user message") + self.assertEqual(message, {"role": "user", "content": "Test user message"}) + + def test_assistant_message(self): + ollama = Ollama(model_config=self.model_config, client=self.client_mock) + message = ollama.assistant_message("Test assistant message") + self.assertEqual(message, {"role": "assistant", "content": "Test assistant message"}) + + @patch.object(Client, 'chat', return_value={'message': {'content': 'Test response'}}) + def test_invoke_success(self, mock_chat): + ollama = Ollama(model_config=self.model_config, client=Client()) + response = ollama.invoke("Test prompt") + + # Check if the response is as expected + self.assertEqual(response, 'Test response') + + # Verify that the chat method was called with the correct arguments + mock_chat.assert_called_once_with( + model=self.model_config['model'], + messages=[{"role": "user", "content": "Test prompt"}], + options=Options(temperature=0.1) + ) + + def test_invoke_empty_prompt(self): + ollama = Ollama(model_config=self.model_config, client=self.client_mock) + with self.assertRaises(ValueError) as context: + ollama.invoke("") + self.assertEqual(str(context.exception), PROMPT_EMPTY_EXCEPTION) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/sqlserver_test.py b/tests/sqlserver_test.py new file mode 100644 index 0000000..fd68159 --- /dev/null +++ b/tests/sqlserver_test.py @@ -0,0 +1,163 @@ +import unittest +from unittest.mock import patch, MagicMock +import pyodbc +import pandas as pd +from mindsql.databases.sqlserver import SQLServer, ERROR_WHILE_RUNNING_QUERY, ERROR_CONNECTING_TO_DB_CONSTANT, \ + INVALID_DB_CONNECTION_OBJECT, CONNECTION_ESTABLISH_ERROR_CONSTANT +from mindsql.databases.sqlserver import log as logger + + +class TestSQLServer(unittest.TestCase): + + @patch('mindsql.databases.sqlserver.pyodbc.connect') + def test_create_connection_success(self, mock_connect): + mock_connect.return_value = MagicMock(spec=pyodbc.Connection) + connection = SQLServer.create_connection( + 'DRIVER={ODBC Driver 17 for SQL Server};SERVER=server_name;DATABASE=database_name;UID=user;PWD=password') + self.assertIsInstance(connection, pyodbc.Connection) + + @patch('mindsql.databases.sqlserver.pyodbc.connect') + def test_create_connection_failure(self, mock_connect): + mock_connect.side_effect = pyodbc.Error('Connection failed') + with self.assertLogs(logger, level='ERROR') as cm: + connection = SQLServer.create_connection( + 'DRIVER={ODBC Driver 17 for SQL Server};SERVER=server_name;DATABASE=database_name;UID=user;PWD=password') + self.assertIsNone(connection) + self.assertTrue(any( + ERROR_CONNECTING_TO_DB_CONSTANT.format("SQL Server", 'Connection failed') in message for message in + cm.output)) + + @patch('mindsql.databases.sqlserver.pyodbc.connect') + def test_execute_sql_success(self, mock_connect): + # Mock the connection and cursor + mock_connection = MagicMock(spec=pyodbc.Connection) + mock_cursor = MagicMock() + + mock_connect.return_value = mock_connection + mock_connection.cursor.return_value = mock_cursor + + # Mock cursor behavior + mock_cursor.execute.return_value = None + mock_cursor.description = [('column1',), ('column2',)] + mock_cursor.fetchall.return_value = [(1, 'a'), (2, 'b')] + + connection = SQLServer.create_connection('fake_connection_string') + sql = "SELECT * FROM table" + sql_server = SQLServer() + result = sql_server.execute_sql(connection, sql) + expected_df = pd.DataFrame(data=[(1, 'a'), (2, 'b')], columns=['column1', 'column2']) + pd.testing.assert_frame_equal(result, expected_df) + + @patch('mindsql.databases.sqlserver.pyodbc.connect') + def test_execute_sql_failure(self, mock_connect): + # Mock the connection and cursor + mock_connection = MagicMock(spec=pyodbc.Connection) + mock_cursor = MagicMock() + + mock_connect.return_value = mock_connection + mock_connection.cursor.return_value = mock_cursor + mock_cursor.execute.side_effect = pyodbc.Error('Query failed') + + connection = SQLServer.create_connection('fake_connection_string') + sql = "SELECT * FROM table" + sql_server = SQLServer() + + with self.assertLogs(logger, level='ERROR') as cm: + result = sql_server.execute_sql(connection, sql) + self.assertIsNone(result) + self.assertTrue(any(ERROR_WHILE_RUNNING_QUERY.format('Query failed') in message for message in cm.output)) + + @patch('mindsql.databases.sqlserver.pyodbc.connect') + def test_get_databases_success(self, mock_connect): + # Mock the connection and cursor + mock_connection = MagicMock(spec=pyodbc.Connection) + mock_cursor = MagicMock() + + mock_connect.return_value = mock_connection + mock_connection.cursor.return_value = mock_cursor + + # Mock cursor behavior + mock_cursor.execute.return_value = None + mock_cursor.fetchall.return_value = [('database1',), ('database2',)] + + connection = SQLServer.create_connection('fake_connection_string') + sql_server = SQLServer() + result = sql_server.get_databases(connection) + self.assertEqual(result, ['database1', 'database2']) + + @patch('mindsql.databases.sqlserver.pyodbc.connect') + def test_get_databases_failure(self, mock_connect): + # Mock the connection and cursor + mock_connection = MagicMock(spec=pyodbc.Connection) + mock_cursor = MagicMock() + + mock_connect.return_value = mock_connection + mock_connection.cursor.return_value = mock_cursor + mock_cursor.execute.side_effect = pyodbc.Error('Query failed') + + connection = SQLServer.create_connection('fake_connection_string') + sql_server = SQLServer() + + with self.assertLogs(logger, level='ERROR') as cm: + result = sql_server.get_databases(connection) + self.assertEqual(result, []) + self.assertTrue(any(ERROR_WHILE_RUNNING_QUERY.format('Query failed') in message for message in cm.output)) + + @patch('mindsql.databases.sqlserver.SQLServer.execute_sql') + def test_get_table_names_success(self, mock_execute_sql): + mock_execute_sql.return_value = pd.DataFrame(data=[('schema1.table1',), ('schema2.table2',)], + columns=['table_name']) + + connection = MagicMock(spec=pyodbc.Connection) + sql_server = SQLServer() + result = sql_server.get_table_names(connection, 'database_name') + expected_df = pd.DataFrame(data=[('schema1.table1',), ('schema2.table2',)], columns=['table_name']) + pd.testing.assert_frame_equal(result, expected_df) + + @patch('mindsql.databases.sqlserver.SQLServer.execute_sql') + def test_get_all_ddls_success(self, mock_execute_sql): + mock_execute_sql.side_effect = [ + pd.DataFrame(data=[('schema1.table1',)], columns=['table_name']), + pd.DataFrame(data=['CREATE TABLE schema1.table1 (...);'], columns=['SQLQuery']) + ] + + connection = MagicMock(spec=pyodbc.Connection) + sql_server = SQLServer() + result = sql_server.get_all_ddls(connection, 'database_name') + + expected_df = pd.DataFrame(data=[{'Table': 'schema1.table1', 'DDL': 'CREATE TABLE schema1.table1 (...);'}]) + pd.testing.assert_frame_equal(result, expected_df) + + def test_validate_connection_success(self): + connection = MagicMock(spec=pyodbc.Connection) + sql_server = SQLServer() + # Should not raise any exception + sql_server.validate_connection(connection) + + def test_validate_connection_failure(self): + sql_server = SQLServer() + + with self.assertRaises(ValueError) as cm: + sql_server.validate_connection(None) + self.assertEqual(str(cm.exception), CONNECTION_ESTABLISH_ERROR_CONSTANT) + + with self.assertRaises(ValueError) as cm: + sql_server.validate_connection("InvalidConnectionObject") + self.assertEqual(str(cm.exception), INVALID_DB_CONNECTION_OBJECT.format("SQL Server")) + + @patch('mindsql.databases.sqlserver.SQLServer.execute_sql') + def test_get_ddl_success(self, mock_execute_sql): + mock_execute_sql.return_value = pd.DataFrame(data=['CREATE TABLE schema1.table1 (...);'], columns=['SQLQuery']) + + connection = MagicMock(spec=pyodbc.Connection) + sql_server = SQLServer() + result = sql_server.get_ddl(connection, 'schema1.table1') + self.assertEqual(result, 'CREATE TABLE schema1.table1 (...);') + + def test_get_dialect(self): + sql_server = SQLServer() + self.assertEqual(sql_server.get_dialect(), 'tsql') + + +if __name__ == '__main__': + unittest.main()