Spaces:
Sleeping
Sleeping
File size: 5,509 Bytes
a91d88d | 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 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | """
Human+ RAG Engine
Loads knowledge base (.md files) → chunks → FAISS vector store.
Cached with st.cache_resource so it only runs once per deployment.
"""
import streamlit as st
from pathlib import Path
from typing import Optional
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_core.documents import Document
# ── Config ───────────────────────────────────────────────────────────
KNOWLEDGE_DIR = Path("/app/knowledge")
KNOWLEDGE_DIR_DEV = Path("knowledge") # fallback saat dev lokal
EMBEDDING_MODEL = "sentence-transformers/all-mpnet-base-v2"
CHUNK_SIZE = 800
CHUNK_OVERLAP = 150
SEPARATORS = ["\n\n## ", "\n\n### ", "\n\n", "\n", " ", "."]
# Retriever config (MMR = Maximal Marginal Relevance — kurangi redundansi)
RETRIEVER_K = 8
RETRIEVER_FETCH_K = 20
RETRIEVER_LAMBDA = 0.7 # 0 = diversity, 1 = relevance
class RAGEngineError(Exception):
"""Raised when RAG engine fails to load or retrieve."""
pass
@st.cache_resource(show_spinner=False, ttl=7200)
def get_vector_store() -> FAISS:
"""
Load knowledge base dan build FAISS vector store.
Cached selama 2 jam — auto-reload jika knowledge base diupdate.
Returns:
FAISS vector store yang siap diquery
Raises:
RAGEngineError: jika folder atau dokumen tidak ditemukan
"""
knowledge_dir = _resolve_knowledge_dir()
documents = _load_markdown_documents(knowledge_dir)
chunks = _split_documents(documents)
embeddings = _build_embeddings()
vector_store = FAISS.from_documents(chunks, embeddings)
return vector_store
def retrieve_context(query: str, vector_store: Optional[FAISS] = None) -> str:
"""
Retrieve relevant knowledge chunks untuk query biomarker/protocol.
Args:
query: teks query (biasanya extracted biomarkers)
vector_store: optional — jika None, akan di-load otomatis
Returns:
context string yang siap dimasukkan ke prompt
"""
if vector_store is None:
vector_store = get_vector_store()
retriever = vector_store.as_retriever(
search_type="mmr",
search_kwargs={
"k": RETRIEVER_K,
"fetch_k": RETRIEVER_FETCH_K,
"lambda_mult": RETRIEVER_LAMBDA,
},
)
try:
docs = retriever.invoke(query)
except Exception as e:
raise RAGEngineError(f"Retrieval gagal: {e}") from e
context = "\n\n".join(
f"[Sumber: {doc.metadata.get('source', 'unknown')}]\n{doc.page_content}"
for doc in docs
)
return context
def get_knowledge_stats() -> dict:
"""Return stats tentang knowledge base yang ter-load (untuk debugging)."""
try:
knowledge_dir = _resolve_knowledge_dir()
files = list(knowledge_dir.glob("*.md"))
return {
"ok": True,
"dir": str(knowledge_dir),
"n_files": len(files),
"filenames": [f.name for f in files],
}
except RAGEngineError as e:
return {"ok": False, "error": str(e)}
# ── Private helpers ──────────────────────────────────────────────────
def _resolve_knowledge_dir() -> Path:
"""Cari knowledge directory — Docker path atau dev local."""
for candidate in [KNOWLEDGE_DIR, KNOWLEDGE_DIR_DEV]:
if candidate.exists() and candidate.is_dir():
return candidate
raise RAGEngineError(
f"Folder knowledge tidak ditemukan. "
f"Dicari di: {KNOWLEDGE_DIR} dan {KNOWLEDGE_DIR_DEV}"
)
def _load_markdown_documents(knowledge_dir: Path) -> list[Document]:
"""Load semua file .md dari knowledge_dir sebagai LangChain Documents."""
documents = []
for file_path in sorted(knowledge_dir.glob("*.md")):
try:
content = file_path.read_text(encoding="utf-8").strip()
if content:
documents.append(
Document(
page_content=content,
metadata={
"source": file_path.name,
"filepath": str(file_path),
},
)
)
except Exception as e:
# Log tapi jangan stop — satu file gagal tidak perlu matikan semua
st.warning(f"⚠️ Gagal membaca {file_path.name}: {e}")
if not documents:
raise RAGEngineError(
f"Tidak ada dokumen .md yang berhasil dibaca dari {knowledge_dir}"
)
return documents
def _split_documents(documents: list[Document]) -> list[Document]:
"""Split documents menjadi chunks yang optimal untuk embedding."""
splitter = RecursiveCharacterTextSplitter(
chunk_size=CHUNK_SIZE,
chunk_overlap=CHUNK_OVERLAP,
separators=SEPARATORS,
)
return splitter.split_documents(documents)
def _build_embeddings() -> HuggingFaceEmbeddings:
"""Build HuggingFace embeddings model."""
return HuggingFaceEmbeddings(
model_name=EMBEDDING_MODEL,
model_kwargs={"device": "cpu"},
encode_kwargs={"normalize_embeddings": True},
)
|