mirror of
https://github.com/Zippland/Bubbles.git
synced 2026-01-19 09:41:17 +08:00
45 lines
1.2 KiB
Python
45 lines
1.2 KiB
Python
#! /usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
|
|
import os
|
|
import google.generativeai as genai
|
|
|
|
|
|
class BardAssistant:
|
|
def __init__(self, conf: dict) -> None:
|
|
self._api_key = conf["api_key"]
|
|
self._model_name = conf["model_name"]
|
|
self._prompt = conf['prompt']
|
|
self._proxy = conf['proxy']
|
|
|
|
genai.configure(api_key=self._api_key)
|
|
self._bard = genai.GenerativeModel(self._model_name)
|
|
|
|
def __repr__(self):
|
|
return 'BardAssistant'
|
|
|
|
@staticmethod
|
|
def value_check(conf: dict) -> bool:
|
|
if conf:
|
|
if conf.get("api_key") and conf.get("model_name") and conf.get("prompt"):
|
|
return True
|
|
return False
|
|
|
|
def get_answer(self, msg: str, sender: str = None) -> str:
|
|
response = self._bard.generate_content([{'role': 'user', 'parts': [msg]}])
|
|
return response.text
|
|
|
|
|
|
if __name__ == "__main__":
|
|
from configuration import Config
|
|
config = Config().BardAssistant
|
|
if not config:
|
|
exit(0)
|
|
|
|
bard_assistant = BardAssistant(config)
|
|
if bard_assistant._proxy:
|
|
os.environ['HTTP_PROXY'] = bard_assistant._proxy
|
|
os.environ['HTTPS_PROXY'] = bard_assistant._proxy
|
|
rsp = bard_assistant.get_answer(bard_assistant._prompt)
|
|
print(rsp)
|