Gemini AI common code and utils

This commit is contained in:
ajaysi
2025-04-11 17:47:55 +05:30
parent b556edb989
commit e41be5789a
18 changed files with 590 additions and 262 deletions

View File

@@ -7,27 +7,29 @@
import os
from pathlib import Path
from dotenv import load_dotenv
from google.api_core import retry
import google.generativeai as genai
from pprint import pprint
from loguru import logger
import sys
from ..gpt_providers.text_generation.main_text_generation import llm_text_gen
def generate_with_retry(model, prompt):
def generate_with_retry(prompt, system_prompt=None):
"""
Generates content from the model with retry handling for errors.
Generates content using the llm_text_gen function with retry handling for errors.
Parameters:
model (GenerativeModel): The generative model to use for content generation.
prompt (str): The prompt to generate content from.
system_prompt (str, optional): Custom system prompt to use instead of the default one.
Returns:
str: The generated content.
"""
try:
# FIXME: Need a progress bar here.
return model.generate_content(prompt, request_options={'retry':retry.Retry()})
# Use llm_text_gen instead of directly calling the model
return llm_text_gen(prompt, system_prompt)
except Exception as e:
print(f"Error generating content: {e}")
logger.error(f"Error generating content: {e}")
return ""
@@ -36,11 +38,12 @@ def ai_essay_generator(essay_title, selected_essay_type, selected_education_leve
Write an Essay using prompt chaining and iterative generation.
Parameters:
persona (str): The persona statement for the author.
story_genre (str): The genre of the story.
characters (str): The characters in the story.
essay_title (str): The title or topic of the essay.
selected_essay_type (str): The type of essay to write.
selected_education_level (str): The education level of the target audience.
selected_num_pages (int): The number of pages or words for the essay.
"""
print(f"Starting to write Essay on {essay_title}..")
logger.info(f"Starting to write Essay on {essay_title}..")
try:
# Define persona and writing guidelines
guidelines = f'''\
@@ -127,59 +130,55 @@ def ai_essay_generator(essay_title, selected_essay_type, selected_education_leve
{guidelines}
'''
# Configure generative AI
load_dotenv(Path('../.env'))
genai.configure(api_key=os.getenv('GEMINI_API_KEY'))
# Initialize the generative model
model = genai.GenerativeModel('gemini-pro')
# Generate prompts
try:
premise = generate_with_retry(model, premise_prompt).text
print(f"The title of the Essay is: {premise}")
premise = generate_with_retry(premise_prompt)
logger.info(f"The title of the Essay is: {premise}")
except Exception as err:
print(f"Essay title Generation Error: {err}")
logger.error(f"Essay title Generation Error: {err}")
return
outline = generate_with_retry(model, outline_prompt.format(premise=premise)).text
print(f"The Outline of the essay is: {outline}\n\n")
outline = generate_with_retry(outline_prompt.format(premise=premise))
logger.info(f"The Outline of the essay is: {outline}\n\n")
if not outline:
print("Failed to generate Essay outline. Exiting...")
logger.error("Failed to generate Essay outline. Exiting...")
return
try:
starting_draft = generate_with_retry(model,
starting_prompt.format(premise=premise, outline=outline)).text
starting_draft = generate_with_retry(
starting_prompt.format(premise=premise, outline=outline))
pprint(starting_draft)
except Exception as err:
print(f"Failed to Generate Essay draft: {err}")
logger.error(f"Failed to Generate Essay draft: {err}")
return
try:
draft = starting_draft
continuation = generate_with_retry(model,
continuation_prompt.format(premise=premise, outline=outline, story_text=draft)).text
continuation = generate_with_retry(
continuation_prompt.format(premise=premise, outline=outline, story_text=draft))
pprint(continuation)
except Exception as err:
print(f"Failed to write the initial draft: {err}")
logger.error(f"Failed to write the initial draft: {err}")
# Add the continuation to the initial draft, keep building the story until we see 'IAMDONE'
try:
draft += '\n\n' + continuation
except Exception as err:
print(f"Failed as: {err} and {continuation}")
logger.error(f"Failed as: {err} and {continuation}")
while 'IAMDONE' not in continuation:
try:
continuation = generate_with_retry(model,
continuation_prompt.format(premise=premise, outline=outline, story_text=draft)).text
continuation = generate_with_retry(
continuation_prompt.format(premise=premise, outline=outline, story_text=draft))
draft += '\n\n' + continuation
except Exception as err:
print(f"Failed to continually write the Essay: {err}")
logger.error(f"Failed to continually write the Essay: {err}")
return
# Remove 'IAMDONE' and print the final story
final = draft.replace('IAMDONE', '').strip()
pprint(final)
return final
except Exception as e:
print(f"Main Essay writing: An error occurred: {e}")
logger.error(f"Main Essay writing: An error occurred: {e}")
return ""