def __init__()

in UI/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)