Update app.py

app.py

@@ -1,21 +1,17 @@
-import time, os, multiprocessing, torch
+import time, os, multiprocessing, torch, requests, asyncio, json, 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
 from webtextcrawler.webtextcrawler import extract_text_from_url
-from duckduckgo_search import DDGS
-import google.generativeai as genai
 import gradio as gr
+from googlesearch import search
 
 torch.set_num_threads(2)
-gemini_key = os.environ.get("GEMINI_KEY")
-genai.configure(api_key=gemini_key)
-gemini = genai.GenerativeModel('gemini-pro')
+openrouter_key = os.environ.get("OPENROUTER_KEY")
 model = EmbeddingModel(use_quantized_onnx_model=True)
 
 def fetch_links(query, max_results=10):
-    with DDGS() as ddgs:
-        return [r['href'] for r in ddgs.text(keywords=query, max_results=max_results)]
+    return list(search(query, num_results=max_results))
 
 def fetch_texts(links):
     with multiprocessing.Pool(10) as pool:
@@ -48,8 +44,22 @@ def generate_search_terms(message, lang):
     else:
         prompt = f"From the following text, generate some search terms: \"{message}\"\nYour answer should be just the most appropriate search term, and nothing else."
 
-    response = gemini.generate_content(prompt)
-    return response.text
+    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}
+             ] }
+    
+    response = requests.post(url, headers=headers, json=body)
+    return response.json()['choices'][0]['message']['content']
 
 async def predict(message, history):
     full_response = ""
@@ -108,10 +118,45 @@ async def predict(message, history):
     yield full_response
 
     full_response += "\nResponse: "
-    streaming_response = gemini.generate_content(prompt, stream=True)
-    for sr in streaming_response:
-        full_response += sr.text
-        yield full_response
+    
+    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}
+             ] }
+
+    async with aiohttp.ClientSession() as session:
+        async with session.post(url, headers=headers, json=body) as response:
+            buffer = ""  # A buffer to hold incomplete lines of data
+            async for chunk in response.content.iter_any():
+                buffer += chunk.decode()
+                while "\n" in buffer:  # Process as long as there are complete lines in the buffer
+                    line, buffer = buffer.split("\n", 1)
+
+                    if line.startswith("data: "):
+                        event_data = line[len("data: "):]
+                        if event_data != '[DONE]':
+                            try:
+                                current_text = json.loads(event_data)['choices'][0]['delta']['content']
+                                full_response += current_text
+                                yield full_response
+                                await asyncio.sleep(0.01)
+                            except Exception:
+                                try:
+                                    current_text = json.loads(event_data)['choices'][0]['text']
+                                    full_response += current_text
+                                    yield full_response
+                                    await asyncio.sleep(0.01)
+                                except Exception:
+                                    pass
 
 gr.ChatInterface(
     predict,