158 lines
4.8 KiB
Python
158 lines
4.8 KiB
Python
import os
|
|
import json
|
|
import numpy as np
|
|
from openai import OpenAI
|
|
import discord
|
|
from discord.ext import commands
|
|
|
|
# --- Configuration ---
|
|
CONFIG_PATH = os.getenv("CONFIG_PATH", "config.json")
|
|
KNOWLEDGE_BASE_PATH = os.getenv(
|
|
"KNOWLEDGE_BASE_PATH", "bot/data/dbt_knowledge.embeddings.json"
|
|
)
|
|
DISCORD_BOT_TOKEN = os.getenv("DISCORD_BOT_TOKEN")
|
|
|
|
|
|
class SimpleVectorStore:
|
|
"""A simple in-memory vector store using NumPy."""
|
|
|
|
def __init__(self):
|
|
self.vectors = []
|
|
self.metadata = []
|
|
|
|
def add(self, vectors, metadatas):
|
|
self.vectors.extend(vectors)
|
|
self.metadata.extend(metadatas)
|
|
|
|
def search(self, query_vector, top_k=5):
|
|
if not self.vectors:
|
|
return []
|
|
|
|
query_vec = np.array(query_vector)
|
|
doc_vecs = np.array(self.vectors)
|
|
norms = np.linalg.norm(doc_vecs, axis=1)
|
|
valid_indices = norms > 0
|
|
scores = np.zeros(len(doc_vecs))
|
|
dot_products = np.dot(doc_vecs, query_vec)
|
|
scores[valid_indices] = dot_products[valid_indices] / (
|
|
norms[valid_indices] * np.linalg.norm(query_vec)
|
|
)
|
|
top_indices = np.argsort(scores)[-top_k:][::-1]
|
|
|
|
results = []
|
|
for idx in top_indices:
|
|
results.append({"metadata": self.metadata[idx], "score": scores[idx]})
|
|
return results
|
|
|
|
|
|
class JurySystem:
|
|
def __init__(self):
|
|
self.config = self.load_config()
|
|
self.client = OpenAI(
|
|
base_url="https://openrouter.ai/api/v1",
|
|
api_key=self.config["openrouter_api_key"],
|
|
)
|
|
self.vector_store = SimpleVectorStore()
|
|
self.load_knowledge_base()
|
|
|
|
def load_config(self):
|
|
with open(CONFIG_PATH, "r") as f:
|
|
return json.load(f)
|
|
|
|
def load_knowledge_base(self):
|
|
print(f"Loading knowledge base from {KNOWLEDGE_BASE_PATH}...")
|
|
try:
|
|
with open(KNOWLEDGE_BASE_PATH, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
vectors = []
|
|
metadata = []
|
|
for item in data:
|
|
vectors.append(item["embedding"])
|
|
metadata.append(
|
|
{"id": item["id"], "source": item["source"], "text": item["text"]}
|
|
)
|
|
self.vector_store.add(vectors, metadata)
|
|
print(f"Loaded {len(vectors)} chunks into vector store.")
|
|
except FileNotFoundError:
|
|
print(f"Error: {KNOWLEDGE_BASE_PATH} not found.")
|
|
raise
|
|
except Exception as e:
|
|
print(f"Error loading knowledge base: {e}")
|
|
raise
|
|
|
|
def process_query(self, query):
|
|
try:
|
|
response = self.client.embeddings.create(
|
|
model="qwen/qwen3-embedding-8b", input=query
|
|
)
|
|
query_emb = response.data[0].embedding
|
|
context_chunks = self.vector_store.search(query_emb, top_k=5)
|
|
|
|
if not context_chunks:
|
|
return "I couldn't find any relevant information in the knowledge base."
|
|
|
|
context_text = "\n\n---\n\n".join(
|
|
[chunk["metadata"]["text"] for chunk in context_chunks]
|
|
)
|
|
|
|
system_prompt = """You are a helpful AI assistant specializing in DBT (Dialectical Behavior Therapy).
|
|
Use the provided context to answer the user's question.
|
|
If the answer is not in the context, say you don't know based on the provided text.
|
|
Be concise and compassionate."""
|
|
|
|
user_prompt = f"Context:\n{context_text}\n\nQuestion: {query}"
|
|
|
|
response = self.client.chat.completions.create(
|
|
model="openai/gpt-4o-mini",
|
|
messages=[
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": user_prompt},
|
|
],
|
|
temperature=0.7,
|
|
)
|
|
|
|
return response.choices[0].message.content
|
|
except Exception as e:
|
|
return f"Error processing query: {e}"
|
|
|
|
|
|
# Initialize the Jury System
|
|
print("Initializing AI Jury System...")
|
|
jury_system = JurySystem()
|
|
print("Jury System ready!")
|
|
|
|
# Discord Bot Setup
|
|
intents = discord.Intents.default()
|
|
intents.message_content = True
|
|
bot = commands.Bot(command_prefix="!", intents=intents)
|
|
|
|
|
|
@bot.event
|
|
async def on_ready():
|
|
print(f"Bot logged in as {bot.user}")
|
|
|
|
|
|
@bot.event
|
|
async def on_message(message):
|
|
if message.author == bot.user:
|
|
return
|
|
|
|
# Process all messages as DBT queries
|
|
if not message.content.startswith("!"):
|
|
async with message.channel.typing():
|
|
response = jury_system.process_query(message.content)
|
|
await message.reply(response)
|
|
|
|
await bot.process_commands(message)
|
|
|
|
|
|
@bot.command(name="ask")
|
|
async def ask_dbt(ctx, *, question):
|
|
"""Ask a DBT-related question"""
|
|
async with ctx.typing():
|
|
response = jury_system.process_query(question)
|
|
await ctx.send(response)
|
|
|
|
|
|
bot.run(DISCORD_BOT_TOKEN)
|