ai commands

Author: unytics

bigfun/bigfunctions.py

@@ -200,6 +200,15 @@ def doc(self):
             self.config['use_case'] = open(self.use_case_filename, encoding='utf-8').read()
         return BIGFUNCTION_DOC_TEMPLATE.render(**self.config)
 
+    def generate_description_improvement(self):
+        prompt = '\n\n'.join([
+            'Improve the description of this bigquery function. Return only the description.',
+            'FUNCTION YAML DEFINITION:',
+            open(self.filename, encoding='utf-8').read(),
+        ])
+        description = generate_content(prompt)
+        print(description)
+
     def generate_use_case(self, overwrite_if_exists=False):
         if os.path.isfile(self.use_case_filename):
             if not overwrite_if_exists:
@@ -320,6 +329,18 @@ def parse_readme(self):
 
         return title, readme, frontmatter
 
+    def generate_readme(self):
+        prompt = '''
+        BigFunctions is a catalog of BigQuery functions. Functions are categorized in folders. Generate the description of the folder containing the following functions:
+
+
+        '''.replace('\n        ', '')
+        prompt += '\n\n--------------------\n\n '.join(['FUNCTION `' + b.name + '``: ' + b.config['description'] for b in self.bigfunctions])
+        print(prompt)
+        readme = generate_content(prompt)
+        with open(f'{self.path}/README.md', 'w', encoding='utf-8') as file:
+            file.write(readme)
+
     @property
     def doc(self):
         return DOC_FOLDER_TEMPLATE.render(**self.dict)

bigfun/cli.py

@@ -178,6 +178,30 @@ def generate_use_case(bigfunction):
         bigfunction.generate_use_case()
 
 
+@docs.command()
+@click.argument('bigfunction')
+def generate_description_improvement(bigfunction):
+    """
+    Generate markdown files for documentation from yaml bigfunctions files
+    """
+    bigfunctions = [bigfunction.strip() for bigfunction in bigfunction.split(',')]
+    if bigfunction == 'ALL':
+        bigfunctions = bf.BIGFUNCTIONS
+    for name in bigfunctions:
+        bigfunction = bf.BigFunction(name)
+        bigfunction.generate_description_improvement()
+
+
+@docs.command()
+@click.argument('folder_path')
+def generate_folder_readme(folder_path):
+    """
+    Generate README.md for bigfunctions folder
+    """
+    folder = bf.Folder(folder_path)
+    folder.generate_readme()
+
+
 @docs.command()
 @click.option('--config', default='config.yaml', help='Path to the config file')
 @click.option('--port', default='8000', help='Port to serve docs on')

bigfun/utils.py

@@ -377,32 +377,9 @@ def download(url, destination_filename):
 
 
 def generate_content(prompt):
-    import vertexai
-    from vertexai.generative_models import GenerativeModel, SafetySetting
-
-    safety_settings = [
-        SafetySetting(
-            category=SafetySetting.HarmCategory.HARM_CATEGORY_HATE_SPEECH,
-            threshold=SafetySetting.HarmBlockThreshold.OFF
-        ),
-        SafetySetting(
-            category=SafetySetting.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
-            threshold=SafetySetting.HarmBlockThreshold.OFF
-        ),
-        SafetySetting(
-            category=SafetySetting.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
-            threshold=SafetySetting.HarmBlockThreshold.OFF
-        ),
-        SafetySetting(
-            category=SafetySetting.HarmCategory.HARM_CATEGORY_HARASSMENT,
-            threshold=SafetySetting.HarmBlockThreshold.OFF
-        ),
-    ]
-
-    vertexai.init()
-    model = GenerativeModel("gemini-1.5-pro-002")
-    response = model.generate_content(
-        prompt,
-        safety_settings=safety_settings,
-    )
+    import google.auth
+    import google.genai
+    _, project = google.auth.default()
+    genai = google.genai.Client(vertexai=True, project=project, location='europe-west1')
+    response = genai.models.generate_content(model='gemini-2.0-flash-001', contents=prompt)
     return response.text