Toggle Menu
Prev Up Next

# Source code for langchain_google_vertexai.llms

    from __future__ import annotations

    from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union

    from google.cloud.aiplatform import telemetry
    from langchain_core.callbacks.manager import (
        AsyncCallbackManagerForLLMRun,
        CallbackManagerForLLMRun,
    )
    from langchain_core.language_models.llms import BaseLLM
    from langchain_core.outputs import Generation, GenerationChunk, LLMResult
    from langchain_core.pydantic_v1 import Field, root_validator
    from vertexai.generative_models import (  # type: ignore[import-untyped]
        Candidate,
        GenerativeModel,
        Image,
    )
    from vertexai.language_models import (  # type: ignore[import-untyped]
        CodeGenerationModel,
        TextGenerationModel,
    )
    from vertexai.language_models._language_models import (  # type: ignore[import-untyped]
        TextGenerationResponse,
    )
    from vertexai.preview.language_models import (  # type: ignore[import-untyped]
        CodeGenerationModel as PreviewCodeGenerationModel,
    )
    from vertexai.preview.language_models import (
        TextGenerationModel as PreviewTextGenerationModel,
    )

    from langchain_google_vertexai._base import (
        _VertexAICommon,
    )
    from langchain_google_vertexai._utils import (
        create_retry_decorator,
        get_generation_info,
        is_codey_model,
        is_gemini_model,
    )

    def _completion_with_retry(
        llm: VertexAI,
        prompt: List[Union[str, Image]],
        stream: bool = False,
        is_gemini: bool = False,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> Any:
        """Use tenacity to retry the completion call."""
        retry_decorator = create_retry_decorator(
            max_retries=llm.max_retries, run_manager=run_manager
        )

        @retry_decorator
        def _completion_with_retry_inner(
            prompt: List[Union[str, Image]], is_gemini: bool = False, **kwargs: Any
        ) -> Any:
            if is_gemini:
                return llm.client.generate_content(
                    prompt,
                    stream=stream,
                    safety_settings=kwargs.pop("safety_settings", None),
                    generation_config=kwargs,
                )
            else:
                if stream:
                    return llm.client.predict_streaming(prompt[0], **kwargs)
                return llm.client.predict(prompt[0], **kwargs)

        with telemetry.tool_context_manager(llm._user_agent):
            return _completion_with_retry_inner(prompt, is_gemini, **kwargs)

    async def _acompletion_with_retry(
        llm: VertexAI,
        prompt: str,
        is_gemini: bool = False,
        stream: bool = False,
        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> Any:
        """Use tenacity to retry the completion call."""
        retry_decorator = create_retry_decorator(
            max_retries=llm.max_retries, run_manager=run_manager
        )

        @retry_decorator
        async def _acompletion_with_retry_inner(
            prompt: str, is_gemini: bool = False, stream: bool = False, **kwargs: Any
        ) -> Any:
            if is_gemini:
                return await llm.client.generate_content_async(
                    prompt,
                    generation_config=kwargs,
                    stream=stream,
                    safety_settings=kwargs.pop("safety_settings", None),
                )
            if stream:
                raise ValueError("Async streaming is supported only for Gemini family!")
            return await llm.client.predict_async(prompt, **kwargs)

        with telemetry.tool_context_manager(llm._user_agent):
            return await _acompletion_with_retry_inner(
                prompt, is_gemini, stream=stream, **kwargs
            )

    class VertexAI(_VertexAICommon, BaseLLM):
        """Google Vertex AI large language models."""

        model_name: str = Field(default="text-bison", alias="model")
        "The name of the Vertex AI large language model."
        tuned_model_name: Optional[str] = None
        """The name of a tuned model. If tuned_model_name is passed
        model_name will be used to determine the model family
        """

        def __init__(self, *, model_name: Optional[str] = None, **kwargs: Any) -> None:
            """Needed for mypy typing to recognize model_name as a valid arg."""
            if model_name:
                kwargs["model_name"] = model_name
            super().__init__(**kwargs)

        class Config:
            """Configuration for this pydantic object."""

            allow_population_by_field_name = True

        @classmethod
        def is_lc_serializable(self) -> bool:
            return True

        @classmethod
        def get_lc_namespace(cls) -> List[str]:
            """Get the namespace of the langchain object."""
            return ["langchain", "llms", "vertexai"]

        @root_validator()
        def validate_environment(cls, values: Dict) -> Dict:
            """Validate that the python package exists in environment."""
            tuned_model_name = values.get("tuned_model_name")
            model_name = values["model_name"]
            safety_settings = values["safety_settings"]
            is_gemini = is_gemini_model(values["model_name"])
            cls._init_vertexai(values)

            if safety_settings and (not is_gemini or tuned_model_name):
                raise ValueError("Safety settings are only supported for Gemini models")

            if is_codey_model(model_name):
                model_cls = CodeGenerationModel
                preview_model_cls = PreviewCodeGenerationModel
            elif is_gemini:
                model_cls = GenerativeModel
                preview_model_cls = GenerativeModel
            else:
                model_cls = TextGenerationModel
                preview_model_cls = PreviewTextGenerationModel

            if tuned_model_name:
                generative_model_name = values["tuned_model_name"]
            else:
                generative_model_name = values["model_name"]

            if is_gemini:
                values["client"] = model_cls(
                    model_name=generative_model_name, safety_settings=safety_settings
                )
                values["client_preview"] = preview_model_cls(
                    model_name=generative_model_name, safety_settings=safety_settings
                )
            else:
                if tuned_model_name:
                    values["client"] = model_cls.get_tuned_model(generative_model_name)
                    values["client_preview"] = preview_model_cls.get_tuned_model(
                        generative_model_name
                    )
                else:
                    values["client"] = model_cls.from_pretrained(generative_model_name)
                    values["client_preview"] = preview_model_cls.from_pretrained(
                        generative_model_name
                    )
            if values["streaming"] and values["n"] > 1:
                raise ValueError("Only one candidate can be generated with streaming!")
            return values

        def _candidate_to_generation(
            self,
            response: Union[Candidate, TextGenerationResponse],
            *,
            stream: bool = False,
            usage_metadata: Optional[Dict] = None,
        ) -> GenerationChunk:
            """Converts a stream response to a generation chunk."""
            generation_info = get_generation_info(
                response,
                self._is_gemini_model,
                stream=stream,
                usage_metadata=usage_metadata,
            )
            try:
                text = response.text
            except AttributeError:
                text = ""
            except ValueError:
                text = ""
            return GenerationChunk(
                text=text,
                generation_info=generation_info,
            )

        def _generate(
            self,
            prompts: List[str],
            stop: Optional[List[str]] = None,
            run_manager: Optional[CallbackManagerForLLMRun] = None,
            stream: Optional[bool] = None,
            **kwargs: Any,
        ) -> LLMResult:
            should_stream = stream if stream is not None else self.streaming
            params = self._prepare_params(stop=stop, stream=should_stream, **kwargs)
            generations: List[List[Generation]] = []
            for prompt in prompts:
                if should_stream:
                    generation = GenerationChunk(text="")
                    for chunk in self._stream(
                        prompt, stop=stop, run_manager=run_manager, **kwargs
                    ):
                        generation += chunk
                    generations.append([generation])
                else:
                    res = _completion_with_retry(
                        self,
                        [prompt],
                        stream=should_stream,
                        is_gemini=self._is_gemini_model,
                        run_manager=run_manager,
                        **params,
                    )
                    if self._is_gemini_model:
                        usage_metadata = res.to_dict().get("usage_metadata")
                    else:
                        usage_metadata = res.raw_prediction_response.metadata
                    generations.append(
                        [
                            self._candidate_to_generation(r, usage_metadata=usage_metadata)
                            for r in res.candidates
                        ]
                    )
            return LLMResult(generations=generations)

        async def _agenerate(
            self,
            prompts: List[str],
            stop: Optional[List[str]] = None,
            run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
            **kwargs: Any,
        ) -> LLMResult:
            params = self._prepare_params(stop=stop, **kwargs)
            generations: List[List[Generation]] = []
            for prompt in prompts:
                res = await _acompletion_with_retry(
                    self,
                    prompt,
                    is_gemini=self._is_gemini_model,
                    run_manager=run_manager,
                    **params,
                )
                if self._is_gemini_model:
                    usage_metadata = res.to_dict().get("usage_metadata")
                else:
                    usage_metadata = res.raw_prediction_response.metadata
                generations.append(
                    [
                        self._candidate_to_generation(r, usage_metadata=usage_metadata)
                        for r in res.candidates
                    ]
                )
            return LLMResult(generations=generations)

        def _stream(
            self,
            prompt: str,
            stop: Optional[List[str]] = None,
            run_manager: Optional[CallbackManagerForLLMRun] = None,
            **kwargs: Any,
        ) -> Iterator[GenerationChunk]:
            params = self._prepare_params(stop=stop, stream=True, **kwargs)
            for stream_resp in _completion_with_retry(
                self,
                [prompt],
                stream=True,
                is_gemini=self._is_gemini_model,
                run_manager=run_manager,
                **params,
            ):
                usage_metadata = None
                if self._is_gemini_model:
                    usage_metadata = stream_resp.to_dict().get("usage_metadata")
                    stream_resp = stream_resp.candidates[0]
                chunk = self._candidate_to_generation(
                    stream_resp, stream=True, usage_metadata=usage_metadata
                )
                yield chunk
                if run_manager:
                    run_manager.on_llm_new_token(
                        chunk.text,
                        chunk=chunk,
                        verbose=self.verbose,
                    )

        async def _astream(
            self,
            prompt: str,
            stop: Optional[List[str]] = None,
            run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
            **kwargs: Any,
        ) -> AsyncIterator[GenerationChunk]:
            params = self._prepare_params(stop=stop, stream=True, **kwargs)
            if not self._is_gemini_model:
                raise ValueError("Async streaming is supported only for Gemini family!")
            async for chunk in await _acompletion_with_retry(
                self,
                prompt,
                stream=True,
                is_gemini=self._is_gemini_model,
                run_manager=run_manager,
                **params,
            ):
                usage_metadata = chunk.to_dict().get("usage_metadata")
                chunk = self._candidate_to_generation(
                    chunk.candidates[0], stream=True, usage_metadata=usage_metadata
                )
                yield chunk
                if run_manager:
                    await run_manager.on_llm_new_token(
                        chunk.text, chunk=chunk, verbose=self.verbose
                    )