def __init__()

in dbai_src/dbai.py [0:0]

• 26 lines of code
• 1 McCabe index (conditional complexity)

    def __init__(
            self,
            proj_id="proj-kous",
            dataset_id="Albertsons",
            tables_list=['camain_oracle_hcm', 'camain_ps']
            ):

        self.proj_id = proj_id
        self.dataset_id = dataset_id
        self.tables_list = tables_list

        self.sql_query_tool = Tool(
            function_declarations=[
                list_tables_func,
                get_table_metadata_func,
                sql_query_func,
                plot_chart_auto_func,
                # plot_chart_func,
            ],
        )

        self.agent = GenerativeModel("gemini-1.5-pro-001",
                                     generation_config={"temperature": 0.05},
                                     safety_settings=safety_settings,
                                     tools=[self.sql_query_tool],
                                     )

        self.bq_client = bigquery.Client(project=self.proj_id)
        self.system_prompt = """
        You are a fluent person who efficiently communicates with the user over
 different Database queries. Please always call the functions at your disposal
 whenever you need to know something, and do not reply unless you feel you have
 all information to answer the question satisfactorily.
 Only use information that you learn from BigQuery, do not make up information.
 Always use date or time functions instead of hard-coded values in SQL
 to reflect true current value.
        """
        self.load_metadata()

        vertexai.init(project=self.proj_id)