mirror of
https://github.com/danielmiessler/fabric
synced 2024-11-08 07:11:06 +00:00
added gemini support
This commit is contained in:
xssdoctor
parent
4ccba83dd3
commit
e2231b3504
@ -158,7 +158,7 @@ def main():
|
||||
print("No patterns found")
|
||||
sys.exit()
|
||||
if args.listmodels:
|
||||
gptmodels, localmodels, claudemodels = standalone.fetch_available_models()
|
||||
gptmodels, localmodels, claudemodels, googlemodels = standalone.fetch_available_models()
|
||||
print("GPT Models:")
|
||||
for model in gptmodels:
|
||||
print(model)
|
||||
@ -168,6 +168,9 @@ def main():
|
||||
print("\nClaude Models:")
|
||||
for model in claudemodels:
|
||||
print(model)
|
||||
print("\nGoogle Models:")
|
||||
for model in googlemodels:
|
||||
print(model)
|
||||
sys.exit()
|
||||
if args.text is not None:
|
||||
text = args.text
|
||||
|
||||
@ -53,12 +53,14 @@ class Standalone:
|
||||
if not self.model:
|
||||
self.model = 'gpt-4-turbo-preview'
|
||||
self.claude = False
|
||||
sorted_gpt_models, ollamaList, claudeList = self.fetch_available_models()
|
||||
sorted_gpt_models, ollamaList, claudeList, googleList = self.fetch_available_models()
|
||||
self.sorted_gpt_models = sorted_gpt_models
|
||||
self.ollamaList = ollamaList
|
||||
self.claudeList = claudeList
|
||||
self.googleList = googleList
|
||||
self.local = self.model in ollamaList
|
||||
self.claude = self.model in claudeList
|
||||
self.google = self.model in googleList
|
||||
|
||||
async def localChat(self, messages, host=''):
|
||||
from ollama import AsyncClient
|
||||
@ -144,6 +146,47 @@ class Standalone:
|
||||
session.save_to_session(
|
||||
system, user, message.content[0].text, self.args.session)
|
||||
|
||||
async def googleChat(self, system, user, copy=False):
|
||||
import google.generativeai as genai
|
||||
self.googleApiKey = os.environ["GOOGLE_API_KEY"]
|
||||
genai.configure(api_key=self.googleApiKey)
|
||||
model = genai.GenerativeModel(
|
||||
model_name=self.model, system_instruction=system)
|
||||
response = model.generate_content(user)
|
||||
print(response.text)
|
||||
if copy:
|
||||
pyperclip.copy(response.text)
|
||||
if self.args.output:
|
||||
with open(self.args.output, "w") as f:
|
||||
f.write(response.text)
|
||||
if self.args.session:
|
||||
from .helper import Session
|
||||
session = Session()
|
||||
session.save_to_session(
|
||||
system, user, response.text, self.args.session)
|
||||
|
||||
async def googleStream(self, system, user, copy=False):
|
||||
import google.generativeai as genai
|
||||
buffer = ""
|
||||
self.googleApiKey = os.environ["GOOGLE_API_KEY"]
|
||||
genai.configure(api_key=self.googleApiKey)
|
||||
model = genai.GenerativeModel(
|
||||
model_name=self.model, system_instruction=system)
|
||||
response = model.generate_content(user, stream=True)
|
||||
for chunk in response:
|
||||
buffer += chunk.text
|
||||
print(chunk.text)
|
||||
if copy:
|
||||
pyperclip.copy(buffer)
|
||||
if self.args.output:
|
||||
with open(self.args.output, "w") as f:
|
||||
f.write(buffer)
|
||||
if self.args.session:
|
||||
from .helper import Session
|
||||
session = Session()
|
||||
session.save_to_session(
|
||||
system, user, buffer, self.args.session)
|
||||
|
||||
def streamMessage(self, input_data: str, context="", host=''):
|
||||
""" Stream a message and handle exceptions.
|
||||
|
||||
@ -209,6 +252,10 @@ class Standalone:
|
||||
elif self.claude:
|
||||
from anthropic import AsyncAnthropic
|
||||
asyncio.run(self.claudeStream(system, user_message))
|
||||
elif self.google:
|
||||
if system == "":
|
||||
system = " "
|
||||
asyncio.run(self.googleStream(system, user_message['content']))
|
||||
else:
|
||||
stream = self.client.chat.completions.create(
|
||||
model=self.model,
|
||||
@ -315,6 +362,10 @@ class Standalone:
|
||||
asyncio.run(self.localChat(messages))
|
||||
elif self.claude:
|
||||
asyncio.run(self.claudeChat(system, user_message))
|
||||
elif self.google:
|
||||
if system == "":
|
||||
system = " "
|
||||
asyncio.run(self.googleChat(system, user_message['content']))
|
||||
else:
|
||||
response = self.client.chat.completions.create(
|
||||
model=self.model,
|
||||
@ -354,6 +405,7 @@ class Standalone:
|
||||
def fetch_available_models(self):
|
||||
gptlist = []
|
||||
fullOllamaList = []
|
||||
googleList = []
|
||||
if "CLAUDE_API_KEY" in os.environ:
|
||||
claudeList = ['claude-3-opus-20240229', 'claude-3-sonnet-20240229',
|
||||
'claude-3-haiku-20240307', 'claude-2.1']
|
||||
@ -389,8 +441,16 @@ class Standalone:
|
||||
fullOllamaList.append(model['name'])
|
||||
except:
|
||||
fullOllamaList = []
|
||||
try:
|
||||
import google.generativeai as genai
|
||||
genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
|
||||
for m in genai.list_models():
|
||||
if 'generateContent' in m.supported_generation_methods:
|
||||
googleList.append(m.name)
|
||||
except:
|
||||
googleList = []
|
||||
|
||||
return gptlist, fullOllamaList, claudeList
|
||||
return gptlist, fullOllamaList, claudeList, googleList
|
||||
|
||||
def get_cli_input(self):
|
||||
""" aided by ChatGPT; uses platform library
|
||||
@ -523,6 +583,7 @@ class Setup:
|
||||
self.env_file = os.path.join(self.config_directory, ".env")
|
||||
self.gptlist = []
|
||||
self.fullOllamaList = []
|
||||
self.googleList = []
|
||||
self.claudeList = ['claude-3-opus-20240229']
|
||||
load_dotenv(self.env_file)
|
||||
try:
|
||||
@ -612,6 +673,31 @@ class Setup:
|
||||
with open(self.env_file, "w") as f:
|
||||
f.write(f"CLAUDE_API_KEY={claude_key}\n")
|
||||
|
||||
def google_key(self, google_key):
|
||||
""" Set the Google API key in the environment file.
|
||||
|
||||
Args:
|
||||
google_key (str): The API key to be set.
|
||||
|
||||
Returns:
|
||||
None
|
||||
|
||||
Raises:
|
||||
OSError: If the environment file does not exist or cannot be accessed.
|
||||
"""
|
||||
google_key = google_key.strip()
|
||||
if os.path.exists(self.env_file) and google_key:
|
||||
with open(self.env_file, "r") as f:
|
||||
lines = f.readlines()
|
||||
with open(self.env_file, "w") as f:
|
||||
for line in lines:
|
||||
if "GOOGLE_API_KEY" not in line:
|
||||
f.write(line)
|
||||
f.write(f"GOOGLE_API_KEY={google_key}\n")
|
||||
elif google_key:
|
||||
with open(self.env_file, "w") as f:
|
||||
f.write(f"GOOGLE_API_KEY={google_key}\n")
|
||||
|
||||
def youtube_key(self, youtube_key):
|
||||
""" Set the YouTube API key in the environment file.
|
||||
|
||||
|
||||
42
poetry.lock
generated
42
poetry.lock
generated
@ -1737,6 +1737,23 @@ gitdb = ">=4.0.1,<5"
|
||||
doc = ["sphinx (==4.3.2)", "sphinx-autodoc-typehints", "sphinx-rtd-theme", "sphinxcontrib-applehelp (>=1.0.2,<=1.0.4)", "sphinxcontrib-devhelp (==1.0.2)", "sphinxcontrib-htmlhelp (>=2.0.0,<=2.0.1)", "sphinxcontrib-qthelp (==1.0.3)", "sphinxcontrib-serializinghtml (==1.1.5)"]
|
||||
test = ["coverage[toml]", "ddt (>=1.1.1,!=1.4.3)", "mock", "mypy", "pre-commit", "pytest (>=7.3.1)", "pytest-cov", "pytest-instafail", "pytest-mock", "pytest-sugar", "typing-extensions"]
|
||||
|
||||
[[package]]
|
||||
name = "google-ai-generativelanguage"
|
||||
version = "0.6.3"
|
||||
description = "Google Ai Generativelanguage API client library"
|
||||
optional = false
|
||||
python-versions = ">=3.7"
|
||||
files = [
|
||||
{file = "google-ai-generativelanguage-0.6.3.tar.gz", hash = "sha256:10a11f1e1bb8470ff50030c1acd729b3aba7a29ade2c30cf1d1c917291366c67"},
|
||||
{file = "google_ai_generativelanguage-0.6.3-py3-none-any.whl", hash = "sha256:55a6698f6c9cbbfde5f9cd288073b6941dd9e3e6dc2176dfa3197f9a4c489895"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
google-api-core = {version = ">=1.34.1,<2.0.dev0 || >=2.11.dev0,<3.0.0dev", extras = ["grpc"]}
|
||||
google-auth = ">=2.14.1,<2.24.0 || >2.24.0,<2.25.0 || >2.25.0,<3.0.0dev"
|
||||
proto-plus = ">=1.22.3,<2.0.0dev"
|
||||
protobuf = ">=3.19.5,<3.20.0 || >3.20.0,<3.20.1 || >3.20.1,<4.21.0 || >4.21.0,<4.21.1 || >4.21.1,<4.21.2 || >4.21.2,<4.21.3 || >4.21.3,<4.21.4 || >4.21.4,<4.21.5 || >4.21.5,<5.0.0dev"
|
||||
|
||||
[[package]]
|
||||
name = "google-api-core"
|
||||
version = "2.18.0"
|
||||
@ -2036,6 +2053,29 @@ files = [
|
||||
[package.extras]
|
||||
testing = ["pytest"]
|
||||
|
||||
[[package]]
|
||||
name = "google-generativeai"
|
||||
version = "0.5.3"
|
||||
description = "Google Generative AI High level API client library and tools."
|
||||
optional = false
|
||||
python-versions = ">=3.9"
|
||||
files = [
|
||||
{file = "google_generativeai-0.5.3-py3-none-any.whl", hash = "sha256:a74509ee219601c74c0561eb4e1c9af6a88594c7dd098d30a18c6592afe62bd9"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
google-ai-generativelanguage = "0.6.3"
|
||||
google-api-core = "*"
|
||||
google-api-python-client = "*"
|
||||
google-auth = ">=2.15.0"
|
||||
protobuf = "*"
|
||||
pydantic = "*"
|
||||
tqdm = "*"
|
||||
typing-extensions = "*"
|
||||
|
||||
[package.extras]
|
||||
dev = ["Pillow", "absl-py", "black", "ipython", "nose2", "pandas", "pytype", "pyyaml"]
|
||||
|
||||
[[package]]
|
||||
name = "google-resumable-media"
|
||||
version = "2.7.0"
|
||||
@ -7061,4 +7101,4 @@ testing = ["coverage (>=5.0.3)", "zope.event", "zope.testing"]
|
||||
[metadata]
|
||||
lock-version = "2.0"
|
||||
python-versions = ">=3.10,<3.13"
|
||||
content-hash = "f8ccf1782ae31b7b9834dc1ff59cfba4a215f349f2fff87df5300eecc8941528"
|
||||
content-hash = "20b4784977b09b2d3ce41679356d7096938b70d3b326fcc54837043e83538e15"
|
||||
|
||||
@ -30,6 +30,7 @@ jwt = "^1.3.1"
|
||||
flask = "^3.0.2"
|
||||
helpers = "^0.2.0"
|
||||
praisonai = "^0.0.18"
|
||||
google-generativeai = "^0.5.3"
|
||||
|
||||
[tool.poetry.group.cli.dependencies]
|
||||
pyyaml = "^6.0.1"
|
||||
|
||||
Loading…
Reference in New Issue
Block a user