Add support for Gemini 1.5 Flash via Gemini API

app.py

@@ -5,7 +5,6 @@ import datetime
 import logging
 import pathlib
 import random
-import sys
 import tempfile
 from typing import List, Union
 
@@ -17,9 +16,6 @@ from langchain_community.chat_message_histories import StreamlitChatMessageHisto
 from langchain_core.messages import HumanMessage
 from langchain_core.prompts import ChatPromptTemplate
 
-sys.path.append('..')
-sys.path.append('../..')
-
 from global_config import GlobalConfig
 from helpers import llm_helper, pptx_helper, text_helper
 
@@ -54,6 +50,60 @@ def _get_prompt_template(is_refinement: bool) -> str:
     return template
 
 
+def are_all_inputs_valid(
+        user_prompt: str,
+        selected_provider: str,
+        selected_model: str,
+        user_key: str,
+) -> bool:
+    """
+    Validate user input and LLM selection.
+
+    :param user_prompt: The prompt.
+    :param selected_provider: The LLM provider.
+    :param selected_model: Name of the model.
+    :param user_key: User-provided API key.
+    :return: `True` if all inputs "look" OK; `False` otherwise.
+    """
+
+    if not text_helper.is_valid_prompt(user_prompt):
+        handle_error(
+            'Not enough information provided!'
+            ' Please be a little more descriptive and type a few words'
+            ' with a few characters :)',
+            False
+        )
+        return False
+
+    if not selected_provider or not selected_model:
+        handle_error('No valid LLM provider and/or model name found!', False)
+        return False
+
+    if not llm_helper.is_valid_llm_provider_model(selected_provider, selected_model, user_key):
+        handle_error(
+            'The LLM settings do not look correct. Make sure that an API key/access token'
+            ' is provided if the selected LLM requires it.',
+            False
+        )
+        return False
+
+    return True
+
+
+def handle_error(error_msg: str, should_log: bool):
+    """
+    Display an error message in the app.
+
+    :param error_msg: The error message to be displayed.
+    :param should_log: If `True`, log the message.
+    """
+
+    if should_log:
+        logger.error(error_msg)
+
+    st.error(error_msg)
+
+
 APP_TEXT = _load_strings()
 
 # Session variables
@@ -80,11 +130,8 @@ with st.sidebar:
     llm_provider_to_use = st.sidebar.selectbox(
         label='2: Select an LLM to use:',
         options=[f'{k} ({v["description"]})' for k, v in GlobalConfig.VALID_MODELS.items()],
-        index=0,
-        help=(
-            'LLM provider codes:\n\n'
-            '- **[hf]**: Hugging Face Inference Endpoint\n'
-        ),
+        index=GlobalConfig.DEFAULT_MODEL_INDEX,
+        help=GlobalConfig.LLM_PROVIDER_HELP,
     ).split(' ')[0]
 
     # The API key/access token
@@ -123,53 +170,28 @@ def set_up_chat_ui():
     with st.expander('Usage Instructions'):
         st.markdown(GlobalConfig.CHAT_USAGE_INSTRUCTIONS)
 
-    st.info(
-        'If you like SlideDeck AI, please consider leaving a heart ❤️ on the'
-        ' [Hugging Face Space](https://huggingface.co/spaces/barunsaha/slide-deck-ai/) or'
-        ' a star ⭐ on [GitHub](https://github.com/barun-saha/slide-deck-ai).'
-        ' Your [feedback](https://forms.gle/JECFBGhjvSj7moBx9) is appreciated.'
-    )
-
-    # view_messages = st.expander('View the messages in the session state')
-
-    st.chat_message('ai').write(
-        random.choice(APP_TEXT['ai_greetings'])
-    )
+    st.info(APP_TEXT['like_feedback'])
+    st.chat_message('ai').write(random.choice(APP_TEXT['ai_greetings']))
 
     history = StreamlitChatMessageHistory(key=CHAT_MESSAGES)
-
-    if _is_it_refinement():
-        template = _get_prompt_template(is_refinement=True)
-    else:
-        template = _get_prompt_template(is_refinement=False)
-
-    prompt_template = ChatPromptTemplate.from_template(template)
+    prompt_template = ChatPromptTemplate.from_template(
+        _get_prompt_template(
+            is_refinement=_is_it_refinement()
+        )
+    )
 
     # Since Streamlit app reloads at every interaction, display the chat history
     # from the save session state
     for msg in history.messages:
-        msg_type = msg.type
-        if msg_type == 'user':
-            st.chat_message(msg_type).write(msg.content)
-        else:
-            st.chat_message(msg_type).code(msg.content, language='json')
+        st.chat_message(msg.type).code(msg.content, language='json')
 
     if prompt := st.chat_input(
         placeholder=APP_TEXT['chat_placeholder'],
         max_chars=GlobalConfig.LLM_MODEL_MAX_INPUT_LENGTH
     ):
-        if not text_helper.is_valid_prompt(prompt):
-            st.error(
-                'Not enough information provided!'
-                ' Please be a little more descriptive and type a few words'
-                ' with a few characters :)'
-            )
-            return
-
         provider, llm_name = llm_helper.get_provider_model(llm_provider_to_use)
 
-        if not provider or not llm_name:
-            st.error('No valid LLM provider and/or model name found!')
+        if not are_all_inputs_valid(prompt, provider, llm_name, api_key_token):
             return
 
         logger.info(
@@ -178,72 +200,76 @@ def set_up_chat_ui():
         )
         st.chat_message('user').write(prompt)
 
-        user_messages = _get_user_messages()
-        user_messages.append(prompt)
-        list_of_msgs = [
-            f'{idx + 1}. {msg}' for idx, msg in enumerate(user_messages)
-        ]
-        list_of_msgs = '\n'.join(list_of_msgs)
-
         if _is_it_refinement():
+            user_messages = _get_user_messages()
+            user_messages.append(prompt)
+            list_of_msgs = [
+                f'{idx + 1}. {msg}' for idx, msg in enumerate(user_messages)
+            ]
             formatted_template = prompt_template.format(
                 **{
-                    'instructions': list_of_msgs,
+                    'instructions': '\n'.join(list_of_msgs),
                     'previous_content': _get_last_response(),
                 }
             )
         else:
-            formatted_template = prompt_template.format(
-                **{
-                    'question': prompt,
-                }
-            )
+            formatted_template = prompt_template.format(**{'question': prompt})
 
         progress_bar = st.progress(0, 'Preparing to call LLM...')
         response = ''
 
         try:
-            for chunk in llm_helper.get_langchain_llm(
-                    provider=provider,
-                    model=llm_name,
-                    max_new_tokens=GlobalConfig.VALID_MODELS[llm_provider_to_use]['max_new_tokens'],
-                    api_key=api_key_token.strip(),
-            ).stream(formatted_template):
-                response += chunk
-
-                # Update the progress bar
-                progress_percentage = min(
-                    len(response) / GlobalConfig.VALID_MODELS[llm_provider_to_use]['max_new_tokens'], 0.95
+            llm = llm_helper.get_langchain_llm(
+                provider=provider,
+                model=llm_name,
+                max_new_tokens=GlobalConfig.VALID_MODELS[llm_provider_to_use]['max_new_tokens'],
+                api_key=api_key_token.strip(),
+            )
+
+            if not llm:
+                handle_error(
+                    'Failed to create an LLM instance! Make sure that you have selected the'
+                    ' correct model from the dropdown list and have provided correct API key'
+                    ' or access token.',
+                    False
                 )
+                return
+
+            for _ in llm.stream(formatted_template):
+                response += _
+
+                # Update the progress bar with an approx progress percentage
                 progress_bar.progress(
-                    progress_percentage,
+                    min(
+                        len(response) / GlobalConfig.VALID_MODELS[
+                            llm_provider_to_use
+                        ]['max_new_tokens'],
+                        0.95
+                    ),
                     text='Streaming content...this might take a while...'
                 )
         except requests.exceptions.ConnectionError:
-            msg = (
+            handle_error(
                 'A connection error occurred while streaming content from the LLM endpoint.'
                 ' Unfortunately, the slide deck cannot be generated. Please try again later.'
-                ' Alternatively, try selecting a different LLM from the dropdown list.'
+                ' Alternatively, try selecting a different LLM from the dropdown list.',
+                True
             )
-            logger.error(msg)
-            st.error(msg)
             return
         except huggingface_hub.errors.ValidationError as ve:
-            msg = (
+            handle_error(
                 f'An error occurred while trying to generate the content: {ve}'
-                '\nPlease try again with a significantly shorter input text.'
+                '\nPlease try again with a significantly shorter input text.',
+                True
             )
-            logger.error(msg)
-            st.error(msg)
             return
         except Exception as ex:
-            msg = (
+            handle_error(
                 f'An unexpected error occurred while generating the content: {ex}'
                 '\nPlease try again later, possibly with different inputs.'
-                ' Alternatively, try selecting a different LLM from the dropdown list.'
+                ' Alternatively, try selecting a different LLM from the dropdown list.',
+                True
             )
-            logger.error(msg)
-            st.error(msg)
             return
 
         history.add_user_message(prompt)
@@ -252,25 +278,20 @@ def set_up_chat_ui():
         # The content has been generated as JSON
         # There maybe trailing ``` at the end of the response -- remove them
         # To be careful: ``` may be part of the content as well when code is generated
-        response_cleaned = text_helper.get_clean_json(response)
-
+        response = text_helper.get_clean_json(response)
         logger.info(
-            'Cleaned JSON response:: original length: %d | cleaned length: %d',
-            len(response), len(response_cleaned)
+            'Cleaned JSON length: %d', len(response)
         )
-        # logger.debug('Cleaned JSON: %s', response_cleaned)
 
         # Now create the PPT file
         progress_bar.progress(
             GlobalConfig.LLM_PROGRESS_MAX,
             text='Finding photos online and generating the slide deck...'
         )
-        path = generate_slide_deck(response_cleaned)
         progress_bar.progress(1.0, text='Done!')
-
         st.chat_message('ai').code(response, language='json')
 
-        if path:
+        if path := generate_slide_deck(response):
             _display_download_button(path)
 
         logger.info(
@@ -291,44 +312,35 @@ def generate_slide_deck(json_str: str) -> Union[pathlib.Path, None]:
     try:
         parsed_data = json5.loads(json_str)
     except ValueError:
-        st.error(
-            'Encountered error while parsing JSON...will fix it and retry'
-        )
-        logger.error(
-            'Caught ValueError: trying again after repairing JSON...'
+        handle_error(
+            'Encountered error while parsing JSON...will fix it and retry',
+            True
         )
         try:
             parsed_data = json5.loads(text_helper.fix_malformed_json(json_str))
         except ValueError:
-            st.error(
+            handle_error(
                 'Encountered an error again while fixing JSON...'
                 'the slide deck cannot be created, unfortunately ☹'
-                '\nPlease try again later.'
+                '\nPlease try again later.',
+                True
             )
-            logger.error(
-                'Caught ValueError: failed to repair JSON!'
-            )
-
             return None
     except RecursionError:
-        st.error(
-            'Encountered an error while parsing JSON...'
+        handle_error(
+            'Encountered a recursion error while parsing JSON...'
             'the slide deck cannot be created, unfortunately ☹'
-            '\nPlease try again later.'
+            '\nPlease try again later.',
+            True
         )
-        logger.error('Caught RecursionError while parsing JSON. Cannot generate the slide deck!')
-
         return None
     except Exception:
-        st.error(
+        handle_error(
             'Encountered an error while parsing JSON...'
             'the slide deck cannot be created, unfortunately ☹'
-            '\nPlease try again later.'
-        )
-        logger.error(
-            'Caught ValueError: failed to parse JSON!'
+            '\nPlease try again later.',
+            True
         )
-
         return None
 
     if DOWNLOAD_FILE_KEY in st.session_state:

global_config.py

@@ -17,17 +17,32 @@ class GlobalConfig:
     A data class holding the configurations.
     """
 
-    VALID_PROVIDERS = {'hf'}
+    PROVIDER_HUGGING_FACE = 'hf'
+    PROVIDER_GOOGLE_GEMINI = 'gg'
+    VALID_PROVIDERS = {PROVIDER_HUGGING_FACE, PROVIDER_GOOGLE_GEMINI}
     VALID_MODELS = {
+        '[gg]gemini-1.5-flash-002': {
+            'description': 'faster response',
+            'max_new_tokens': 8192,
+            'paid': True,
+        },
         '[hf]mistralai/Mistral-7B-Instruct-v0.2': {
             'description': 'faster, shorter',
-            'max_new_tokens': 8192
+            'max_new_tokens': 8192,
+            'paid': False,
         },
         '[hf]mistralai/Mistral-Nemo-Instruct-2407': {
             'description': 'longer response',
-            'max_new_tokens': 12228
+            'max_new_tokens': 10240,
+            'paid': False,
         },
     }
+    LLM_PROVIDER_HELP = (
+        'LLM provider codes:\n\n'
+        '- **[gg]**: Google Gemini API\n'
+        '- **[hf]**: Hugging Face Inference Endpoint\n'
+    )
+    DEFAULT_MODEL_INDEX = 1
     LLM_MODEL_TEMPERATURE = 0.2
     LLM_MODEL_MIN_OUTPUT_LENGTH = 100
     LLM_MODEL_MAX_INPUT_LENGTH = 400  # characters

helpers/llm_helper.py

@@ -1,13 +1,18 @@
+"""
+Helper functions to access LLMs.
+"""
 import logging
 import re
+import sys
 from typing import Tuple, Union
 
 import requests
 from requests.adapters import HTTPAdapter
 from urllib3.util import Retry
-
 from langchain_community.llms.huggingface_endpoint import HuggingFaceEndpoint
-from langchain_core.language_models import LLM
+from langchain_core.language_models import BaseLLM
+
+sys.path.append('..')
 
 from global_config import GlobalConfig
 
@@ -49,30 +54,26 @@ def get_provider_model(provider_model: str) -> Tuple[str, str]:
     return '', ''
 
 
-def get_hf_endpoint(repo_id: str, max_new_tokens: int, api_key: str = '') -> LLM:
+def is_valid_llm_provider_model(provider: str, model: str, api_key: str) -> bool:
     """
-    Get an LLM via the HuggingFaceEndpoint of LangChain.
-
-    :param repo_id: The model name.
-    :param max_new_tokens: The max new tokens to generate.
-    :param api_key: [Optional] Hugging Face access token.
-    :return: The HF LLM inference endpoint.
+    Verify whether LLM settings are proper.
+    This function does not verify whether `api_key` is correct. It only confirms that the key has
+    at least five characters. Key verification is done when the LLM is created.
+
+    :param provider: Name of the LLM provider.
+    :param model: Name of the model.
+    :param api_key: The API key or access token.
+    :return: `True` if the settings "look" OK; `False` otherwise.
     """
 
-    logger.debug('Getting LLM via HF endpoint: %s', repo_id)
+    if not provider or not model or provider not in GlobalConfig.VALID_PROVIDERS:
+        return False
+
+    if provider in [GlobalConfig.PROVIDER_GOOGLE_GEMINI, ]:
+        if not api_key or len(api_key) < 5:
+            return False
 
-    return HuggingFaceEndpoint(
-        repo_id=repo_id,
-        max_new_tokens=max_new_tokens,
-        top_k=40,
-        top_p=0.95,
-        temperature=GlobalConfig.LLM_MODEL_TEMPERATURE,
-        repetition_penalty=1.03,
-        streaming=True,
-        huggingfacehub_api_token=api_key or GlobalConfig.HUGGINGFACEHUB_API_TOKEN,
-        return_full_text=False,
-        stop_sequences=['</s>'],
-    )
+    return True
 
 
 def get_langchain_llm(
@@ -80,22 +81,19 @@ def get_langchain_llm(
         model: str,
         max_new_tokens: int,
         api_key: str = ''
-) -> Union[LLM, None]:
+) -> Union[BaseLLM, None]:
     """
     Get an LLM based on the provider and model specified.
 
     :param provider: The LLM provider. Valid values are `hf` for Hugging Face.
-    :param model:
-    :param max_new_tokens:
-    :param api_key:
-    :return:
+    :param model: The name of the LLM.
+    :param max_new_tokens: The maximum number of tokens to generate.
+    :param api_key: API key or access token to use.
+    :return: An instance of the LLM or `None` in case of any error.
     """
-    if not provider or not model or provider not in GlobalConfig.VALID_PROVIDERS:
-        return None
 
-    if provider == 'hf':
+    if provider == GlobalConfig.PROVIDER_HUGGING_FACE:
         logger.debug('Getting LLM via HF endpoint: %s', model)
-
         return HuggingFaceEndpoint(
             repo_id=model,
             max_new_tokens=max_new_tokens,
@@ -109,6 +107,27 @@ def get_langchain_llm(
             stop_sequences=['</s>'],
         )
 
+    if provider == GlobalConfig.PROVIDER_GOOGLE_GEMINI:
+        from google.generativeai.types.safety_types import HarmBlockThreshold, HarmCategory
+        from langchain_google_genai import GoogleGenerativeAI
+
+        return GoogleGenerativeAI(
+            model=model,
+            temperature=GlobalConfig.LLM_MODEL_TEMPERATURE,
+            max_tokens=max_new_tokens,
+            timeout=None,
+            max_retries=2,
+            google_api_key=api_key,
+            safety_settings={
+                HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT:
+                    HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
+                HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
+                HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
+                HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT:
+                    HarmBlockThreshold.BLOCK_LOW_AND_ABOVE
+            }
+        )
+
     return None
 
 

requirements.txt

@@ -10,6 +10,7 @@ pydantic==2.9.1
 langchain~=0.3.7
 langchain-core~=0.3.0
 langchain-community==0.3.0
+langchain-google-genai==2.0.6
 streamlit~=1.38.0
 
 python-pptx
@@ -19,7 +20,6 @@ requests~=2.32.3
 
 transformers~=4.44.0
 torch==2.4.0
-langchain-community
 
 urllib3~=2.2.1
 lxml~=4.9.3

strings.json

@@ -33,5 +33,6 @@
         "Looks like you have a looming deadline. Can I help you get started with your slide deck?",
         "Hello! What topic do you have on your mind today?"
     ],
-    "chat_placeholder": "Write the topic or instructions here"
+    "chat_placeholder": "Write the topic or instructions here",
+    "like_feedback": "If you like SlideDeck AI, please consider leaving a heart ❤\uFE0F on the [Hugging Face Space](https://huggingface.co/spaces/barunsaha/slide-deck-ai/) or a star ⭐ on [GitHub](https://github.com/barun-saha/slide-deck-ai). Your [feedback](https://forms.gle/JECFBGhjvSj7moBx9) is appreciated."
 }