GreyDGL / PentestGPT

A GPT-empowered penetration testing tool
MIT License
7.28k stars 883 forks source link

Google AI Studio Model Free #239

Open l33tm3 opened 4 months ago

l33tm3 commented 4 months ago

Is your feature request related to a problem? Please describe. I think it would be a good idea to add the connection to Google's AI Studio since it is free and has models like Gemini 1.5 and 1.5 Pro as well as Gemma and 1.5 Flash

Describe the solution you'd like

"""
Install the Google AI Python SDK

$ pip install google-generativeai

See the getting started guide for more information:
https://ai.google.dev/gemini-api/docs/get-started/python
"""

import os

import google.generativeai as genai

genai.configure(api_key=os.environ["GEMINI_API_KEY"])

# Create the model
# See https://ai.google.dev/api/python/google/generativeai/GenerativeModel
generation_config = {
  "temperature": 0.9,
  "top_p": 1,
  "top_k": undefined,
  "max_output_tokens": 2048,
  "response_mime_type": "text/plain",
}

model = genai.GenerativeModel(
  model_name="gemini-1.5-flash",
  generation_config=generation_config,
  # safety_settings = Adjust safety settings
  # See https://ai.google.dev/gemini-api/docs/safety-settings
)

chat_session = model.start_chat(
  history=[
  ]
)

response = chat_session.send_message("")

print(response.text)

Link: https://aistudio.google.com/

l33tm3 commented 4 months ago

@GreyDGL

GreyDGL commented 2 months ago

I was not aware of this. Will take a look

Yellowwaves commented 2 months ago

@GreyDGL I did a program like yours but mine was automatic. I used the session system like this. And yep, you have 300$ on google ai for free.

import google.generativeai as genai

def initialize_chat(api_key, generation_config, safety_settings):
    genai.configure(api_key=api_key)
    model = genai.GenerativeModel(model_name='gemini-1.5-flash', generation_config=generation_config, safety_settings=safety_settings
    )
    return model.start_chat(history=[])

def send_and_print_message(session_name, chat, message): # Récuperation du session_name pour déboguer
    try:
        response = chat.send_message(message, stream=False) # stream false permet que le message s'envoie d'un coup
        response_text = ""
        for chunk in response:
                response_text += chunk.text
        logs(session_name, message, response_text)  # Enregistre le message et la réponse dans le log

        return response_text
    except Exception as e:
        print(f"An error occurred: {e}")
        return None

def create_chat_session(api_key):
    # Load API key and configurations
    if not api_key:
        raise ValueError("API key not found. Please set your GEMINI_API_KEY in the environment.")

    generation_config = {
        "temperature": 0.7,
        "top_p": 1,
        "top_k": 1,
        "max_output_tokens": 2048,
    }

    safety_settings = {
        "HARM_CATEGORY_HARASSMENT": "BLOCK_NONE",
        "HARM_CATEGORY_HATE_SPEECH": "BLOCK_NONE",
        "HARM_CATEGORY_SEXUALLY_EXPLICIT": "BLOCK_NONE",
        "HARM_CATEGORY_DANGEROUS_CONTENT": "BLOCK_NONE",
    }

    # Initialize and return the chat session
    return initialize_chat(api_key, generation_config, safety_settings)