Update app.py

app.py

@@ -61,6 +61,21 @@ cross_encoder_tokenizer = AutoTokenizer.from_pretrained(cross_encoder_checkpoint
 cross_encoder = pipeline('text-classification', model=cross_encoder_model, tokenizer=cross_encoder_tokenizer)
 
 
+def update_llm_api_key(api_key, llm_type):
+    if llm_type == 'ChatGPT':
+        chatgpt = OpenAI(
+            api_key=api_key,
+        )
+        llms['ChatGPT'] = chatgpt
+    else:
+        gigachat = GigaChat(
+            credentials=api_key,
+            scope='GIGACHAT_API_PERS',
+            verify_ssl_certs=False
+        )
+        llms['GigaChat'] = gigachat
+
+
 def encode(docs):
     if type(docs) == str:
         docs = [docs]
@@ -147,11 +162,21 @@ def check_request_validity(func):
 
 @check_request_validity
 def fn(
+    chatgpt_key,
+    gigachat_key,
     query,
     llm_type,
     llm_answer_type,
     use_llm_for_retriever
 ):
+    chatgpt_key = chatgpt_key.strip()
+    gigachat_key = gigachat_key.strip()
+    
+    if chatgpt_key:
+        update_llm_api_key(chatgpt_key, 'ChatGPT')
+    if gigachat_key:
+        update_llm_api_key(gigachat_key, 'GigaChat')
+    
     # обогатим запрос с помощью LLM, чтобы поисковым моделям было проще найти нужный фрагмент с ответом
     retriever_ranker_query = update_query_with_llm(query, llm_type, use_llm_for_retriever)
 
@@ -178,11 +203,14 @@ def fn(
 demo = gr.Interface(
     fn=fn,
     inputs=[
+        gr.Textbox(lines=3, label='Ключ Open AI API', placeholder='Введите ключ'),
+        gr.Textbox(lines=3, label='Ключ GigaChat API', placeholder='Введите ключ'),
         gr.Textbox(lines=3, label='Запрос', placeholder='Введите запрос'),
         gr.Dropdown(label='Тип LLM', choices=['ChatGPT', 'GigaChat', 'Без LLM'], value='ChatGPT'),
         gr.Dropdown(label='Тип итогового ответа LLM', choices=['Только цифры штрафа', 'Развернутый ответ'], value='Только цифры штрафа'),
         gr.Checkbox(label="Использовать LLM для Retriever'а", value=True, info="При использовании LLM для Retriever'а ко входному запросу будет добавляться промежуточный ответ LLM на запрос. Это способствует повышению качества поиска ответа."),
         gr.Checkbox(label="Использовать LLM для проверки валидности запроса", value=False)
+        
     ],
     outputs=[
         gr.Textbox(label='Ответ'),