crewplus 0.2.31__py3-none-any.whl → 0.2.33__py3-none-any.whl

crewplus/services/gemini_chat_model.py

@@ -4,6 +4,7 @@ import logging
 from typing import Any, Dict, Iterator, List, Optional, AsyncIterator, Union, Tuple
 from google import genai
 from google.genai import types
+from google.oauth2 import service_account
 import base64
 import requests
 from langchain_core.language_models import BaseChatModel
@@ -31,8 +32,15 @@ class GeminiChatModel(BaseChatModel):
     with the LangChain ecosystem.
 
     It supports standard invocation, streaming, and asynchronous operations.
-    API keys can be provided directly or loaded from the `GOOGLE_API_KEY`
-    environment variable.
+    By default, it uses Google AI with an API key. It can also be configured to use
+    Google Cloud Vertex AI.
+
+    **Authentication:**
+    - **Google AI (Default):** The `google_api_key` parameter or the `GOOGLE_API_KEY` 
+      environment variable is used.
+    - **Vertex AI:** To use Vertex AI, set `use_vertex_ai=True` and provide
+      GCP configuration (`project_id`, `location`). Authentication is handled
+      via `service_account_file`, `credentials`, or Application Default Credentials (ADC).
 
     **Tracing Integration:**
     Tracing (e.g., with Langfuse) is automatically enabled when the respective
@@ -53,6 +61,11 @@ class GeminiChatModel(BaseChatModel):
         top_k (Optional[int]): The top-k sampling parameter.
         logger (Optional[logging.Logger]): An optional logger instance.
         enable_tracing (Optional[bool]): Enable/disable all tracing (auto-detect if None).
+        use_vertex_ai (bool): If True, uses Vertex AI instead of Google AI Platform. Defaults to False.
+        project_id (Optional[str]): GCP Project ID, required for Vertex AI.
+        location (Optional[str]): GCP Location for Vertex AI (e.g., "us-central1").
+        service_account_file (Optional[str]): Path to GCP service account JSON for Vertex AI.
+        credentials (Optional[Any]): GCP credentials object for Vertex AI (alternative to file).
 
     Example:
         .. code-block:: python
@@ -185,6 +198,22 @@ class GeminiChatModel(BaseChatModel):
                 "Hello without tracing",
                 config={"metadata": {"tracing_disabled": True}}
             )
+
+    Example (Vertex AI):
+        .. code-block:: python
+
+            # Assumes GCP environment is configured (e.g., gcloud auth application-default login)
+            # or environment variables are set:
+            # os.environ["GCP_PROJECT_ID"] = "your-gcp-project-id"
+            # os.environ["GCP_LOCATION"] = "us-central1"
+            # os.environ["GCP_SERVICE_ACCOUNT_FILE"] = "path/to/your/service-account-key.json"
+
+            vertex_model = GeminiChatModel(
+                model_name="gemini-1.5-flash-001",
+                use_vertex_ai=True,
+            )
+            response = vertex_model.invoke("Hello from Vertex AI!")
+            print(response.content)
     """
     
     # Model configuration
@@ -195,6 +224,13 @@ class GeminiChatModel(BaseChatModel):
     top_p: Optional[float] = Field(default=None, description="Top-p sampling parameter")
     top_k: Optional[int] = Field(default=None, description="Top-k sampling parameter")
     
+    # Vertex AI specific configuration
+    use_vertex_ai: bool = Field(default=False, description="Use Vertex AI instead of Google AI Platform")
+    project_id: Optional[str] = Field(default=None, description="Google Cloud Project ID for Vertex AI")
+    location: Optional[str] = Field(default=None, description="Google Cloud Location for Vertex AI (e.g., 'us-central1')")
+    service_account_file: Optional[str] = Field(default=None, description="Path to Google Cloud service account key file")
+    credentials: Optional[Any] = Field(default=None, description="Google Cloud credentials object", exclude=True)
+
     # Configuration for tracing and logging
     logger: Optional[logging.Logger] = Field(default=None, description="Optional logger instance", exclude=True)
     enable_tracing: Optional[bool] = Field(default=None, description="Enable tracing (auto-detect if None)")
@@ -213,6 +249,19 @@ class GeminiChatModel(BaseChatModel):
                 self.logger.addHandler(logging.StreamHandler())
                 self.logger.setLevel(logging.INFO)
         
+        self._initialize_client()
+        
+        self._tracing_manager = TracingManager(self)
+
+    def _initialize_client(self):
+        """Initializes the Google GenAI client for either Google AI or Vertex AI."""
+        if self.use_vertex_ai:
+            self._init_vertex_ai_client()
+        else:
+            self._init_google_ai_client()
+
+    def _init_google_ai_client(self):
+        """Initializes the client for Google AI Platform."""
         # Get API key from environment if not provided
         if self.google_api_key is None:
             api_key = os.getenv("GOOGLE_API_KEY")
@@ -222,13 +271,51 @@ class GeminiChatModel(BaseChatModel):
         # Initialize the Google GenAI client
         if self.google_api_key:
             self._client = genai.Client(api_key=self.google_api_key.get_secret_value())
-            self.logger.info(f"Initialized GeminiChatModel with model: {self.model_name}")
+            self.logger.info(f"Initialized GeminiChatModel with model: {self.model_name} for Google AI")
         else:
             error_msg = "Google API key is required. Set GOOGLE_API_KEY environment variable or pass google_api_key parameter."
             self.logger.error(error_msg)
             raise ValueError(error_msg)
+
+    def _init_vertex_ai_client(self):
+        """Initializes the client for Vertex AI."""
+        # Get config from environment if not provided
+        if self.project_id is None:
+            self.project_id = os.getenv("GCP_PROJECT_ID")
+        if self.location is None:
+            self.location = os.getenv("GCP_LOCATION")
         
-        self._tracing_manager = TracingManager(self)
+        if not self.project_id or not self.location:
+            error_msg = "For Vertex AI, 'project_id' and 'location' are required."
+            self.logger.error(error_msg)
+            raise ValueError(error_msg)
+
+        creds = self.credentials
+        if creds is None:
+            # Get service account file from env if not provided
+            sa_file = self.service_account_file or os.getenv("GCP_SERVICE_ACCOUNT_FILE")
+            if sa_file:
+                try:
+                    creds = service_account.Credentials.from_service_account_file(
+                        sa_file,
+                        scopes=['https://www.googleapis.com/auth/cloud-platform']
+                    )
+                except Exception as e:
+                    error_msg = f"Failed to load credentials from service account file '{sa_file}': {e}"
+                    self.logger.error(error_msg)
+                    raise ValueError(error_msg)
+        
+        # If creds is still None, the client will use Application Default Credentials (ADC).
+        
+        self._client = genai.Client(
+            project=self.project_id,
+            location=self.location,
+            credentials=creds,
+        )
+        self.logger.info(
+            f"Initialized GeminiChatModel with model: {self.model_name} for Vertex AI "
+            f"(Project: {self.project_id}, Location: {self.location})"
+        )
 
     def get_model_identifier(self) -> str:
         """Return a string identifying this model for tracing and logging."""

crewplus/services/model_load_balancer.py

@@ -188,6 +188,25 @@ class ModelLoadBalancer:
             if 'max_tokens' in model_config:
                 kwargs['max_tokens'] = model_config['max_tokens']
             return GeminiChatModel(**kwargs)
+        elif provider == 'vertex-ai':
+            deployment_name = model_config['deployment_name']
+            
+            # Handle the 'model_name@location' format for deployment_name
+            model_name_for_gemini = deployment_name.split('@')[0] if '@' in deployment_name else deployment_name
+
+            kwargs = {
+                'use_vertex_ai': True,
+                'model_name': model_name_for_gemini,
+                'project_id': model_config['project_id'],
+                'location': model_config['location'],
+            }
+            if 'service_account_file' in model_config:
+                kwargs['service_account_file'] = model_config['service_account_file']
+            if 'temperature' in model_config:
+                kwargs['temperature'] = model_config['temperature']
+            if 'max_tokens' in model_config:
+                kwargs['max_tokens'] = model_config['max_tokens']
+            return GeminiChatModel(**kwargs)
         else:
             self.logger.error(f"Unsupported provider: {provider}")
             raise ValueError(f"Unsupported provider: {provider}")

{crewplus-0.2.31.dist-info → crewplus-0.2.33.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: crewplus
-Version: 0.2.31
+Version: 0.2.33
 Summary: Base services for CrewPlus AI applications
 Author-Email: Tim Liu <tim@opsmateai.com>
 License: MIT

{crewplus-0.2.31.dist-info → crewplus-0.2.33.dist-info}/RECORD

@@ -1,13 +1,13 @@
-crewplus-0.2.31.dist-info/METADATA,sha256=IyfwlRN7t1zCjAlVhAobrV8gxmJ1z3VuBe-tbmv6Za0,5327
-crewplus-0.2.31.dist-info/WHEEL,sha256=9P2ygRxDrTJz3gsagc0Z96ukrxjr-LFBGOgv3AuKlCA,90
-crewplus-0.2.31.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
-crewplus-0.2.31.dist-info/licenses/LICENSE,sha256=2_NHSHRTKB_cTcT_GXgcenOCtIZku8j343mOgAguTfc,1087
+crewplus-0.2.33.dist-info/METADATA,sha256=epoFGaE7ko9Y0IXOPMSbJqbkNkbizG0pszDxkqr7LB8,5327
+crewplus-0.2.33.dist-info/WHEEL,sha256=9P2ygRxDrTJz3gsagc0Z96ukrxjr-LFBGOgv3AuKlCA,90
+crewplus-0.2.33.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
+crewplus-0.2.33.dist-info/licenses/LICENSE,sha256=2_NHSHRTKB_cTcT_GXgcenOCtIZku8j343mOgAguTfc,1087
 crewplus/__init__.py,sha256=m46HkZL1Y4toD619NL47Sn2Qe084WFFSFD7e6VoYKZc,284
 crewplus/services/__init__.py,sha256=V1CG8b2NOmRzNgQH7BPl4KVxWSYJH5vfEsW1wVErKNE,375
 crewplus/services/azure_chat_model.py,sha256=WMSf4BDO8UcP7ZASNGRJxdTEnuWBmCRSY_4yx_VMbok,5499
-crewplus/services/gemini_chat_model.py,sha256=oVLL07VEjwrHOeP56YSPLnldbfjvTVRrsTI6xcPNt1E,35224
+crewplus/services/gemini_chat_model.py,sha256=XlJeaic1lq31lMx1EPKCe4AlYaHDR57tXOdu8V8j6dU,39623
 crewplus/services/init_services.py,sha256=EBpDkIwzuujmdlqjyWvdLQCfhQmfS_OKFz-9Ji8nmAU,1628
-crewplus/services/model_load_balancer.py,sha256=2LOiPUslRKreBZTndKOx6WH8VtJT8mep37GsDINQTJs,10780
+crewplus/services/model_load_balancer.py,sha256=PU3wn8lh6pGeVFn62SURW_1lIGXbUnAklM1EWcJMhLU,11752
 crewplus/services/tracing_manager.py,sha256=aCU9N4Jvh8pDD3h8kWX4O-Ax8xwdLHnQ4wJ3sf-vLwA,6289
 crewplus/utils/__init__.py,sha256=2Gk1n5srFJQnFfBuYTxktdtKOVZyNrFcNaZKhXk35Pw,142
 crewplus/utils/schema_action.py,sha256=GDaBoVFQD1rXqrLVSMTfXYW1xcUu7eDcHsn57XBSnIg,422
@@ -20,4 +20,4 @@ docs/GeminiChatModel.md,sha256=zZYyl6RmjZTUsKxxMiC9O4yV70MC4TD-IGUmWhIDBKA,8677
 docs/ModelLoadBalancer.md,sha256=aGHES1dcXPz4c7Y8kB5-vsCNJjriH2SWmjBkSGoYKiI,4398
 docs/VDBService.md,sha256=Dw286Rrf_fsi13jyD3Bo4Sy7nZ_G7tYm7d8MZ2j9hxk,9375
 docs/index.md,sha256=3tlc15uR8lzFNM5WjdoZLw0Y9o1P1gwgbEnOdIBspqc,1643
-crewplus-0.2.31.dist-info/RECORD,,
+crewplus-0.2.33.dist-info/RECORD,,