mistral_pgsql_function_calling_ethux.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ab418a2f-cb71-4276-a805-a434a40944e0",
   "metadata": {},
   "source": [
    "# Mistral Large V2 / 2407 Function calling - Retrieve information from a Postgresql Database\n",
    "\n",
    "This tutorial is written as example how to use Microsoft AutoGen in Combination with Mistral Nemo to query a Postgres Database.\n",
    "It gives a short and simple overview how to do this, it is also possible to write extra tools and add them to query other DB's or maybe VectorDB's with a perfect prompt ;)! \n",
    "\n",
    "Besides query's and prompting a good context is also really important, thats why I added a extra function to always provide the context of all avaible. \n",
    "\n",
    "**PLEASE NOTE:** Try to prevent using to many tools and contexts together, but use different 'chat models' instead of a single big model to do everything.\n",
    "\n",
    "\n",
    "\n",
    "First we have to install postgres locally before we can can continue, go to the following page;\n",
    "\n",
    "- https://postgresapp.com/\n",
    "\n",
    "After you installed postgres you need to import a database dump if you don't have any data within postgres yet.\n",
    "\n",
    "- https://wiki.postgresql.org/wiki/Sample_Databases\n",
    "\n",
    "\n",
    "I did go for the following Database dumb; name.basics.tsv, found here; https://datasets.imdbws.com/\n",
    "\n",
    "After downloading the database dump you can create a table and import the data into it.\n",
    "\n",
    "```\n",
    "# Change the names and data types within ( ) to your choice or to match your dataset.\n",
    "CREATE TABLE imdb ( nconst TEXT, primaryName TEXT, birthYear INT, deathYear INT, primaryProfession TEXT, knownForTitles TEXT);\n",
    "COPY imdb FROM '/Users/root/Downloads/name.basics.tsv' WITH (HEADER true);\n",
    "\n",
    "# Test if the data is imported\n",
    "SELECT * FROM imdb LIMIT 10;\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da899631-32e4-4612-8a07-ae7a9bcb7195",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install pyautogen psycopg2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86b2b919-397b-416d-9ca5-69e618e8121e",
   "metadata": {},
   "source": [
    "# Import all required packages\n",
    "\n",
    "Lets importat all required packages, in this case we need autogen, the postgresql package and some other libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "938180a6-7c14-4acd-b062-992b04309baa",
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogen import ConversableAgent, register_function\n",
    "from typing import List, Optional, Union, Dict, Any\n",
    "import psycopg2\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46b3e4ac-1b93-4387-90b7-8658fbe99e68",
   "metadata": {},
   "source": [
    "# The Postgress Function Tool\n",
    "\n",
    "The Postgress function tool takes a valid json input, based on this input the query is executed.\n",
    "After that the function returns the output to the LLM, which will respond on that with a message to the user.\n",
    "\n",
    "Because we pre-defined how to use the tool, it is not possible to delete , update or create any records, read access only!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b66935a-d31d-42b1-8e04-892045355d6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "postgres_user = input(\"Enter your Postgres user: \")\n",
    "postgres_password = input(\"Enter your Postgres password: \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d82c2eee-a081-49f4-9fad-24f1f9657907",
   "metadata": {},
   "outputs": [],
   "source": [
    "def execute_postgres_query(\n",
    "    table_name: str,\n",
    "    columns: List[str],\n",
    "    filters: Optional[Dict[str, Any]] = None,\n",
    "    sort_column: Optional[str] = None,\n",
    "    sort_order: Optional[str] = None,\n",
    "    limit: Optional[int] = 150,  # Default limit of 150 rows, you can edit this yourself if needed, the AI will also be able to change this.\n",
    "):\n",
    "    # Validate input\n",
    "    if not table_name:\n",
    "        return \"Error: table_name is required\"\n",
    "    if not columns:\n",
    "        return \"Error: columns is required\"\n",
    "    if sort_column and not sort_order:\n",
    "        return \"Error: sort_order is required when sort_column is specified\"\n",
    "\n",
    "    # Generate SQL query\n",
    "    query = f\"SELECT {', '.join(columns)} FROM {table_name}\"\n",
    "    params = []\n",
    "\n",
    "    if filters:\n",
    "        filter_conditions = []\n",
    "        for column, value in filters.items():\n",
    "            if isinstance(value, str) and value.startswith('%') and value.endswith('%'):\n",
    "                filter_conditions.append(f\"{column} LIKE %s\")\n",
    "                params.append(value)\n",
    "            elif isinstance(value, list):\n",
    "                filter_conditions.append(f\"{column} NOT IN %s\")\n",
    "                params.append(tuple(value))\n",
    "            else:\n",
    "                filter_conditions.append(f\"{column} = %s\")\n",
    "                params.append(value)\n",
    "        query += \" WHERE \" + \" AND \".join(filter_conditions)\n",
    "\n",
    "    if sort_column and sort_order:\n",
    "        query += f\" ORDER BY {sort_column} {sort_order}\"\n",
    "    if limit:\n",
    "        query += f\" LIMIT {limit}\"\n",
    "\n",
    "    # Execute SQL query\n",
    "    conn = psycopg2.connect(database=\"postgres\", user=postgres_user, password=postgres_password, host=\"localhost\", port=\"5433\")\n",
    "    cur = conn.cursor()\n",
    "    cur.execute(query, params)\n",
    "    results = cur.fetchall()\n",
    "    cur.close()\n",
    "    conn.close()\n",
    "    return results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "282111b8-eda6-4b37-a316-9635b6678d71",
   "metadata": {},
   "source": [
    "# Define API keys\n",
    "\n",
    "Fill in your Mistral API key to access Mistral-large-2407! :D\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da4ba448-12d0-4865-8099-3b0c41b446c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "mistral_key = input(\"Enter your Mistral AI key: \")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d60345db-d288-40d3-aa6a-63305e1a4758",
   "metadata": {},
   "source": [
    "# Get All tables\n",
    "\n",
    "We want to prompt the LLM with the context of all tables, this has to be up to date, so we create a sepperate function which queries the postgres tool but with pre defined input to query all tables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33aec574-4e6d-480e-869c-7f3e82cdadaa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_all_tables():\n",
    "    # Exclude default PostgreSQL schemas\n",
    "    excluded_schemas = ['information_schema', 'pg_catalog']\n",
    "\n",
    "    # Query to get all tables excluding the default schemas\n",
    "    table_columns = ['table_schema', 'table_name']\n",
    "    table_name = 'information_schema.tables'\n",
    "    filters = {'table_schema': excluded_schemas}\n",
    "    sort_column = 'table_schema'\n",
    "    sort_order = 'ASC'\n",
    "\n",
    "    # Execute the query to get all tables\n",
    "    tables_query_result = execute_postgres_query(\n",
    "        table_name,\n",
    "        table_columns,\n",
    "        filters,\n",
    "        sort_column,\n",
    "        sort_order\n",
    "    )\n",
    "\n",
    "    # Parse the results of the tables query\n",
    "    tables = [{'table_schema': row[0], 'table_name': row[1]} for row in tables_query_result]\n",
    "\n",
    "    # Prepare a list to store table information with columns\n",
    "    table_info = []\n",
    "\n",
    "    # Iterate over each table to get its columns\n",
    "    for table in tables:\n",
    "        schema_name = table['table_schema']\n",
    "        table_name = table['table_name']\n",
    "\n",
    "        # Query to get columns for the current table\n",
    "        columns_columns = ['column_name']\n",
    "        columns_table_name = 'information_schema.columns'\n",
    "        columns_filters = {'table_schema': schema_name, 'table_name': table_name}\n",
    "        columns_sort_column = 'ordinal_position'\n",
    "        columns_sort_order = 'ASC'\n",
    "\n",
    "        # Execute the query to get columns\n",
    "        columns_query_result = execute_postgres_query(\n",
    "            columns_table_name,\n",
    "            columns_columns,\n",
    "            columns_filters,\n",
    "            columns_sort_column,\n",
    "            columns_sort_order\n",
    "        )\n",
    "\n",
    "        # Parse the results of the columns query\n",
    "        columns = [row[0] for row in columns_query_result]\n",
    "\n",
    "        # Add table information with columns to the list\n",
    "        table_info.append({\n",
    "            'table_schema': schema_name,\n",
    "            'table_name': table_name,\n",
    "            'columns': columns\n",
    "        })\n",
    "\n",
    "    return table_info"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb49afae-cf21-4a71-a313-708cefddfd3e",
   "metadata": {},
   "source": [
    "# Execute user queries\n",
    "\n",
    "Now everything is set to use the chat and query Postgresql Database using Mistral-Large-2407"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31357af0-ddb1-4d83-8599-a72a4a55cd79",
   "metadata": {},
   "outputs": [],
   "source": [
    "def chatbot(mistral_key):\n",
    "    config_list = [\n",
    "        {\n",
    "            'model': 'mistral-large-2407',\n",
    "            'base_url': 'https://api.mistral.ai/v1',\n",
    "            \"api_key\": mistral_key,\n",
    "            \"tool_choice\": \"auto\",\n",
    "        },\n",
    "    ]\n",
    "\n",
    "    llm_config={\n",
    "        \"config_list\": config_list,\n",
    "        \"temperature\": 0.1\n",
    "    }\n",
    "\n",
    "    user = ConversableAgent(\n",
    "        \"user\",\n",
    "        llm_config=False,\n",
    "        is_termination_msg=lambda msg: \"tool_calls\" not in msg,\n",
    "        human_input_mode=\"NEVER\",\n",
    "        max_consecutive_auto_reply=3,\n",
    "    )\n",
    "    \n",
    "    assistant = ConversableAgent(\n",
    "    name=\"assistant\",\n",
    "        llm_config=llm_config,\n",
    "        system_message=\"You are an helpful AI assistant, you use your Postgres tool to query the database. Keep in mind the possibility of to long contexts lengths when using limits wrong.\"\n",
    "    )\n",
    "\n",
    "    assistant.register_for_llm(name=\"postgres_query\", description=\"Useful for when you need query the postgres db\")(execute_postgres_query)\n",
    "    user.register_for_execution(name=\"postgres_query\")(execute_postgres_query)\n",
    "\n",
    "    LLM_CONTEXT = get_all_tables()\n",
    "\n",
    "    user.send(f\"This are all the available tables; \\n\\n  {LLM_CONTEXT} \\n\\n \", assistant, request_reply=False)\n",
    "    assistant.send(\"Thanks for the additonal context of all existing tables!\", user, request_reply=False)\n",
    "    \n",
    "    while True:\n",
    "        task = input(\"Enter the query for the LLM ('exit' to quit): \")\n",
    "        if task.lower() == 'exit':\n",
    "            break\n",
    "        context_task = f\"{task}\"\n",
    "        user.initiate_chat(assistant, message=context_task, clear_history=False)\n",
    "\n",
    "chatbot(mistral_key)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "047c0b8d-2076-4b89-8ae7-cd41aa983744",
   "metadata": {},
   "source": [
    "# Conclusion\n",
    "\n",
    "With Mistral Large V2 or any equivelant model it is relatively easy to create a function to query a Postgres DB without giving it full access to delete records. \n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}