from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np
import os

kb_path = os.path.join(os.path.dirname(__file__), "mental_health_kb.txt") if '__file__' in globals() else "mental_health_kb.txt"

with open(kb_path, "r", encoding="utf-8") as f:
    docs = [d.strip() for d in f.read().split("\n\n") if d.strip()]

vectorizer = TfidfVectorizer()
tfidf_matrix = vectorizer.fit_transform(docs)

def retrieve(query):
    q_vec = vectorizer.transform([query])
    scores = cosine_similarity(q_vec, tfidf_matrix)[0]
    top_indices = np.argsort(scores)[::-1][:3]
    results = [docs[i] for i in top_indices if scores[i] > 0]
    return "\n\n".join(results)[:1500] if results else ""