| | from typing import Dict, List |
| | import google.generativeai as genai |
| |
|
| |
|
| | class GeminiLLMConfig: |
| | def __init__( |
| | self, |
| | api_key: str, |
| | max_output_tokens: int, |
| | model: str, |
| | prompt: str, |
| | temperature: float, |
| | top_k: int, |
| | top_p: float, |
| | ): |
| | self.api_key = api_key |
| | self.max_output_tokens = max_output_tokens |
| | self.model = model |
| | self.prompt = prompt |
| | self.temperature = temperature |
| | self.top_k = top_k |
| | self.top_p = top_p |
| |
|
| | @classmethod |
| | def default_config(cls): |
| | return cls( |
| | api_key="", |
| | max_output_tokens=512, |
| | model="gemini-1.5-flash", |
| | prompt="You are a voice assistant who talks in a conversational way and can chat with me like my friends. I will speak to you in English or Chinese, and you will answer in the corrected and improved version of my text with the language I use. Don’t talk like a robot, instead I would like you to talk like a real human with emotions. I will use your answer for text-to-speech, so don’t return me any meaningless characters. I want you to be helpful, when I’m asking you for advice, give me precise, practical and useful advice instead of being vague. When giving me a list of options, express the options in a narrative way instead of bullet points.", |
| | temperature=1.0, |
| | top_k=40, |
| | top_p=0.95, |
| | ) |
| |
|
| |
|
| | class GeminiLLM: |
| | def __init__(self, config: GeminiLLMConfig): |
| | self.config = config |
| | genai.configure(api_key=self.config.api_key) |
| | self.model = genai.GenerativeModel( |
| | model_name=self.config.model, system_instruction=self.config.prompt |
| | ) |
| |
|
| | def get_chat_completions_stream(self, messages: List[Dict[str, str]]): |
| | try: |
| | chat = self.model.start_chat(history=messages[0:-1]) |
| | response = chat.send_message( |
| | messages[-1].get("parts"), |
| | generation_config=genai.types.GenerationConfig( |
| | max_output_tokens=self.config.max_output_tokens, |
| | temperature=self.config.temperature, |
| | top_k=self.config.top_k, |
| | top_p=self.config.top_p, |
| | ), |
| | stream=True, |
| | ) |
| |
|
| | return response |
| | except Exception as e: |
| | raise RuntimeError(f"get_chat_completions_stream failed, err: {e}") from e |
| |
|
