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 from webtextcrawler.webtextcrawler import extract_text_from_url import gradio as gr from googlesearch import search torch.set_num_threads(2) openrouter_key = os.environ.get("sk-proj-sbgYj5kgnU35y0xjMSEyT3BlbkFJRMigEKIR9YdLqyx4y5bD") model = EmbeddingModel(use_quantized_onnx_model=True) def fetch_links(query, max_results=10): return list(search(query, num_results=max_results)) def fetch_texts(links): with multiprocessing.Pool(10) as pool: texts = pool.map(extract_text_from_url, links) return '\n'.join([t for t in texts if t]) def index_and_search(query, text): start = time.time() query_embedding = model.extract_embeddings(query) # Indexing vector_db = VectorDatabase() 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) retrieval_time = time.time() - start return '\n'.join([s['sentence'] for s in search_results[2]]), embedding_time, retrieval_time 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." 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) 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 = "" 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() 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 full_response += f"Text extraction took: {webcrawl_time:.4f} seconds\n" full_response += "\nIndexing in vector database and building prompt...\n" yield full_response context, embedding_time, retrieval_time = index_and_search(message, text) 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{context}\n\nAnswer: \"{message}\"\n(You can use the context to answer)\n(Your answer should be complete, detailed and well-structured)" 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 full_response += "\nGenerating response...\n" 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} ] } 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, title="Live Web Chat", description="", retry_btn=None, undo_btn=None, examples=[ 'What is the current sentiment of the Brazil election?', 'Compare the current economies of China and India?', 'What are new shoe design trends in 2024', ] ).launch()