Initialize the Vertex AI wrapper. Args: model_name: Name of the model to use (e.g. "gemini-1.5-pro") temperature: Temperature for generation between 0 and 1 print_cost: Whether to print the cost of the completion verbose: Whether to pr
(
self,
model_name: str = "gemini-1.5-pro",
temperature: float = 0.7,
print_cost: bool = False,
verbose: bool = False,
use_langfuse: bool = False
)
| 12 | """Wrapper for Vertex AI to support Gemini models.""" |
| 13 | |
| 14 | def __init__( |
| 15 | self, |
| 16 | model_name: str = "gemini-1.5-pro", |
| 17 | temperature: float = 0.7, |
| 18 | print_cost: bool = False, |
| 19 | verbose: bool = False, |
| 20 | use_langfuse: bool = False |
| 21 | ): |
| 22 | """Initialize the Vertex AI wrapper. |
| 23 | |
| 24 | Args: |
| 25 | model_name: Name of the model to use (e.g. "gemini-1.5-pro") |
| 26 | temperature: Temperature for generation between 0 and 1 |
| 27 | print_cost: Whether to print the cost of the completion |
| 28 | verbose: Whether to print verbose output |
| 29 | use_langfuse: Whether to enable Langfuse logging |
| 30 | """ |
| 31 | self.model_name = model_name |
| 32 | self.temperature = temperature |
| 33 | self.print_cost = print_cost |
| 34 | self.verbose = verbose |
| 35 | |
| 36 | # Initialize Vertex AI |
| 37 | project_id = os.getenv("GOOGLE_CLOUD_PROJECT") |
| 38 | location = os.getenv("GOOGLE_CLOUD_LOCATION", "us-central1") |
| 39 | if not project_id: |
| 40 | raise ValueError("No GOOGLE_CLOUD_PROJECT found in environment variables") |
| 41 | |
| 42 | vertexai.init(project=project_id, location=location) |
| 43 | self.model = GenerativeModel(model_name) |
| 44 | |
| 45 | def __call__(self, messages: List[Dict[str, Any]], metadata: Optional[Dict[str, Any]] = None) -> str: |
| 46 | """Process messages and return completion. |
nothing calls this directly
no outgoing calls
no test coverage detected