Stephen Hu 2 mesi fa
parent
commit
7713e14d6a
Nessun account collegato all'indirizzo email del committer
1 ha cambiato i file con 14 aggiunte e 3 eliminazioni
  1. 14
    3
      rag/llm/chat_model.py

+ 14
- 3
rag/llm/chat_model.py Vedi File

@@ -1099,9 +1099,20 @@ class GeminiChat(Base):

if system:
self.model._system_instruction = content_types.to_content(system)
response = self.model.generate_content(hist, generation_config=gen_conf)
ans = response.text
return ans, response.usage_metadata.total_token_count
retry_count = 0
max_retries = 3
while retry_count < max_retries:
try:
response = self.model.generate_content(hist, generation_config=gen_conf)
ans = response.text
return ans, response.usage_metadata.total_token_count
except Exception as e:
retry_count += 1
if retry_count >= max_retries:
raise e
else:
import time
time.sleep(50)

def chat_streamly(self, system, history, gen_conf={}, **kwargs):
from google.generativeai.types import content_types

Loading…
Annulla
Salva