Update app.py

app.py

@@ -1,4 +1,11 @@
-import time, os, multiprocessing, torch, requests, asyncio, json, aiohttp
+import time
+import os
+import multiprocessing
+import torch
+import requests
+import asyncio
+import json
+import aiohttp
 from minivectordb.embedding_model import EmbeddingModel
 from minivectordb.vector_database import VectorDatabase
 from text_util_en_pt.cleaner import structurize_text, detect_language, Language
@@ -24,17 +31,17 @@ def index_and_search(query, text):
 
     # Indexing
     vector_db = VectorDatabase()
-    sentences = [ s['sentence'] for s in structurize_text(text)]
+    sentences = [s['sentence'] for s in structurize_text(text)]
 
     for idx, sentence in enumerate(sentences):
         sentence_embedding = model.extract_embeddings(sentence)
         vector_db.store_embedding(idx + 1, sentence_embedding, {'sentence': sentence})
-    
+
     embedding_time = time.time() - start
 
     # Retrieval
     start = time.time()
-    search_results = vector_db.find_most_similar(query_embedding, k = 30)
+    search_results = vector_db.find_most_similar(query_embedding, k=30)
     retrieval_time = time.time() - start
     return '\n'.join([s['sentence'] for s in search_results[2]]), embedding_time, retrieval_time
 
@@ -45,21 +52,30 @@ def generate_search_terms(message, lang):
         prompt = f"From the following text, generate some search terms: \"{message}\"\nYour answer should be just the most appropriate search term, and nothing else."
 
     url = "https://openrouter.ai/api/v1/chat/completions"
-    headers = { "Content-Type": "application/json",
-                "Authorization": f"Bearer {openrouter_key}" }
-    body = { "stream": False,
-             "models": [
-                    "mistralai/mistral-7b-instruct:free",
-                    "openchat/openchat-7b:free"
-             ],
-             "route": "fallback",
-             "max_tokens": 1024,
-             "messages": [
-                 {"role": "user", "content": prompt}
-             ] }
-    
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {openrouter_key}"
+    }
+    body = {
+        "stream": False,
+        "models": [
+            "mistralai/mistral-7b-instruct:free",
+            "openchat/openchat-7b:free"
+        ],
+        "route": "fallback",
+        "max_tokens": 1024,
+        "messages": [
+            {"role": "user", "content": prompt}
+        ]
+    }
+
     response = requests.post(url, headers=headers, json=body)
-    return response.json()['choices'][0]['message']['content']
+    response_json = response.json()
+    try:
+        return response_json['choices'][0]['message']['content']
+    except KeyError:
+        print(f"Error: 'choices' key not found in the response. Response: {response_json}")
+        return None
 
 async def predict(message, history):
     full_response = ""
@@ -91,7 +107,7 @@ async def predict(message, history):
     for link in links:
         full_response += f"{link}\n"
         yield full_response
-    
+
     full_response += "\nExtracting text from web pages...\n"
     yield full_response
     start = time.time()
@@ -118,20 +134,24 @@ async def predict(message, history):
     yield full_response
 
     full_response += "\nResponse: "
-    
+
     url = "https://openrouter.ai/api/v1/chat/completions"
-    headers = { "Content-Type": "application/json",
-                "Authorization": f"Bearer {openrouter_key}" }
-    body = { "stream": True,
-             "models": [
-                    "mistralai/mistral-7b-instruct:free",
-                    "openchat/openchat-7b:free"
-             ],
-             "route": "fallback",
-             "max_tokens": 1024,
-             "messages": [
-                 {"role": "user", "content": prompt}
-             ] }
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {openrouter_key}"
+    }
+    body = {
+        "stream": True,
+        "models": [
+            "mistralai/mistral-7b-instruct:free",
+            "openchat/openchat-7b:free"
+        ],
+        "route": "fallback",
+        "max_tokens": 1024,
+        "messages": [
+            {"role": "user", "content": prompt}
+        ]
+    }
 
     async with aiohttp.ClientSession() as session:
         async with session.post(url, headers=headers, json=body) as response:
@@ -169,4 +189,4 @@ gr.ChatInterface(
         'Compare the current economies of China and India?',
         'What are new shoe design trends in 2024',
     ]
-).launch()
+).launch()