Update app.py

app.py

@@ -1,22 +1,23 @@
-import time, aiohttp, asyncio, json, os, multiprocessing, torch
+import time, os, multiprocessing
 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
 
-torch.set_num_threads(2)
-
-openrouter_key = os.environ.get("OPENROUTER_KEY")
+gemini_key = os.environ.get("GEMINI_KEY")
+genai.configure(api_key=gemini_key)
+gemini = genai.GenerativeModel('gemini-pro')
 model = EmbeddingModel(use_quantized_onnx_model=True)
 
-def fetch_links(query, max_results=5):
+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)]
 
 def fetch_texts(links):
-    with multiprocessing.Pool(5) as pool:
+    with multiprocessing.Pool(10) as pool:
         texts = pool.map(extract_text_from_url, links)
     return '\n'.join([t for t in texts if t])
 
@@ -36,83 +37,80 @@ def index_and_search(query, text):
 
     # Retrieval
     start = time.time()
-    search_results = vector_db.find_most_similar(query_embedding, k = 12)
+    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
 
-def retrieval_pipeline(query):
+def generate_search_terms(message, lang):
+    if lang == Language.ptbr:
+        prompt = f"A partir do texto a seguir, gere alguns termos de pesquisa: \"{message}\"\nSua resposta deve ser apenas o termo de busca mais adequado, e nada mais."
+    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
+
+async def predict(message, history):
+    full_response = ""
+
+    query_language = detect_language(message)
+
+    start = time.time()
+    full_response += "Generating search terms...\n"
+    yield full_response
+    search_query = generate_search_terms(message, query_language)
+    search_terms_time = time.time() - start
+
+    full_response += f"Search terms: \"{search_query}\"\n"
+    yield full_response
+    full_response += f"Search terms took: {search_terms_time:.4f} seconds\n"
+    yield full_response
+
     start = time.time()
-    links = fetch_links(query)
+    full_response += "\nSearching the web...\n"
+    yield full_response
+    links = fetch_links(search_query)
     websearch_time = time.time() - start
 
+    full_response += f"Web search took: {websearch_time:.4f} seconds\n"
+    yield full_response
+
+    full_response += f"Links visited:\n"
+    yield full_response
+    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()
     text = fetch_texts(links)
     webcrawl_time = time.time() - start
 
-    context, embedding_time, retrieval_time = index_and_search(query, text)
+    full_response += f"Text extraction took: {webcrawl_time:.4f} seconds\n"
 
-    return context, websearch_time, webcrawl_time, embedding_time, retrieval_time, links
+    full_response += "\nIndexing in vector database and building prompt...\n"
+    yield full_response
 
-async def predict(message, history):
-    context, websearch_time, webcrawl_time, embedding_time, retrieval_time, links = retrieval_pipeline(message)
+    context, embedding_time, retrieval_time = index_and_search(message, text)
 
-    if detect_language(message) == Language.ptbr:
-        prompt = f"Contexto:\n\n{context}\n\nBaseado no contexto, responda: {message}"
+    if query_language == Language.ptbr:
+        prompt = f"Contexto:\n{context}\n\nResponda: \"{message}\"\n(Você pode utilizar o contexto para responder)\n(Sua resposta deve ser completa, detalhada e bem estruturada)"
     else:
-        prompt = f"Context:\n\n{context}\n\nBased on the context, answer: {message}"
-
-    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", 
-                    "nousresearch/nous-capybara-7b:free",
-                    "huggingfaceh4/zephyr-7b-beta:free"
-             ],
-             "route": "fallback",
-             "max_tokens": 768,
-             "messages": [
-                 {"role": "user", "content": prompt}
-             ] }
+        prompt = f"Context:\n{context}\n\nAnswer: \"{message}\"\n(You can use the context to answer)\n(Your answer should be complete, detailed and well-structured)"
 
-    full_response = ""
-    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
-    
-    final_metadata_block = ""
+    full_response += f"Embedding time: {embedding_time:.4f} seconds\n"
+    full_response += f"Retrieval from VectorDB time: {retrieval_time:.4f} seconds\n"
+    yield full_response
 
-    final_metadata_block += f"Links visited:\n"
-    for link in links:
-        final_metadata_block += f"{link}\n"
-    final_metadata_block += f"\nWeb search time: {websearch_time:.4f} seconds\n"
-    final_metadata_block += f"\nText extraction: {webcrawl_time:.4f} seconds\n"
-    final_metadata_block += f"\nEmbedding time: {embedding_time:.4f} seconds\n"
-    final_metadata_block += f"\nRetrieval from VectorDB time: {retrieval_time:.4f} seconds"
+    full_response += "\nGenerating response...\n"
+    yield full_response
 
-    yield f"{full_response}\n\n{final_metadata_block}"
+    full_response += "\nResponse: "
+    streaming_response = gemini.generate_content(prompt, stream=True)
+    for sr in streaming_response:
+        full_response += sr.text
+        yield full_response
 
 gr.ChatInterface(
     predict,