From b96ba6f8310189bede352894c24a8c0484ba4eeb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=BB=84=E8=85=BE?=
 <101850389+hangters@users.noreply.github.com>
Date: Fri, 12 Jul 2024 13:41:58 +0800
Subject: [PATCH] add Gemini key (#1480)

### What problem does this PR solve?

#1036

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)

---------

Co-authored-by: Zhedong Cen <cenzhedong2@126.com>
---
 rag/llm/__init__.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/rag/llm/__init__.py b/rag/llm/__init__.py
index 65a76e8f1..3fdab822d 100644
--- a/rag/llm/__init__.py
+++ b/rag/llm/__init__.py
@@ -33,7 +33,7 @@ EmbeddingModel = {
     "BAAI": DefaultEmbedding,
     "Mistral": MistralEmbed,
     "Bedrock": BedrockEmbed,
-    "Groq": GroqChat
+    "Gemini":GeminiEmbed,
 }
 
 
@@ -44,7 +44,8 @@ CvModel = {
     "Xinference": XinferenceCV,
     "Tongyi-Qianwen": QWenCV,
     "ZHIPU-AI": Zhipu4V,
-    "Moonshot": LocalCV
+    "Moonshot": LocalCV,
+    'Gemini':GeminiCV
 }
 
 
@@ -61,7 +62,9 @@ ChatModel = {
     "BaiChuan": BaiChuanChat,
     "MiniMax": MiniMaxChat,
     "Mistral": MistralChat,
-    "Bedrock": BedrockChat
+    'Gemini' : GeminiChat,
+    "Bedrock": BedrockChat,
+    "Groq": GroqChat
 }