Gemini API : Error when calling .generate_content(prompt)

Question

I am rather new to developing with LLMs. I am trying to write an RAG based app in Streamlit with Gemini's API using google embeddings. The goal of the app is to retrieve information from PDFs (which is my RAG's knowledge base) when answering the user's requests. I used Streamlit for the interfacing, FAISS for the vector DB.

Here's the code with which I make the calling :

# Get the user's question 
question = st.chat_input("Ask here")

# Handle the user's question
if question:
    vectordb = st.session_state.get("vectordb", None)
    if not vectordb:
        with st.message("assistant"):
            st.write("You need to provide a PDF")
            st.stop()

    # Search the vectordb for similar content to the user's question
    search_results = vectordb.similarity_search(question, k=3)
    # sources of result
    pdf_extract = "/n ".join([result.page_content for result in search_results])

    # Update the prompt with the pdf extract
    prompt[0] = {
        "role": "system",
        "content": prompt_template.format(pdf_extract=pdf_extract),
    }

    # Add the user's question to the prompt and display it
    prompt.append({"role": "user", "content": question})

    with st.chat_message("user"):
        st.write(question)

    # Display an empty assistant message while waiting for the response
    with st.chat_message("assistant"):
        botmsg = st.empty()

    response = []
    result = ""
    
    response = model.generate_content(prompt)
    botmsg.write(result)

model is

model = genai.GenerativeModel('gemini-pro')

After setting up API's key and all. I did not include this part of the code as I believe is not relevant to the issue.

When I execute my script, I face this error :

KeyError: "Could not recognize the intended type of the `dict`. A `Content` should have a 'parts' key. A `Part` should have a 'inline_data' or a 'text' key. A `Blob` should have 'mime_type' and 'data' keys. Got keys: ['role', 'content']"

I had this previous code which uses openAI's API and OpenAI Embedding model, and it works :

   for chunk in st.session_state.chat.send_message(prompt):
        text = chunk.choices[0].get("delta", {}).get("content")
        if text is not None:
            response.append(text)
            result = "".join(response).strip()
            botmsg.write(result)

Gemini's corresponding code is this : ( present in the previous code snippet )

    response = model.generate_content(prompt)
    botmsg.write(result)

How can I solve this ? I don't understand the error, I am pretty sure there is something I am missing in the code. If someone encountered this, or has an idea about why this is occurring, I'd appreciate.

Thank you in advance.

Answer 1

Here's a basic test script to get you goin. Hope this helps :)

from vertexai.preview.generative_models import GenerativeModel, GenerationConfig, Part, Content, HarmCategory, HarmBlockThreshold, ResponseBlockedError


class Gemini:
    
    def __init__(self, streaming: bool = False):
    
        # Load the model
        self.model = GenerativeModel("gemini-pro")
        self.streaming = streaming

        # Generation config
        self.config = GenerationConfig(
            temperature=0.0,
            top_p=0.8,
            top_k=32,
            candidate_count=1,
            max_output_tokens=2048,
        )

        # Safety config
        self.safety_config = {   
            HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_ONLY_HIGH,
            HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_ONLY_HIGH,
            HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_ONLY_HIGH,
            HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_ONLY_HIGH,
            HarmCategory.HARM_CATEGORY_UNSPECIFIED: HarmBlockThreshold.BLOCK_ONLY_HIGH,
        }
        

    def invoke(self, prompt: dict):
        
        # Format the user input as a Content object
        user_input = Content(
            role="user", 
            parts=[
                Part.from_text(prompt["input"]),
            ]
        )
        
        response = self.model.generate_content(
            contents=user_input,
            generation_config=self.config,
            safety_settings=self.safety_config,
            stream=self.streaming, 
            tools=[]
        )
        
        return response
    

# Initialise the model
model = Gemini(streaming=False)

# Test the model
generated_response = model.invoke({"input": "ELI5 in 50 words or less: Why is the sky blue?"})


# Print the response with basic exception handling
try:
    if hasattr(generated_response.candidates[0].content.parts[0], 'text'):
        print(generated_response.candidates[0].content.parts[0].text)
except ValueError as e:
    print("An error occured: ", e)