Spaces:
Running
Running
| from typing import cast | |
| from langflow.base.models.model import LCModelComponent | |
| from langflow.field_typing import LanguageModel | |
| from langflow.inputs import MessageTextInput | |
| from langflow.inputs.inputs import HandleInput | |
| from langflow.io import BoolInput, FileInput, FloatInput, IntInput, StrInput | |
| class ChatVertexAIComponent(LCModelComponent): | |
| display_name = "Vertex AI" | |
| description = "Generate text using Vertex AI LLMs." | |
| icon = "VertexAI" | |
| name = "VertexAiModel" | |
| inputs = [ | |
| *LCModelComponent._base_inputs, | |
| FileInput( | |
| name="credentials", | |
| display_name="Credentials", | |
| info="JSON credentials file. Leave empty to fallback to environment variables", | |
| file_types=["json"], | |
| ), | |
| MessageTextInput(name="model_name", display_name="Model Name", value="gemini-1.5-pro"), | |
| StrInput(name="project", display_name="Project", info="The project ID.", advanced=True), | |
| StrInput(name="location", display_name="Location", value="us-central1", advanced=True), | |
| IntInput(name="max_output_tokens", display_name="Max Output Tokens", advanced=True), | |
| IntInput(name="max_retries", display_name="Max Retries", value=1, advanced=True), | |
| FloatInput(name="temperature", value=0.0, display_name="Temperature"), | |
| IntInput(name="top_k", display_name="Top K", advanced=True), | |
| FloatInput(name="top_p", display_name="Top P", value=0.95, advanced=True), | |
| BoolInput(name="verbose", display_name="Verbose", value=False, advanced=True), | |
| HandleInput( | |
| name="output_parser", | |
| display_name="Output Parser", | |
| info="The parser to use to parse the output of the model", | |
| advanced=True, | |
| input_types=["OutputParser"], | |
| ), | |
| ] | |
| def build_model(self) -> LanguageModel: | |
| try: | |
| from langchain_google_vertexai import ChatVertexAI | |
| except ImportError as e: | |
| msg = "Please install the langchain-google-vertexai package to use the VertexAIEmbeddings component." | |
| raise ImportError(msg) from e | |
| location = self.location or None | |
| if self.credentials: | |
| from google.cloud import aiplatform | |
| from google.oauth2 import service_account | |
| credentials = service_account.Credentials.from_service_account_file(self.credentials) | |
| project = self.project or credentials.project_id | |
| # ChatVertexAI sometimes skip manual credentials initialization | |
| aiplatform.init( | |
| project=project, | |
| location=location, | |
| credentials=credentials, | |
| ) | |
| else: | |
| project = self.project or None | |
| credentials = None | |
| return cast( | |
| "LanguageModel", | |
| ChatVertexAI( | |
| credentials=credentials, | |
| location=location, | |
| project=project, | |
| max_output_tokens=self.max_output_tokens or None, | |
| max_retries=self.max_retries, | |
| model_name=self.model_name, | |
| temperature=self.temperature, | |
| top_k=self.top_k or None, | |
| top_p=self.top_p, | |
| verbose=self.verbose, | |
| ), | |
| ) | |