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

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"

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)

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
 

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

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"
+]

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

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)

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)

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.