[FEAT][Gemini] [FEAT][Telemtry]

pull/299/head
Kye 1 year ago
parent 4ece24851f
commit 5b4d7c485c

@ -97,6 +97,7 @@ nav:
- DistilWhisperModel: "swarms/models/distilled_whisperx.md" - DistilWhisperModel: "swarms/models/distilled_whisperx.md"
- ElevenLabsText2SpeechTool: "swarms/models/elevenlabs.md" - ElevenLabsText2SpeechTool: "swarms/models/elevenlabs.md"
- OpenAITTS: "swarms/models/openai_tts.md" - OpenAITTS: "swarms/models/openai_tts.md"
- Gemini: "swarms/models/gemini.md"
- swarms.structs: - swarms.structs:
- Overview: "swarms/structs/overview.md" - Overview: "swarms/structs/overview.md"
- AutoScaler: "swarms/swarms/autoscaler.md" - AutoScaler: "swarms/swarms/autoscaler.md"

@ -0,0 +1,14 @@
from swarms.models.gemini import Gemini
# Initialize the model
model = Gemini(
gemini_api_key="A",
)
# Establish the prompt and image
task = "What is your name"
img = "images/github-banner-swarms.png"
# Run the model
out = model.run("What is your name?", img)
print(out)

@ -1,6 +1,8 @@
import asyncio import asyncio
import base64 import base64
import concurrent.futures import concurrent.futures
import logging
import os
import time import time
from abc import abstractmethod from abc import abstractmethod
from concurrent.futures import ThreadPoolExecutor from concurrent.futures import ThreadPoolExecutor
@ -96,10 +98,6 @@ class BaseMultiModalModel:
self.meta_prompt = meta_prompt self.meta_prompt = meta_prompt
self.chat_history = [] self.chat_history = []
def __call__(self, task: str, img: str, *args, **kwargs):
"""Run the model"""
return self.run(task, img, *args, **kwargs)
@abstractmethod @abstractmethod
def run( def run(
self, task: Optional[str], img: Optional[str], *args, **kwargs self, task: Optional[str], img: Optional[str], *args, **kwargs
@ -107,7 +105,21 @@ class BaseMultiModalModel:
"""Run the model""" """Run the model"""
pass pass
async def arun(self, task: str, img: str): def __call__(
self, task: str = None, img: str = None, *args, **kwargs
):
"""Call the model
Args:
task (str): _description_
img (str): _description_
Returns:
_type_: _description_
"""
return self.run(task, img, *args, **kwargs)
async def arun(self, task: str, img: str, *args, **kwargs):
"""Run the model asynchronously""" """Run the model asynchronously"""
pass pass

@ -3,11 +3,14 @@ import subprocess as sp
from pathlib import Path from pathlib import Path
from dotenv import load_dotenv from dotenv import load_dotenv
from PIL import Image
from swarms.models.base_multimodal_model import BaseMultiModalModel from swarms.models.base_multimodal_model import BaseMultiModalModel
try: try:
import google.generativeai as genai import google.generativeai as genai
from google.generativeai.types import GenerationConfig
except ImportError as error: except ImportError as error:
print(f"Error importing google.generativeai: {error}") print(f"Error importing google.generativeai: {error}")
print("Please install the google.generativeai package") print("Please install the google.generativeai package")
@ -39,13 +42,24 @@ class Gemini(BaseMultiModalModel):
"""Gemini model """Gemini model
Args: Args:
BaseMultiModalModel (class): Base multimodal model class model_name (str, optional): _description_. Defaults to "gemini-pro".
model_name (str, optional): model name. Defaults to "gemini-pro". gemini_api_key (str, optional): _description_. Defaults to get_gemini_api_key_env.
gemini_api_key (str, optional): Gemini API key. Defaults to None. return_safety (bool, optional): _description_. Defaults to False.
candidates (bool, optional): _description_. Defaults to False.
stream (bool, optional): _description_. Defaults to False.
candidate_count (int, optional): _description_. Defaults to 1.
stop_sequence ([type], optional): _description_. Defaults to ['x'].
max_output_tokens (int, optional): _description_. Defaults to 100.
temperature (float, optional): _description_. Defaults to 0.9.
Methods: Methods:
run: run the Gemini model run: Run the Gemini model
process_img: process the image process_img: Process the image
chat: Chat with the Gemini model
list_models: List the Gemini models
stream_tokens: Stream the tokens
process_img_pil: Process img
Examples: Examples:
@ -59,20 +73,67 @@ class Gemini(BaseMultiModalModel):
def __init__( def __init__(
self, self,
model_name: str = "gemini-pro", model_name: str = "gemini-pro-vision",
gemini_api_key: str = get_gemini_api_key_env, gemini_api_key: str = None,
return_safety: bool = False,
candidates: bool = False,
stream: bool = False,
candidate_count: int = 1,
stop_sequence=["x"],
max_output_tokens: int = 100,
temperature: float = 0.9,
*args, *args,
**kwargs, **kwargs,
): ):
super().__init__(model_name, *args, **kwargs) super().__init__(model_name, *args, **kwargs)
self.model_name = model_name self.model_name = model_name
self.gemini_api_key = gemini_api_key self.gemini_api_key = gemini_api_key
self.safety = return_safety
self.candidates = candidates
self.stream = stream
self.candidate_count = candidate_count
self.stop_sequence = stop_sequence
self.max_output_tokens = max_output_tokens
self.temperature = temperature
# Prepare the generation config
self.generation_config = GenerationConfig(
candidate_count=candidate_count,
# stop_sequence=stop_sequence,
max_output_tokens=max_output_tokens,
temperature=temperature,
)
# Initialize the model # Initialize the model
self.model = genai.GenerativeModel( self.model = genai.GenerativeModel(
model_name, *args, **kwargs model_name, *args, **kwargs
) )
# Check for the key
if self.gemini_api_key is None:
raise ValueError("Please provide a Gemini API key")
def system_prompt(
self,
system_prompt: str = None,
task: str = None,
*args,
**kwargs,
):
"""System prompt
Args:
system_prompt (str, optional): _description_. Defaults to None.
"""
PROMPT = f"""
{system_prompt}
{task}
"""
return PROMPT
def run( def run(
self, self,
task: str = None, task: str = None,
@ -91,18 +152,33 @@ class Gemini(BaseMultiModalModel):
""" """
try: try:
if img: if img:
process_img = self.process_img(img, *args, **kwargs) # process_img = self.process_img(img, *args, **kwargs)
process_img = self.process_img_pil(img)
response = self.model.generate_content( response = self.model.generate_content(
content=[task, process_img], *args, **kwargs contents=[task, process_img],
generation_config=self.generation_config,
stream=self.stream,
*args,
**kwargs,
) )
# if self.candidates:
# return response.candidates
# elif self.safety:
# return response.safety
# else:
# return response.text
return response.text return response.text
else: else:
response = self.model.generate_content( response = self.model.generate_content(
task, *args, **kwargs task, *args, **kwargs
) )
return response return response.text
except Exception as error: except Exception as error:
print(f"Error running Gemini model: {error}") print(f"Error running Gemini model: {error}")
print(f"Please check the task and image: {task}, {img}")
raise error
def process_img( def process_img(
self, self,
@ -158,3 +234,35 @@ class Gemini(BaseMultiModalModel):
response1 = response.text response1 = response.text
print(response1) print(response1)
response = chat.send_message(img, *args, **kwargs) response = chat.send_message(img, *args, **kwargs)
def list_models(self) -> str:
"""List the Gemini models
Returns:
str: _description_
"""
for m in genai.list_models():
if "generateContent" in m.supported_generation_methods:
print(m.name)
def stream_tokens(self, content: str = None):
"""Stream the tokens
Args:
content (t, optional): _description_. Defaults to None.
"""
for chunk in content:
print(chunk.text)
print("_" * 80)
def process_img_pil(self, img: str = None):
"""Process img
Args:
img (str, optional): _description_. Defaults to None.
Returns:
_type_: _description_
"""
img = Image.open(img)
return img

@ -0,0 +1,14 @@
from swarms.telemetry.log_all import log_all_calls, log_calls
from swarms.telemetry.posthog_utils import log_activity_posthog
from swarms.telemetry.user_utils import generate_user_id, get_machine_id, get_system_info, generate_unique_identifier
__all__ = [
"log_all_calls",
"log_calls",
"log_activity_posthog",
"generate_user_id",
"get_machine_id",
"get_system_info",
"generate_unique_identifier"
]

@ -0,0 +1,33 @@
import logging
import types
# Set up logging
logging.basicConfig(level=logging.INFO)
# Log all calls to functions in this module
def log_all_calls(module):
"""
Decorate all functions of a module to log calls to them.
"""
for name, obj in vars(module).items():
if isinstance(obj, types.FunctionType):
setattr(module, name, log_calls(obj))
# Log all calls to a function
def log_calls(func):
"""
Decorate a function to log calls to it.
"""
def wrapper(*args, **kwargs):
logging.info(
f"Calling function {func.__name__} with args {args} and"
f" kwargs {kwargs}"
)
result = func(*args, **kwargs)
logging.info(f"Function {func.__name__} returned {result}")
return result
return wrapper

@ -0,0 +1,74 @@
import functools
import os
from dotenv import load_dotenv
from posthog import Posthog
from swarms.telemetry.user_utils import generate_unique_identifier
# Load environment variables
load_dotenv()
# Initialize Posthog client
def init_posthog(debug: bool = True, *args, **kwargs):
"""Initialize Posthog client.
Args:
debug (bool, optional): Whether to enable debug mode. Defaults to True.
"""
api_key = os.getenv("POSTHOG_API_KEY")
host = os.getenv("POSTHOG_HOST")
posthog = Posthog(api_key, host=host, *args, **kwargs)
if debug:
posthog.debug = True
return posthog
def log_activity_posthog(event_name: str, **event_properties):
"""Log activity to Posthog.
Args:
event_name (str): Name of the event to log.
**event_properties: Properties of the event to log.
Examples:
>>> from swarms.telemetry.posthog_utils import log_activity_posthog
>>> @log_activity_posthog("test_event", test_property="test_value")
... def test_function():
... print("Hello, world!")
>>> test_function()
Hello, world!
>>> # Check Posthog dashboard for event "test_event" with property
>>> # "test_property" set to "test_value".
"""
def decorator_log_activity(func):
@functools.wraps(func)
def wrapper_log_activity(*args, **kwargs):
result = func(*args, **kwargs)
# Assuming you have a way to get the user id
distinct_user_id = generate_unique_identifier()
# Capture the event
init_posthog.capture(
distinct_user_id, event_name, event_properties
)
return result
return wrapper_log_activity
return decorator_log_activity
@log_activity_posthog('function_executed', function_name='my_function')
def my_function():
# Function logic here
return "Function executed successfully!"
out = my_function()
print(out)

@ -0,0 +1,60 @@
import hashlib
import platform
import uuid
import socket
# Helper functions
def generate_user_id():
"""Generate user id
Returns:
_type_: _description_
"""
return str(uuid.uuid4())
def get_machine_id():
"""Get machine id
Returns:
_type_: _description_
"""
raw_id = platform.node()
hashed_id = hashlib.sha256(
raw_id.encode()
).hexdigest()
return hashed_id
def get_system_info():
"""
Gathers basic system information.
Returns:
dict: A dictionary containing system-related information.
"""
info = {
"platform": platform.system(),
"platform_release": platform.release(),
"platform_version": platform.version(),
"architecture": platform.machine(),
"hostname": socket.gethostname(),
"ip_address": socket.gethostbyname(socket.gethostname()),
"mac_address": ':'.join(['{:02x}'.format((uuid.getnode() >> elements) & 0xff) for elements in range(0, 2 * 6, 8)][::-1]),
"processor": platform.processor(),
"python_version": platform.python_version()
}
return info
def generate_unique_identifier():
"""Generate unique identifier
Returns:
str: unique id
"""
system_info = get_system_info()
unique_id = uuid.uuid5(
uuid.NAMESPACE_DNS,
str(system_info)
)
return str(unique_id)

@ -216,3 +216,98 @@ def test_gemini_run_mock_img_processing_exception(
assert response is None assert response is None
mock_generate_content.assert_not_called() mock_generate_content.assert_not_called()
mock_process_img.assert_called_with(img=img) mock_process_img.assert_called_with(img=img)
# Test Gemini run method with mocked image processing and different exception
@patch("swarms.models.gemini.Gemini.process_img")
@patch("swarms.models.gemini.genai.GenerativeModel.generate_content")
def test_gemini_run_mock_img_processing_different_exception(
mock_generate_content,
mock_process_img,
mock_gemini_api_key,
mock_genai_model,
):
model = Gemini()
task = "A dog"
img = "dog.png"
mock_process_img.side_effect = ValueError("Test exception")
with pytest.raises(ValueError):
model.run(task=task, img=img)
mock_generate_content.assert_not_called()
mock_process_img.assert_called_with(img=img)
# Test Gemini run method with mocked image processing and no exception
@patch("swarms.models.gemini.Gemini.process_img")
@patch("swarms.models.gemini.genai.GenerativeModel.generate_content")
def test_gemini_run_mock_img_processing_no_exception(
mock_generate_content,
mock_process_img,
mock_gemini_api_key,
mock_genai_model,
):
model = Gemini()
task = "A bird"
img = "bird.png"
mock_generate_content.return_value = "A bird is flying"
response = model.run(task=task, img=img)
assert response == "A bird is flying"
mock_generate_content.assert_called_once()
mock_process_img.assert_called_with(img=img)
# Test Gemini chat method
@patch("swarms.models.gemini.Gemini.chat")
def test_gemini_chat(mock_chat):
model = Gemini()
mock_chat.return_value = "Hello, Gemini!"
response = model.chat("Hello, Gemini!")
assert response == "Hello, Gemini!"
mock_chat.assert_called_once()
# Test Gemini list_models method
@patch("swarms.models.gemini.Gemini.list_models")
def test_gemini_list_models(mock_list_models):
model = Gemini()
mock_list_models.return_value = ["model1", "model2"]
response = model.list_models()
assert response == ["model1", "model2"]
mock_list_models.assert_called_once()
# Test Gemini stream_tokens method
@patch("swarms.models.gemini.Gemini.stream_tokens")
def test_gemini_stream_tokens(mock_stream_tokens):
model = Gemini()
mock_stream_tokens.return_value = ["token1", "token2"]
response = model.stream_tokens()
assert response == ["token1", "token2"]
mock_stream_tokens.assert_called_once()
# Test Gemini process_img_pil method
@patch("swarms.models.gemini.Gemini.process_img_pil")
def test_gemini_process_img_pil(mock_process_img_pil):
model = Gemini()
img = "bird.png"
mock_process_img_pil.return_value = "processed image"
response = model.process_img_pil(img)
assert response == "processed image"
mock_process_img_pil.assert_called_with(img)
# Repeat the above tests for different scenarios or different methods in your Gemini class
# until you have 15 tests in total.

Loading…
Cancel
Save