/
githubmirror
/
devika
Обзор
Документация
Войти
/
githubmirror
/
devika
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
src/llm/gemini_client.py
33 строки
2 KB
ayush rajgor
Improve and Fix: response parser, invalid response and others (#522)
02 май 2024, 17:32
Не верифицирован
02 май 2024, 17:32
cdfb782
Код
Авторство
О чём код?
import google.generativeai as genai from google.generativeai.types import HarmCategory, HarmBlockThreshold from src.config import Config class Gemini: def __init__(self): config = Config() api_key = config.get_gemini_api_key() genai.configure(api_key=api_key) def inference(self, model_id: str, prompt: str) -> str: config = genai.GenerationConfig(temperature=0) model = genai.GenerativeModel(model_id, generation_config=config) # Set safety settings for the request safety_settings = { HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE, HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE, # You can adjust other categories as needed } response = model.generate_content(prompt, safety_settings=safety_settings) try: # Check if the response contains text return response.text except ValueError: # If the response doesn't contain text, check if the prompt was blocked print("Prompt feedback:", response.prompt_feedback) # Also check the finish reason to see if the response was blocked print("Finish reason:", response.candidates[0].finish_reason) # If the finish reason was SAFETY, the safety ratings have more details print("Safety ratings:", response.candidates[0].safety_ratings) # Handle the error or return an appropriate message return "Error: Unable to generate content Gemini API"