|
import json |
|
import logging |
|
import textwrap |
|
import uuid |
|
|
|
import google.generativeai as genai |
|
import gradio as gr |
|
import PIL |
|
import requests |
|
|
|
from modules.presets import i18n |
|
|
|
from ..index_func import construct_index |
|
from ..utils import count_token |
|
from .base_model import BaseLLMModel |
|
|
|
|
|
class GoogleGeminiClient(BaseLLMModel): |
|
def __init__(self, model_name, api_key, user_name="") -> None: |
|
super().__init__(model_name=model_name, user=user_name) |
|
self.api_key = api_key |
|
if "vision" in model_name.lower(): |
|
self.multimodal = True |
|
else: |
|
self.multimodal = False |
|
self.image_paths = [] |
|
|
|
def _get_gemini_style_input(self): |
|
self.history.extend([{"role": "image", "content": i} for i in self.image_paths]) |
|
self.image_paths = [] |
|
messages = [] |
|
for item in self.history: |
|
if item["role"] == "image": |
|
messages.append(PIL.Image.open(item["content"])) |
|
else: |
|
messages.append(item["content"]) |
|
return messages |
|
|
|
def to_markdown(self, text): |
|
text = text.replace("•", " *") |
|
return textwrap.indent(text, "> ", predicate=lambda _: True) |
|
|
|
def handle_file_upload(self, files, chatbot, language): |
|
if files: |
|
if self.multimodal: |
|
for file in files: |
|
if file.name: |
|
self.image_paths.append(file.name) |
|
chatbot = chatbot + [((file.name,), None)] |
|
return None, chatbot, None |
|
else: |
|
construct_index(self.api_key, file_src=files) |
|
status = i18n("索引构建完成") |
|
return gr.update(), chatbot, status |
|
|
|
def get_answer_at_once(self): |
|
genai.configure(api_key=self.api_key) |
|
messages = self._get_gemini_style_input() |
|
model = genai.GenerativeModel(self.model_name) |
|
response = model.generate_content(messages) |
|
try: |
|
return self.to_markdown(response.text), len(response.text) |
|
except ValueError: |
|
return ( |
|
i18n("由于下面的原因,Google 拒绝返回 Gemini 的回答:\n\n") |
|
+ str(response.prompt_feedback), |
|
0, |
|
) |
|
|
|
def get_answer_stream_iter(self): |
|
genai.configure(api_key=self.api_key) |
|
messages = self._get_gemini_style_input() |
|
model = genai.GenerativeModel(self.model_name) |
|
response = model.generate_content(messages, stream=True) |
|
partial_text = "" |
|
for i in response: |
|
response = i.text |
|
partial_text += response |
|
yield partial_text |
|
self.all_token_counts[-1] = count_token(partial_text) |
|
yield partial_text |
|
|