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

# select Author with matching first_name and last_name
statement = select(Author).where(Author.first_name == first_name, Author.last_name == last_name)

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()

from utils import (
   get_book_id_by_title,
   get_books_for_author,
   get_num_in_stock_for_book,
)

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",
           },
       },
   },
)

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]

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],
)

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)