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 GoogleModelFamily, _VertexAICommon
from langchain_google_vertexai._utils import (
create_retry_decorator,
get_generation_info,
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
)
[docs]
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")
safety_settings = values["safety_settings"]
values["model_family"] = GoogleModelFamily(values["model_name"])
is_gemini = is_gemini_model(values["model_family"])
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 values["model_family"] == GoogleModelFamily.CODEY:
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
)
