my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	import importlib
#2	import logging
#3	from typing import Any, Optional
#4
#5	from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
#6	from langchain_google_vertexai import ChatVertexAI
#7
#8	from embedchain.config import BaseLlmConfig
#9	from embedchain.helpers.json_serializable import register_deserializable
#10	from embedchain.llm.base import BaseLlm
#11
#12	logger = logging.getLogger(__name__)
#13
#14
#15	@register_deserializable
#16	class VertexAILlm(BaseLlm):
#17	def __init__(self, config: Optional[BaseLlmConfig] = None):
#18	try:
#19	importlib.import_module("vertexai")
#20	except ModuleNotFoundError:
#21	raise ModuleNotFoundError(
#22	"The required dependencies for VertexAI are not installed."
#23	'Please install with `pip install --upgrade "embedchain[vertexai]"`'
#24	) from None
#25	super().__init__(config=config)
#26
#27	def get_llm_model_answer(self, prompt) -> tuple[str, Optional[dict[str, Any]]]:
#28	if self.config.token_usage:
#29	response, token_info = self._get_answer(prompt, self.config)
#30	model_name = "vertexai/" + self.config.model
#31	if model_name not in self.config.model_pricing_map:
#32	raise ValueError(
#33	f"Model {model_name} not found in `model_prices_and_context_window.json`. \
#34	You can disable token usage by setting `token_usage` to False."
#35	)
#36	total_cost = (
#37	self.config.model_pricing_map[model_name]["input_cost_per_token"] * token_info["prompt_token_count"]
#38	) + self.config.model_pricing_map[model_name]["output_cost_per_token"] * token_info[
#39	"candidates_token_count"
#40	]
#41	response_token_info = {
#42	"prompt_tokens": token_info["prompt_token_count"],
#43	"completion_tokens": token_info["candidates_token_count"],
#44	"total_tokens": token_info["prompt_token_count"] + token_info["candidates_token_count"],
#45	"total_cost": round(total_cost, 10),
#46	"cost_currency": "USD",
#47	}
#48	return response, response_token_info
#49	return self._get_answer(prompt, self.config)
#50
#51	@staticmethod
#52	def _get_answer(prompt: str, config: BaseLlmConfig) -> str:
#53	if config.top_p and config.top_p != 1:
#54	logger.warning("Config option `top_p` is not supported by this model.")
#55
#56	if config.stream:
#57	callbacks = config.callbacks if config.callbacks else [StreamingStdOutCallbackHandler()]
#58	llm = ChatVertexAI(
#59	temperature=config.temperature, model=config.model, callbacks=callbacks, streaming=config.stream
#60	)
#61	else:
#62	llm = ChatVertexAI(temperature=config.temperature, model=config.model)
#63
#64	messages = VertexAILlm._get_messages(prompt)
#65	chat_response = llm.invoke(messages)
#66	if config.token_usage:
#67	return chat_response.content, chat_response.response_metadata["usage_metadata"]
#68	return chat_response.content
#69

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public