Spaces:
Runtime error
Runtime error
File size: 3,158 Bytes
1a0b19c 26060c4 1a0b19c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | import asyncio
import os
import discord
import requests
from discord.ext import commands
from hf_client import call_hf
from index_local_runtime import LocalIndexRuntime
from prompts import SYSTEM_PROMPT, build_user_prompt
DISCORD_TOKEN = os.getenv("DISCORD_TOKEN")
BOT_PREFIX = os.getenv("BOT_PREFIX", "!")
REINDEX_API_TOKEN = os.getenv("REINDEX_API_TOKEN")
if not DISCORD_TOKEN:
raise RuntimeError("DISCORD_TOKEN nao definido.")
index_rt = LocalIndexRuntime()
intents = discord.Intents.default()
intents.message_content = True
bot = commands.Bot(command_prefix=BOT_PREFIX, intents=intents)
def _format_context(hits):
return "\n\n---\n\n".join(
[f"[{h['source']}] (score={h['score']:.3f})\n{h['text']}" for h in hits]
)
async def _build_answer(question: str) -> str:
index_rt.maybe_reload()
hits = index_rt.search(question, k=4)
if not hits:
return "Nao encontrei isso nos documentos."
context = _format_context(hits)
messages = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": build_user_prompt(question, context)},
]
return await asyncio.to_thread(call_hf, messages)
@bot.event
async def on_ready():
print(f"[BOT] logged in as {bot.user}")
if index_rt.exists():
index_rt.load()
else:
print("[BOT] indice ainda nao existe; rode !reindex")
@bot.command(name="rag")
async def rag_cmd(ctx, *, question: str):
try:
answer = await _build_answer(question)
await ctx.reply(answer[:1900])
except Exception as exc: # noqa: BLE001
await ctx.reply(f"Falha ao responder: {type(exc).__name__}: {exc}")
@bot.command(name="reindex")
@commands.has_permissions(administrator=True)
async def reindex_cmd(ctx):
await ctx.reply("Iniciando reindex... (pode demorar)")
try:
headers = {}
if REINDEX_API_TOKEN:
headers["Authorization"] = f"Bearer {REINDEX_API_TOKEN}"
r = await asyncio.to_thread(
requests.post,
"http://127.0.0.1:7860/reindex",
headers=headers,
timeout=3600,
)
txt = r.text
await ctx.reply(("Reindex concluido\n" + txt)[:1900])
if index_rt.exists():
index_rt.load()
except Exception as exc: # noqa: BLE001
await ctx.reply(f"Falha no reindex: {type(exc).__name__}: {exc}")
@bot.event
async def on_message(message: discord.Message):
if message.author.bot:
return
await bot.process_commands(message)
if bot.user and bot.user in message.mentions:
content = (
message.content.replace(f"<@{bot.user.id}>", "")
.replace(f"<@!{bot.user.id}>", "")
.strip()
)
if not content:
await message.reply("Me mencione e pergunte. Ex: @bot como faco X?")
return
try:
answer = await _build_answer(content)
await message.reply(answer[:1900])
except Exception as exc: # noqa: BLE001
await message.reply(f"Falha ao responder: {type(exc).__name__}: {exc}")
def run_bot():
bot.run(DISCORD_TOKEN)
|