lollms-client 1.5.6__py3-none-any.whl → 1.7.10__py3-none-any.whl

lollms_client/llm_bindings/gemini/__init__.py

@@ -1,6 +1,7 @@
-# bindings/gemini/binding.py
+# bindings/gemini/__init__.py
 import base64
 import os
+import re
 from io import BytesIO
 from pathlib import Path
 from typing import Optional, Callable, List, Union, Dict
@@ -16,7 +17,7 @@ import pipmaster as pm
 pm.ensure_packages(["google-generativeai", "pillow", "tiktoken", "protobuf"])
 
 import google.generativeai as genai
-from PIL import Image, ImageDraw # ImageDraw is used in the test script below
+from PIL import Image, ImageDraw 
 import tiktoken
 
 BindingName = "GeminiBinding"
@@ -43,7 +44,7 @@ class GeminiBinding(LollmsLLMBinding):
             service_key (str): Google AI Studio API key.
         """
         super().__init__(BindingName, **kwargs)
-        self.model_name = kwargs.get("model_name", None)
+        self.model_name = kwargs.get("model_name", "gemini-1.5-pro-latest")
         self.service_key = kwargs.get("service_key", None)
 
         if not self.service_key:
@@ -91,35 +92,34 @@ class GeminiBinding(LollmsLLMBinding):
                      split:Optional[bool]=False, 
                      user_keyword:Optional[str]="!@>user:",
                      ai_keyword:Optional[str]="!@>assistant:",
+                     think: Optional[bool] = False,
+                     reasoning_effort: Optional[str] = "low", # low, medium, high
+                     reasoning_summary: Optional[bool] = False, # auto
                      ) -> Union[str, dict]:
         """
         Generate text using the Gemini model.
-
-        Args:
-            prompt (str): The input prompt for text generation.
-            images (Optional[List[str]]): List of image file paths or base64 strings.
-            system_prompt (str): The system prompt to guide the model.
-            ... other LollmsLLMBinding parameters ...
-
-        Returns:
-            Union[str, dict]: Generated text or error dictionary.
         """
         if not self.client:
             return {"status": False, "error": "Gemini client not initialized."}
 
-        # Gemini uses 'system_instruction' for GenerativeModel, not part of the regular message list.
-        model = self.client.GenerativeModel(
-            model_name=self.model_name,
-            system_instruction=system_prompt if system_prompt else None
-        )
+        # Handle 'think' parameter logging
+        if think and "thinking" not in str(self.model_name).lower():
+             ASCIIColors.info(f"Thinking requested but model '{self.model_name}' may not be a thinking model. Proceeding.")
+
+        # Gemini uses 'system_instruction' for GenerativeModel
+        try:
+            model = self.client.GenerativeModel(
+                model_name=self.model_name,
+                system_instruction=system_prompt if system_prompt else None
+            )
+        except Exception as e:
+            return {"status": False, "error": f"Failed to initialize GenerativeModel: {e}"}
 
         generation_config = self.get_generation_config(temperature, top_p, top_k, n_predict)
 
         # Prepare content for the API call
         content_parts = []
         if split:
-            # Note: The 'split' logic for Gemini should ideally build a multi-turn history,
-            # but for `generate_text`, we'll treat the last user part as the main prompt.
             discussion_messages = self.split_discussion(prompt, user_keyword, ai_keyword)
             if discussion_messages:
                 last_message = discussion_messages[-1]['content']
@@ -135,6 +135,9 @@ class GeminiBinding(LollmsLLMBinding):
                     if is_image_path(image_data):
                         img = Image.open(image_data)
                     else: # Assume base64
+                        if image_data.startswith("data:image"):
+                             # Remove prefix if present
+                             image_data = image_data.split(",")[1]
                         img = Image.open(BytesIO(base64.b64decode(image_data)))
                     content_parts.append(img)
                 except Exception as e:
@@ -155,17 +158,15 @@ class GeminiBinding(LollmsLLMBinding):
                     try:
                         chunk_text = chunk.text
                     except ValueError:
-                        # Handle potential empty parts in the stream
                         chunk_text = ""
 
                     if chunk_text:
                         full_response_text += chunk_text
                         if streaming_callback:
                             if not streaming_callback(chunk_text, MSG_TYPE.MSG_TYPE_CHUNK):
-                                break # Callback requested stop
+                                break 
                 return full_response_text
             else:
-                # Check for safety blocks
                 if response.prompt_feedback.block_reason:
                     error_msg = f"Content blocked due to: {response.prompt_feedback.block_reason.name}"
                     ASCIIColors.warning(error_msg)
@@ -177,6 +178,133 @@ class GeminiBinding(LollmsLLMBinding):
             trace_exception(ex)
             return {"status": False, "error": error_message}
 
+    def generate_from_messages(self,
+                        messages: List[Dict],
+                        n_predict: Optional[int] = None,
+                        stream: Optional[bool] = None,
+                        temperature: Optional[float] = None,
+                        top_k: Optional[int] = None,
+                        top_p: Optional[float] = None,
+                        repeat_penalty: Optional[float] = None,
+                        repeat_last_n: Optional[int] = None,
+                        seed: Optional[int] = None,
+                        n_threads: Optional[int] = None,
+                        ctx_size: int | None = None,
+                        streaming_callback: Optional[Callable[[str, MSG_TYPE], None]] = None,
+                        think: Optional[bool] = False,
+                        reasoning_effort: Optional[bool] = "low", 
+                        reasoning_summary: Optional[bool] = "auto",
+                        **kwargs
+                        ) -> Union[str, dict]:
+        """
+        Generate content using a list of messages. This is the low-level method 
+        that handles chat history structure for Gemini.
+        """
+        if not self.client:
+            return {"status": False, "error": "Gemini client not initialized."}
+
+        gen_config = self.get_generation_config(temperature, top_p, top_k, n_predict)
+
+        system_instruction = None
+        gemini_contents = []
+
+        # Parse messages
+        for msg in messages:
+            role = msg.get("role", "user")
+            content = msg.get("content", "")
+            
+            # Gemini specific role mapping
+            if role == "system":
+                # Gemini takes system prompt at initialization, not in the content list
+                if isinstance(content, str):
+                    system_instruction = content
+                elif isinstance(content, list):
+                     # Extract text from list if system prompt is complex (rare)
+                     text_parts = [p.get("text", "") for p in content if p.get("type") == "text"]
+                     system_instruction = "\n".join(text_parts)
+                continue
+            
+            gemini_role = "model" if role == "assistant" else "user"
+            parts = []
+
+            # Parse content (Text or List of multimodal)
+            if isinstance(content, str):
+                parts.append(content)
+            elif isinstance(content, list):
+                for item in content:
+                    item_type = item.get("type")
+                    if item_type == "text":
+                        parts.append(item.get("text", ""))
+                    elif item_type in ["input_image", "image_url"]:
+                        # Extract Base64
+                        base64_data = None
+                        url_data = item.get("image_url", item.get("input_image"))
+                        
+                        if isinstance(url_data, dict):
+                            # Format: {"url": "data:image/..."} or {"base64": "..."}
+                            if "base64" in url_data:
+                                base64_data = url_data["base64"]
+                            elif "url" in url_data:
+                                base64_data = url_data["url"]
+                        elif isinstance(url_data, str):
+                            base64_data = url_data
+
+                        if base64_data:
+                            # Clean "data:image/png;base64," prefix
+                            if "base64," in base64_data:
+                                base64_data = base64_data.split("base64,")[1]
+                            try:
+                                img = Image.open(BytesIO(base64.b64decode(base64_data)))
+                                parts.append(img)
+                            except Exception as e:
+                                ASCIIColors.warning(f"Failed to decode image in message: {e}")
+
+            if parts:
+                gemini_contents.append({"role": gemini_role, "parts": parts})
+
+        # Initialize Model with system instruction
+        try:
+            model = self.client.GenerativeModel(
+                model_name=self.model_name,
+                system_instruction=system_instruction
+            )
+        except Exception as e:
+            return {"status": False, "error": f"Failed to initialize GenerativeModel: {e}"}
+
+        full_response_text = ""
+
+        try:
+            # Generate content based on the full history
+            response = model.generate_content(
+                contents=gemini_contents,
+                generation_config=gen_config,
+                stream=stream
+            )
+
+            if stream:
+                for chunk in response:
+                    try:
+                        chunk_text = chunk.text
+                    except ValueError:
+                        chunk_text = "" # Safety filter or empty chunk
+                    
+                    if chunk_text:
+                        full_response_text += chunk_text
+                        if streaming_callback:
+                            if not streaming_callback(chunk_text, MSG_TYPE.MSG_TYPE_CHUNK):
+                                break
+                return full_response_text
+            else:
+                if response.prompt_feedback.block_reason:
+                    error_msg = f"Content blocked due to: {response.prompt_feedback.block_reason.name}"
+                    return {"status": False, "error": error_msg}
+                return response.text
+
+        except Exception as ex:
+            error_message = f"An unexpected error occurred with Gemini API: {str(ex)}"
+            trace_exception(ex)
+            return {"status": False, "error": error_message}
+
     def chat(self,
              discussion: LollmsDiscussion,
              branch_tip_id: Optional[str] = None,
@@ -190,7 +318,10 @@ class GeminiBinding(LollmsLLMBinding):
              seed: Optional[int] = None,
              n_threads: Optional[int] = None,
              ctx_size: Optional[int] = None,
-             streaming_callback: Optional[Callable[[str, MSG_TYPE], None]] = None
+             streaming_callback: Optional[Callable[[str, MSG_TYPE], None]] = None,
+             think: Optional[bool] = False,
+             reasoning_effort: Optional[str] = "low", # low, medium, high
+             reasoning_summary: Optional[bool] = False, # auto
              ) -> Union[str, dict]:
         """
         Conduct a chat session with the Gemini model using a LollmsDiscussion object.
@@ -198,22 +329,17 @@ class GeminiBinding(LollmsLLMBinding):
         if not self.client:
              return {"status": "error", "message": "Gemini client not initialized."}
 
-        # 1. Manually export discussion to Gemini's format.
-        # Gemini uses 'user' and 'model' roles.
-        # The system prompt is handled separately at model initialization.
         system_prompt = discussion.system_prompt
         messages = discussion.get_messages(branch_tip_id)
         
         history = []
         for msg in messages:
-            role = 'user' if msg.sender_type == "user" else 'assistant'
+            role = 'user' if msg.sender_type == "user" else 'model' 
             
-            # Handle multimodal content in the message
             content_parts = []
             if msg.content:
                 content_parts.append(msg.content)
             
-            # Check for images associated with this message
             if msg.images:
                 for file_path in msg.images:
                      if is_image_path(file_path):
@@ -221,39 +347,61 @@ class GeminiBinding(LollmsLLMBinding):
                             content_parts.append(Image.open(file_path))
                          except Exception as e:
                             ASCIIColors.warning(f"Could not load image {file_path}: {e}")
+                     else:
+                        try:
+                            b64_data = file_path
+                            if b64_data.startswith("data:image"):
+                                b64_data = b64_data.split(",")[1]
+                            content_parts.append(Image.open(BytesIO(base64.b64decode(b64_data))))
+                        except:
+                            pass
 
             if content_parts:
                 history.append({'role': role, 'parts': content_parts})
         
-        model = self.client.GenerativeModel(
-            model_name=self.model_name,
-            system_instruction=system_prompt
-        )
+        try:
+            model = self.client.GenerativeModel(
+                model_name=self.model_name,
+                system_instruction=system_prompt if system_prompt else None
+            )
+        except Exception as e:
+            return {"status": "error", "message": f"Failed to initialize GenerativeModel: {e}"}
         
-        # History must not be empty and should not contain consecutive roles of the same type.
-        # We also need to separate the final prompt from the history.
+        # Organize history for chat session
+        # ChatSession expects history excluding the very last message if we use send_message
         if not history:
             return {"status": "error", "message": "Cannot start chat with an empty discussion."}
 
-        chat_history = history[:-1] if len(history) > 1 else []
-        last_prompt_parts = history[-1]['parts']
+        # Filter and consolidate consecutive roles
+        consolidated_history = []
+        current_role = None
+        current_parts = []
 
-        # Ensure history is valid (no consecutive same roles)
-        valid_history = []
-        if chat_history:
-            valid_history.append(chat_history[0])
-            for i in range(1, len(chat_history)):
-                if chat_history[i]['role'] != chat_history[i-1]['role']:
-                    valid_history.append(chat_history[i])
+        for msg in history:
+            if msg['role'] == current_role:
+                current_parts.extend(msg['parts'])
+                consolidated_history[-1]['parts'] = current_parts
+            else:
+                current_role = msg['role']
+                current_parts = list(msg['parts']) 
+                consolidated_history.append({'role': current_role, 'parts': current_parts})
+
+        # Separate the last message for the send_message call
+        last_message = consolidated_history[-1]
+        chat_history = consolidated_history[:-1]
+
+        # Only 'user' can send_message in Gemini chat. 
+        # If the last message in history is from 'model', the flow is slightly broken for standard chat,
+        # but we assume the standard Lollms loop (User -> AI).
         
-        chat_session = model.start_chat(history=valid_history)
+        chat_session = model.start_chat(history=chat_history)
         
         generation_config = self.get_generation_config(temperature, top_p, top_k, n_predict)
         
         full_response_text = ""
         try:
             response = chat_session.send_message(
-                content=last_prompt_parts,
+                content=last_message['parts'],
                 generation_config=generation_config,
                 stream=stream
             )
@@ -284,11 +432,6 @@ class GeminiBinding(LollmsLLMBinding):
             return {"status": "error", "message": error_message}
             
     def tokenize(self, text: str) -> list:
-        """
-        Tokenize the input text.
-        Note: Gemini doesn't expose a public tokenizer API.
-        Using tiktoken for a rough estimate, NOT accurate for Gemini.
-        """
         try:
             encoding = tiktoken.get_encoding("cl100k_base")
             return encoding.encode(text)
@@ -296,10 +439,6 @@ class GeminiBinding(LollmsLLMBinding):
             return list(text.encode('utf-8'))
 
     def detokenize(self, tokens: list) -> str:
-        """
-        Detokenize a list of tokens.
-        Note: Based on the placeholder tokenizer.
-        """
         try:
             encoding = tiktoken.get_encoding("cl100k_base")
             return encoding.decode(tokens)
@@ -307,9 +446,6 @@ class GeminiBinding(LollmsLLMBinding):
             return bytes(tokens).decode('utf-8', errors='ignore')
 
     def count_tokens(self, text: str) -> int:
-        """
-        Count tokens from a text using the Gemini API.
-        """
         if not self.client or not self.model_name:
             ASCIIColors.warning("Cannot count tokens, Gemini client or model_name not set.")
             return -1
@@ -318,24 +454,19 @@ class GeminiBinding(LollmsLLMBinding):
             return model.count_tokens(text).total_tokens
         except Exception as e:
             ASCIIColors.error(f"Failed to count tokens with Gemini API: {e}")
-            # Fallback to tiktoken for a rough estimate
             return len(self.tokenize(text))
 
     def embed(self, text: str, **kwargs) -> List[float]:
-        """
-        Get embeddings for the input text using Gemini API.
-        """
         if not self.client:
              raise Exception("Gemini client not initialized.")
         
-        # Default to a known Gemini embedding model
         model_to_use = kwargs.get("model", "models/embedding-001")
         
         try:
             response = self.client.embed_content(
                 model=model_to_use,
                 content=text,
-                task_type="retrieval_document" # or "semantic_similarity", etc.
+                task_type="retrieval_document" 
             )
             return response['embedding']
         except Exception as ex:
@@ -343,18 +474,16 @@ class GeminiBinding(LollmsLLMBinding):
             raise Exception(f"Gemini embedding failed: {str(ex)}") from ex
 
     def get_model_info(self) -> dict:
-        """Return information about the current Gemini model setup."""
         return {
             "name": self.binding_name,
             "version": genai.__version__,
             "host_address": "https://generativelanguage.googleapis.com",
             "model_name": self.model_name,
             "supports_structured_output": False,
-            "supports_vision": "vision" in self.model_name or "gemini-1.5" in self.model_name,
+            "supports_vision": True,
         }
 
-    def listModels(self) -> List[Dict[str, str]]:
-        """Lists available generative models from the Gemini service."""
+    def list_models(self) -> List[Dict[str, str]]:
         if not self.client:
             ASCIIColors.error("Gemini client not initialized. Cannot list models.")
             return []
@@ -362,7 +491,6 @@ class GeminiBinding(LollmsLLMBinding):
             ASCIIColors.debug("Listing Gemini models...")
             model_info_list = []
             for m in self.client.list_models():
-                # We are interested in models that can generate content.
                 if 'generateContent' in m.supported_generation_methods:
                     model_info_list.append({
                         'model_name': m.name,
@@ -376,7 +504,6 @@ class GeminiBinding(LollmsLLMBinding):
             return []
 
     def load_model(self, model_name: str) -> bool:
-        """Set the model name for subsequent operations."""
         self.model_name = model_name
         ASCIIColors.info(f"Gemini model set to: {model_name}. It will be used on the next API call.")
         return True
@@ -392,106 +519,81 @@ if __name__ == '__main__':
 
     # --- Configuration ---
     test_model_name = "gemini-1.5-pro-latest"
-    test_vision_model_name = "gemini-1.5-pro-latest" # or gemini-pro-vision
-    test_embedding_model = "models/embedding-001"
-    
-    # This variable is global to the script's execution
     full_streamed_text = ""
 
     try:
         # --- Initialization ---
-        ASCIIColors.cyan("\n--- Initializing Binding ---")
         binding = GeminiBinding(model_name=test_model_name)
         ASCIIColors.green("Binding initialized successfully.")
-        ASCIIColors.info(f"Using google-generativeai version: {genai.__version__}")
-
-        # --- List Models ---
-        ASCIIColors.cyan("\n--- Listing Models ---")
-        models = binding.listModels()
-        if models:
-            ASCIIColors.green(f"Found {len(models)} generative models. First 5:")
-            for m in models[:5]:
-                print(m['model_name'])
-        else:
-            ASCIIColors.warning("No models found or failed to list models.")
-
-        # --- Count Tokens ---
-        ASCIIColors.cyan("\n--- Counting Tokens ---")
-        sample_text = "Hello, world! This is a test."
-        token_count = binding.count_tokens(sample_text)
-        ASCIIColors.green(f"Token count for '{sample_text}': {token_count}")
-
-        # --- Text Generation (Non-Streaming) ---
-        ASCIIColors.cyan("\n--- Text Generation (Non-Streaming) ---")
-        prompt_text = "Explain the importance of bees in one paragraph."
-        ASCIIColors.info(f"Prompt: {prompt_text}")
-        generated_text = binding.generate_text(prompt_text, n_predict=100, stream=False)
-        if isinstance(generated_text, str):
-            ASCIIColors.green(f"Generated text:\n{generated_text}")
-        else:
-            ASCIIColors.error(f"Generation failed: {generated_text}")
 
-        # --- Text Generation (Streaming) ---
-        ASCIIColors.cyan("\n--- Text Generation (Streaming) ---")
-        
-        def stream_callback(chunk: str, msg_type: int):
-            # FIX: Use 'global' to modify the variable in the module's scope
-            global full_streamed_text
-            ASCIIColors.green(chunk, end="", flush=True)
-            full_streamed_text += chunk
-            return True
-        
-        # Reset for this test
-        full_streamed_text = ""
-        ASCIIColors.info(f"Prompt: {prompt_text}")
-        result = binding.generate_text(prompt_text, n_predict=150, stream=True, streaming_callback=stream_callback)
-        print("\n--- End of Stream ---")
-        # 'result' is the full text after streaming, which should match our captured text.
-        ASCIIColors.green(f"Full streamed text (for verification): {result}")
-
-        # --- Embeddings ---
-        ASCIIColors.cyan("\n--- Embeddings ---")
+        # --- Test generate_from_messages ---
+        ASCIIColors.cyan("\n--- Testing generate_from_messages ---")
+        messages = [
+            {"role": "system", "content": "You are a pirate."},
+            {"role": "user", "content": "How are you today?"}
+        ]
+        response = binding.generate_from_messages(messages, n_predict=50)
+        ASCIIColors.green(f"Pirate Response: {response}")
+
+    except Exception as e:
+        ASCIIColors.error(f"An error occurred during testing: {e}")
+        trace_exception(e)
+
+    ASCIIColors.yellow("\nGeminiBinding test finished.")
+
+    def test_connection(self) -> dict:
+        """
+        Tests the connection to the Gemini API using the provided key.
+        """
+        if not self.client:
+             return {"status": False, "message": "Client not configured. Check API Key."}
         try:
-            embedding_text = "Lollms is a cool project."
-            embedding_vector = binding.embed(embedding_text, model=test_embedding_model)
-            ASCIIColors.green(f"Embedding for '{embedding_text}' (first 5 dims): {embedding_vector[:5]}...")
-            ASCIIColors.info(f"Embedding vector dimension: {len(embedding_vector)}")
+            # Attempt to list 1 model to verify auth
+            list(self.client.list_models(page_size=1))
+            return {"status": True, "message": "Connection successful! API Key is valid."}
         except Exception as e:
-            ASCIIColors.warning(f"Could not get embedding: {e}")
+            return {"status": False, "message": f"Connection failed: {str(e)}"}
 
-        # --- Vision Model Test ---
-        dummy_image_path = "gemini_dummy_test_image.png"
+    def list_tuned_models(self) -> dict:
+        """
+        Lists tuned models available for this API key.
+        """
+        if not self.client:
+             return {"status": False, "message": "Client not configured."}
         try:
-            img = Image.new('RGB', (200, 50), color = ('blue'))
-            d = ImageDraw.Draw(img)
-            d.text((10,10), "Test Image", fill=('yellow'))
-            img.save(dummy_image_path)
-            ASCIIColors.info(f"Created dummy image: {dummy_image_path}")
-
-            ASCIIColors.cyan(f"\n--- Vision Generation (using {test_vision_model_name}) ---")
-            binding.load_model(test_vision_model_name)
-            vision_prompt = "What color is the text and what does it say?"
-            ASCIIColors.info(f"Vision Prompt: {vision_prompt} with image {dummy_image_path}")
+            tuned_models = []
+            # iterate over the tuned models generator
+            for model in self.client.list_tuned_models():
+                tuned_models.append(model.name)
             
-            vision_response = binding.generate_text(
-                prompt=vision_prompt,
-                images=[dummy_image_path],
-                n_predict=50,
-                stream=False
-            )
-            if isinstance(vision_response, str):
-                ASCIIColors.green(f"Vision model response: {vision_response}")
-            else:
-                ASCIIColors.error(f"Vision generation failed: {vision_response}")
+            if not tuned_models:
+                 return {"status": True, "message": "No tuned models found on this account.", "models": []}
+            
+            # Join them for a pretty message, but return the list for the UI
+            msg = f"Found: {', '.join(tuned_models)}"
+            return {"status": True, "message": msg, "models": tuned_models}
         except Exception as e:
-            ASCIIColors.error(f"Error during vision test: {e}")
-            trace_exception(e)
-        finally:
-            if os.path.exists(dummy_image_path):
-                os.remove(dummy_image_path)
+            return {"status": False, "message": f"Could not list tuned models: {str(e)}", "models": []}
 
-    except Exception as e:
-        ASCIIColors.error(f"An error occurred during testing: {e}")
-        trace_exception(e)
-
-    ASCIIColors.yellow("\nGeminiBinding test finished.")
+    def get_active_model_info(self) -> dict:
+        """
+        Retrieves details about the currently selected model name.
+        """
+        if not self.client or not self.model_name:
+            return {"status": False, "message": "Client or model name not set."}
+        
+        try:
+            model_info = self.client.get_model(self.model_name)
+            
+            info = {
+                "display_name": model_info.display_name,
+                "description": model_info.description,
+                "input_token_limit": model_info.input_token_limit,
+                "output_token_limit": model_info.output_token_limit,
+                "supported_methods": model_info.supported_generation_methods
+            }
+            
+            msg = f"Input Limit: {info['input_token_limit']}, Output Limit: {info['output_token_limit']}"
+            return {"status": True, "message": msg, "info": info}
+        except Exception as e:
+            return {"status": False, "message": f"Failed to get info for {self.model_name}: {str(e)}"}