04.12.2024
Artificial Intelligence meets Python SQLModel
Artificial Intelligence Leverages Database and API
Artificial Intelligence is being discussed everywhere - we at Blueshoe are no exception. In this blog post, we build a simpler yet impressive case with Google's Vertex AI and the Python package SQLModel. We demonstrate how databases and APIs can be integrated into AI models. The scenario should enable simply posing queries in natural language to our database and receiving answers.
Table of Contents
Vertex AI and SQLModel
The Vertex AI platform is a Google Cloud product that offers developers the opportunity to integrate generative AI into their applications. The possibilities are numerous: Input and output can be via text, audio, and video, and various AI models are available. For this blog post, we limit ourselves to text as input and output medium. Vertex AI Tutorials like this show how easy integration into projects can be.
SQLModel is a Python package based on Pydantic (a validation library) that enables storing Python data models in databases or reading from them. Essentially an ORM (=Object-relational mapping) on Pydantic basis, which is excellently suited for an AI database. Combining Artificial Intelligence with Python becomes particularly easy.
Data Model and Data Generation
We take the following scenario: We are a bookstore and maintain a database with books. It is known how many copies of each book are in stock. Each book has a known author.
The data model looks like this:
class Author(SQLModel, table=True):
id: Optional[int] = Field(default=None, primary_key=True)
first_name: str
last_name: str
birthday: date
class Book(SQLModel, table=True):
id: Optional[int] = Field(default=None, primary_key=True)
name: str
author_id: int = Field(default=None, foreign_key="author.id")
num_in_stock: int = 0
Books and authors are connected through a classic foreign key relationship. With SQLModel, queries to the database can now be easily formulated to build an AI API - e.g., selecting all authors with a specific first and last name:
# select Author with matching first_name and last_name
statement = select(Author).where(Author.first_name == first_name, Author.last_name == last_name)
So far, so good. Via Polyfactory, we will generate some test data:
class AuthorFactory(ModelFactory[Author]):
__model__ = Author
__faker__ = Faker(locale="de_DE")
id = Use(lambda: None)
@classmethod
def first_name(cls) -> str:
return cls.__faker__.first_name()
@classmethod
def last_name(cls) -> str:
return cls.__faker__.last_name()
class BookFactory(ModelFactory[Book]):
__model__ = Book
__faker__ = Faker(locale="de_DE")
id = Use(lambda: None)
@classmethod
def name(cls) -> str:
return cls.__faker__.catch_phrase()
def create_authors():
for i in range(0, 1000):
author = AuthorFactory.build()
with Session(engine) as session:
session.add(author)
session.commit()
for i in range(1, randrange(1, 10)):
create_book(author_id=author.id)
def create_book(author_id):
book = BookFactory.build(author_id=author_id)
with Session(engine) as session:
session.add(book)
session.commit()
# init sqlite
engine = create_engine("sqlite:///database.db")
SQLModel.metadata.create_all(engine)
# call create fcnt
create_authors()
And we already have 1000 different authors with between 1 and 10 books. The books get assigned complete availability quantities (number in stock).
Michael SchilonkaWe can also build AI solutions for you.
Connection with Vertex AI
First of all, it would have to be defined what functionality our agent should have access to. We assume that the agent should be able to retrieve books for a specific author and also know the number of copies in stock for a specific book.
The following helper functions are being added:
from utils import (
get_book_id_by_title,
get_books_for_author,
get_num_in_stock_for_book,
)
Vertex AI now needs information about the functions, what these do and which input parameters can be expected. This is done using FunctionDeclarations:
get_authors_for_book_func = FunctionDeclaration(
name=GET_BOOK_BY_AUTHOR,
description="Get a list of book names for an author.",
parameters={
"type": "object",
"properties": {
"first_name": {
"type": "string",
"description": "The first name of the author",
},
"last_name": {
"type": "string",
"description": "The last name of the author",
},
},
},
)
The function itself looks like the following:
def get_books_for_author(first_name: str, last_name: str):
with Session(engine) as session:
statement = select(Author).where(Author.first_name == first_name, Author.last_name == last_name)
author = session.exec(statement).first()
statement = select(Book).where(Book.author_id == author.id)
books = session.exec(statement=statement).all()
return [book.name for book in books]
The functions are combined into a "Tool" and made available to the model:
tools = Tool(
function_declarations=[
get_authors_for_book_func,
get_book_id_by_title_func,
get_num_in_stock_for_book_func,
],
)
model = GenerativeModel(
model_name="gemini-1.5-pro-002",
generation_config=GenerationConfig(temperature=0),
tools=[tools],
)
The model can now classify a user's input and provide an assessment of which functions should be used. This can then be accessed:
if function_calls:
api_responses = []
for func in function_calls:
if func.name == GET_BOOK_BY_AUTHOR:
api_responses.append(
{
"name": func.name,
"content": get_books_for_author(
first_name=func.args["first_name"],
last_name=func.args["last_name"],
),
}
)
elif func.name == GET_BOOK_ID_BY_TITLE:
api_responses.append(
{
"name": func.name,
"content": get_book_id_by_title(
title=func.args["title"],
),
}
)
elif func.name == GET_NUM_IN_STOCK_FOR_BOOK:
api_responses.append(
{
"name": func.name,
"content": get_num_in_stock_for_book(
title=func.args["title"],
),
}
)
# Return the API response to Gemini
for api_response in api_responses:
name = api_response.pop("name")
response = chat_session.send_message(
[
Part.from_function_response(
name=name,
response=api_response,
)
]
)
print(response.text)
The model thus extracts the parameters from the user's input. These can then be used in the function call.
Now the chat can start!
Which books are written by Guenther Hendriks?
Guenther Hendriks has written the following books: Digitized optimal circuit, Future-proofed content-based groupware, Universal clear-thinking Local Area Network, Digitized scalable service-desk, and Innovative full-range protocol.
How many copies of the first of those books are available?
There are 2089 copies of "Digitized optimal circuit" in stock.
Conclusion
The Vertex AI API generally works well. Currently, the response always takes a few seconds, which depending on the use case might potentially be too long. The interface is intuitive and easy to operate. For some inquiries, no response comes from Vertex AI, which then must be handled.
Vertex AI makes it possible to make existing APIs and databases accessible via natural language, which works cool and fast, although it still needs some refinement. It allows an efficient use of the AI database and integration into various systems. Especially for developers who want to program AI with Python, this is an exciting solution.
ADRESS
