Using an AI model to ask questions to your database

# MY FIRST AI PROGRAM

from openai import OpenAI
from masoniteorm.query import QueryBuilder

question = input("Enter your question: ")

client = OpenAI(
    base_url="http://localhost:11434/v1",
    api_key="ANYTHING",  # If using OpenAI, use your API key
)

prompt = """
You are an AI assistant that answers questions in a friendly and clear manner.
Your goal is to answer questions about a database with the following structure:

table: products
fields:
- id: int
- name: varchar, product name
- price: int, product price
- stock: int, product stock
- brand: varchar, product brand

To solve user questions, you have the following tools:

1. get_data:
- Description: Gets data from the products table based on provided filters.
- Parameters:
    - filters: json, filters for the query. Example: {"name": "product1"}
    - Returns: list of dictionaries, each dictionary represents a product

Responses should be in JSON format and if you need to use a tool, return a response with this format:

{
    "tool": "get_data",
    "parameters": {
        "filters": {
            "brand": "Apple"
        },
    }
}
"""

response = client.chat.completions.create(
    model="mistral",
    messages=[
        {"role": "system", "content": prompt},
        {"role": "user", "content": question},
    ],
)

ai_response = response.choices[0].message.content

print(ai_response)

# AI DATABASE ASSISTANT

import json
import os
from openai import OpenAI
from masoniteorm.query import QueryBuilder
from masoniteorm.query.grammars import MySQLGrammar
from datetime import datetime
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Initialize database connection
QueryBuilder.set_connection_details({
    "default": "mysql",
    "mysql": {
        "host": "localhost",
        "driver": "mysql",
        "database": "product_db",
        "user": os.environ.get("DB_USER", "root"),
        "password": os.environ.get("DB_PASSWORD", ""),
        "port": 3306
    }
})

# Tool functions
def get_data(filters=None, limit=10, columns=None, order_by=None):
    query = QueryBuilder().table("products")

    if filters:
        for key, value in filters.items():
            if isinstance(value, dict) and "operator" in value:
                query.where(key, value["operator"], value["value"])
            else:
                query.where(key, "=", value)

    if columns:
        query.select(columns)

    if order_by:
        query.order_by(order_by["column"], order_by.get("direction", "asc"))

    return query.limit(limit).get()

def aggregate_data(operation, column, filters=None):
    query = QueryBuilder().table("products")

    if filters:
        for key, value in filters.items():
            query.where(key, "=", value)

    if operation == "sum":
        return query.sum(column)
    elif operation == "avg":
        return query.avg(column)
    elif operation == "min":
        return query.min(column)
    elif operation == "max":
        return query.max(column)
    elif operation == "count":
        return query.count()

    return None

# Process AI response and handle tool requests
def process_ai_response(ai_response, question):
    try:
        # Check if the response is a tool call
        tool_request = json.loads(ai_response)
        if "tool" in tool_request:
            # Execute the requested tool
            if tool_request["tool"] == "get_data":
                data = get_data(**tool_request["parameters"])

                # Send the data back to the AI for final formatting
                final_response = client.chat.completions.create(
                    model="mistral",
                    messages=[
                        {"role": "system", "content": prompt},
                        {"role": "user", "content": question},
                        {"role": "assistant", "content": ai_response},
                        {"role": "system", "content": f"Tool result: {json.dumps(data)}"}
                    ],
                )
                return final_response.choices[0].message.content

            elif tool_request["tool"] == "aggregate_data":
                result = aggregate_data(**tool_request["parameters"])

                # Send the result back to the AI for final formatting
                final_response = client.chat.completions.create(
                    model="mistral",
                    messages=[
                        {"role": "system", "content": prompt},
                        {"role": "user", "content": question},
                        {"role": "assistant", "content": ai_response},
                        {"role": "system", "content": f"Tool result: {result}"}
                    ],
                )
                return final_response.choices[0].message.content
    except json.JSONDecodeError:
        # If not a tool call, return the original response
        pass
    except Exception as e:
        return f"Error processing request: {str(e)}"

    return ai_response

# Main function
def ask_database(question):
    client = OpenAI(
        base_url="http://localhost:11434/v1",
        api_key="ANYTHING",
    )

    prompt = """
    You are an AI assistant that answers questions about product inventory in a friendly and clear manner.
    Your goal is to answer questions about a database with the following structure:

    table: products
    fields:
    - id: int
    - name: varchar, product name
    - price: int, product price in dollars
    - stock: int, current inventory count
    - brand: varchar, product manufacturer
    - category: varchar, product category
    - created_at: datetime, when the product was added

    To solve user questions, you have the following tools:

    1. get_data:
    - Description: Gets data from the products table based on provided filters.
    - Parameters:
        - filters: json, filters for the query. Example: {"name": "product1", "price": {"operator": ">", "value": 100}}
        - limit: int, maximum number of results to return (default 10)
        - columns: list, specific columns to return (default all)
        - order_by: json, ordering specification. Example: {"column": "price", "direction": "desc"}
    - Returns: list of dictionaries, each dictionary represents a product

    2. aggregate_data:
    - Description: Performs aggregate operations on the products table.
    - Parameters:
        - operation: string, one of ["sum", "avg", "min", "max", "count"]
        - column: string, the column to perform the operation on
        - filters: json, filters for the query (optional)
    - Returns: numeric result of the aggregate operation

    When answering:
    1. Determine if you need data from the database to answer the question
    2. Select the appropriate tool and parameters
    3. Format your tool request using exact JSON format
    4. When responding with data to the user, format it in a friendly, readable way

    If you need to use a tool, return ONLY a response with this format:
    {
        "tool": "tool_name",
        "parameters": {
            "param1": value1,
            "param2": value2
        }
    }
    """

    response = client.chat.completions.create(
        model="mistral",
        messages=[
            {"role": "system", "content": prompt},
            {"role": "user", "content": question},
        ],
    )

    ai_response = response.choices[0].message.content
    return process_ai_response(ai_response, question)

if __name__ == "__main__":
    question = input("Enter your question about the product inventory: ")
    response = ask_database(question)
    print(response)

# Create and activate virtual environment
python -m venv venv
source venv/bin/activate  # On Windows: venv\Scripts\activate

# Install required packages
pip install openai masoniteorm python-dotenv

ollama pull mistral

CREATE DATABASE product_db;

USE product_db;

CREATE TABLE products (
  id INT AUTO_INCREMENT PRIMARY KEY,
  name VARCHAR(255) NOT NULL,
  price INT NOT NULL,
  stock INT NOT NULL,
  brand VARCHAR(255),
  category VARCHAR(255),
  created_at DATETIME DEFAULT CURRENT_TIMESTAMP
);

-- Add some sample data
INSERT INTO products (name, price, stock, brand, category) VALUES
  ('Laptop Pro', 1299, 45, 'Apple', 'Electronics'),
  ('Smartphone X', 899, 120, 'Samsung', 'Electronics'),
  ('Wireless Earbuds', 149, 200, 'Apple', 'Audio'),
  ('Smart Watch', 249, 75, 'Fitbit', 'Wearables');

ollama serve

python db_assistant.py