65 lines
2.1 KiB
Python
65 lines
2.1 KiB
Python
import importlib
|
|
import logging
|
|
import os
|
|
from typing import Optional
|
|
|
|
import google.generativeai as genai
|
|
|
|
from embedchain.config import BaseLlmConfig
|
|
from embedchain.helpers.json_serializable import register_deserializable
|
|
from embedchain.llm.base import BaseLlm
|
|
|
|
|
|
@register_deserializable
|
|
class GoogleLlm(BaseLlm):
|
|
def __init__(self, config: Optional[BaseLlmConfig] = None):
|
|
if "GOOGLE_API_KEY" not in os.environ:
|
|
raise ValueError("Please set the GOOGLE_API_KEY environment variable.")
|
|
|
|
try:
|
|
importlib.import_module("google.generativeai")
|
|
except ModuleNotFoundError:
|
|
raise ModuleNotFoundError(
|
|
"The required dependencies for GoogleLlm are not installed."
|
|
'Please install with `pip install --upgrade "embedchain[google]"`'
|
|
) from None
|
|
|
|
super().__init__(config)
|
|
genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
|
|
|
|
def get_llm_model_answer(self, prompt):
|
|
if self.config.system_prompt:
|
|
raise ValueError("GoogleLlm does not support `system_prompt`")
|
|
return GoogleLlm._get_answer(prompt, self.config)
|
|
|
|
@staticmethod
|
|
def _get_answer(prompt: str, config: BaseLlmConfig):
|
|
model_name = config.model or "gemini-pro"
|
|
logging.info(f"Using Google LLM model: {model_name}")
|
|
model = genai.GenerativeModel(model_name=model_name)
|
|
|
|
generation_config_params = {
|
|
"candidate_count": 1,
|
|
"max_output_tokens": config.max_tokens,
|
|
"temperature": config.temperature or 0.5,
|
|
}
|
|
|
|
if config.top_p >= 0.0 and config.top_p <= 1.0:
|
|
generation_config_params["top_p"] = config.top_p
|
|
else:
|
|
raise ValueError("`top_p` must be > 0.0 and < 1.0")
|
|
|
|
generation_config = genai.types.GenerationConfig(**generation_config_params)
|
|
|
|
response = model.generate_content(
|
|
prompt,
|
|
generation_config=generation_config,
|
|
stream=config.stream,
|
|
)
|
|
|
|
if config.stream:
|
|
for chunk in response:
|
|
yield chunk.text
|
|
else:
|
|
return response.text
|