diff --git a/mkdocs.yml b/mkdocs.yml index 9e39a3cc..ab173cba 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -97,6 +97,7 @@ nav: - DistilWhisperModel: "swarms/models/distilled_whisperx.md" - ElevenLabsText2SpeechTool: "swarms/models/elevenlabs.md" - OpenAITTS: "swarms/models/openai_tts.md" + - Gemini: "swarms/models/gemini.md" - swarms.structs: - Overview: "swarms/structs/overview.md" - AutoScaler: "swarms/swarms/autoscaler.md" diff --git a/playground/models/gemini.py b/playground/models/gemini.py new file mode 100644 index 00000000..6edffde5 --- /dev/null +++ b/playground/models/gemini.py @@ -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) diff --git a/swarms/models/base_multimodal_model.py b/swarms/models/base_multimodal_model.py index 396bd56f..c31931f2 100644 --- a/swarms/models/base_multimodal_model.py +++ b/swarms/models/base_multimodal_model.py @@ -1,6 +1,8 @@ import asyncio import base64 import concurrent.futures +import logging +import os import time from abc import abstractmethod from concurrent.futures import ThreadPoolExecutor @@ -96,10 +98,6 @@ class BaseMultiModalModel: self.meta_prompt = meta_prompt self.chat_history = [] - def __call__(self, task: str, img: str, *args, **kwargs): - """Run the model""" - return self.run(task, img, *args, **kwargs) - @abstractmethod def run( self, task: Optional[str], img: Optional[str], *args, **kwargs @@ -107,7 +105,21 @@ class BaseMultiModalModel: """Run the model""" 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""" pass diff --git a/swarms/models/gemini.py b/swarms/models/gemini.py index 9a052bb2..e9b29346 100644 --- a/swarms/models/gemini.py +++ b/swarms/models/gemini.py @@ -3,11 +3,14 @@ import subprocess as sp from pathlib import Path from dotenv import load_dotenv +from PIL import Image + from swarms.models.base_multimodal_model import BaseMultiModalModel try: import google.generativeai as genai + from google.generativeai.types import GenerationConfig except ImportError as error: print(f"Error importing google.generativeai: {error}") print("Please install the google.generativeai package") @@ -39,13 +42,24 @@ class Gemini(BaseMultiModalModel): """Gemini model Args: - BaseMultiModalModel (class): Base multimodal model class - model_name (str, optional): model name. Defaults to "gemini-pro". - gemini_api_key (str, optional): Gemini API key. Defaults to None. + model_name (str, optional): _description_. Defaults to "gemini-pro". + gemini_api_key (str, optional): _description_. Defaults to get_gemini_api_key_env. + 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: - run: run the Gemini model - process_img: process the image + run: Run the Gemini model + 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: @@ -59,20 +73,67 @@ class Gemini(BaseMultiModalModel): def __init__( self, - model_name: str = "gemini-pro", - gemini_api_key: str = get_gemini_api_key_env, + model_name: str = "gemini-pro-vision", + 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, **kwargs, ): super().__init__(model_name, *args, **kwargs) self.model_name = model_name 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 self.model = genai.GenerativeModel( 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( self, task: str = None, @@ -91,18 +152,33 @@ class Gemini(BaseMultiModalModel): """ try: 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( - 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 else: response = self.model.generate_content( task, *args, **kwargs ) - return response + return response.text except Exception as error: print(f"Error running Gemini model: {error}") + print(f"Please check the task and image: {task}, {img}") + raise error def process_img( self, @@ -158,3 +234,35 @@ class Gemini(BaseMultiModalModel): response1 = response.text print(response1) 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 diff --git a/swarms/telemetry/__init__.py b/swarms/telemetry/__init__.py new file mode 100644 index 00000000..2e33a86a --- /dev/null +++ b/swarms/telemetry/__init__.py @@ -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" +] \ No newline at end of file diff --git a/swarms/telemetry/log_all.py b/swarms/telemetry/log_all.py new file mode 100644 index 00000000..e8a1b06c --- /dev/null +++ b/swarms/telemetry/log_all.py @@ -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 diff --git a/swarms/telemetry/posthog_utils.py b/swarms/telemetry/posthog_utils.py new file mode 100644 index 00000000..d8bb51ec --- /dev/null +++ b/swarms/telemetry/posthog_utils.py @@ -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) \ No newline at end of file diff --git a/swarms/telemetry/user_utils.py b/swarms/telemetry/user_utils.py new file mode 100644 index 00000000..613d85cc --- /dev/null +++ b/swarms/telemetry/user_utils.py @@ -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) \ No newline at end of file diff --git a/tests/models/test_gemini.py b/tests/models/test_gemini.py index bbde19cb..c6f3e023 100644 --- a/tests/models/test_gemini.py +++ b/tests/models/test_gemini.py @@ -216,3 +216,98 @@ def test_gemini_run_mock_img_processing_exception( assert response is None mock_generate_content.assert_not_called() 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.