my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

my-project-public — gitlawb

#1	import logging
#2	import os
#3	from collections.abc import Generator
#4	from typing import Any, Optional, Union
#5
#6	try:
#7	import google.generativeai as genai
#8	except ImportError:
#9	raise ImportError("GoogleLlm requires extra dependencies. Install with `pip install google-generativeai`") from None
#10
#11	from embedchain.config import BaseLlmConfig
#12	from embedchain.helpers.json_serializable import register_deserializable
#13	from embedchain.llm.base import BaseLlm
#14
#15	logger = logging.getLogger(__name__)
#16
#17
#18	@register_deserializable
#19	class GoogleLlm(BaseLlm):
#20	def __init__(self, config: Optional[BaseLlmConfig] = None):
#21	super().__init__(config)
#22	if not self.config.api_key and "GOOGLE_API_KEY" not in os.environ:
#23	raise ValueError("Please set the GOOGLE_API_KEY environment variable or pass it in the config.")
#24
#25	api_key = self.config.api_key or os.getenv("GOOGLE_API_KEY")
#26	genai.configure(api_key=api_key)
#27
#28	def get_llm_model_answer(self, prompt):
#29	if self.config.system_prompt:
#30	raise ValueError("GoogleLlm does not support `system_prompt`")
#31	response = self._get_answer(prompt)
#32	return response
#33
#34	def _get_answer(self, prompt: str) -> Union[str, Generator[Any, Any, None]]:
#35	model_name = self.config.model or "gemini-pro"
#36	logger.info(f"Using Google LLM model: {model_name}")
#37	model = genai.GenerativeModel(model_name=model_name)
#38
#39	generation_config_params = {
#40	"candidate_count": 1,
#41	"max_output_tokens": self.config.max_tokens,
#42	"temperature": self.config.temperature or 0.5,
#43	}
#44
#45	if 0.0 <= self.config.top_p <= 1.0:
#46	generation_config_params["top_p"] = self.config.top_p
#47	else:
#48	raise ValueError("`top_p` must be > 0.0 and < 1.0")
#49
#50	generation_config = genai.types.GenerationConfig(**generation_config_params)
#51
#52	response = model.generate_content(
#53	prompt,
#54	generation_config=generation_config,
#55	stream=self.config.stream,
#56	)
#57	if self.config.stream:
#58	# TODO: Implement streaming
#59	response.resolve()
#60	return response.text
#61	else:
#62	return response.text
#63

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public