def get_model_obj()

in Runtime_env/app/orchestration/agent.py [0:0]

• 27 lines of code
• 6 McCabe index (conditional complexity)

    def get_model_obj(self):
        """
        Helper method to retrieve the model object based on the model name 
        and config.

        Returns:
            An LLM object for the Agent to use.

        Exception:
            The model_name is not found.
        """
        try:
            if "claude" in self.model_name:
                return ChatAnthropicVertex(
                    model_name=self.model_name,
                    max_retries=self.max_retries,
                    ## causes issues with the pydantic model with default None value:
                    # max_output_tokens=self.max_output_tokens,
                    ## for now, claude is only available in us-east5 and europe-west1
                    # location = self.location,
                    location="us-east5",
                    temperature=self.temperature,
                    top_p=self.top_p,
                    top_k=self.top_k,
                    verbose=self.verbose
                )
            else:
                return ChatVertexAI(
                    model_name=self.model_name,
                    location=self.location,
                    max_retries=self.max_retries,
                    max_output_tokens=self.max_output_tokens,
                    temperature=self.temperature,
                    top_p=self.top_p,
                    top_k=self.top_k,
                    verbose=self.verbose
                )
        except exceptions.NotFound as e:
            raise exceptions.NotFound(f"Resource not found. {e}")
        except Exception as e:
            raise RuntimeError(f"Error encountered initalizing model resource. {e}") from e