From 52bc0832da785f6906a992c34ea3deb121bff001 Mon Sep 17 00:00:00 2001 From: Eli <43382407+eli64s@users.noreply.github.com> Date: Thu, 29 Feb 2024 21:36:32 -0600 Subject: [PATCH] Code clean up. --- readmeai/models/vertex.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/readmeai/models/vertex.py b/readmeai/models/vertex.py index cf0e5fc0..03727da9 100644 --- a/readmeai/models/vertex.py +++ b/readmeai/models/vertex.py @@ -20,10 +20,6 @@ from readmeai.models.tokens import token_handler from readmeai.utils.text_cleaner import clean_response -_location = os.environ.get("VERTEXAI_LOCATION") -_project_id = os.environ.get("VERTEXAI_PROJECT") -vertexai.init(location=_location, project=_project_id) - class VertexAIHandler(BaseModelHandler): """Google Cloud Vertex AI LLM API implementation.""" @@ -35,10 +31,13 @@ def __init__(self, config_loader: ConfigLoader) -> None: def _model_settings(self): """Initializes the Vertex AI LLM settings.""" - self.model = GenerativeModel(self.config.llm.model) self.temperature = self.config.llm.temperature self.tokens = self.config.llm.tokens self.top_p = self.config.llm.top_p + self.location = os.environ.get("VERTEXAI_LOCATION") + self.project_id = os.environ.get("VERTEXAI_PROJECT") + vertexai.init(location=self.location, project=self.project_id) + self.model = GenerativeModel(self.config.llm.model) async def _build_payload( self, prompt: str, tokens: int