Browse Source

added gemini support

pull/390/head
xssdoctor 6 months ago
parent
commit
e2231b3504
  1. 5
      installer/client/cli/fabric.py
  2. 90
      installer/client/cli/utils.py
  3. 42
      poetry.lock
  4. 1
      pyproject.toml

5
installer/client/cli/fabric.py

@ -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

90
installer/client/cli/utils.py

@ -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

@ -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"

1
pyproject.toml

@ -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…
Cancel
Save