Add Gradio app with RAG pipeline for chemical regulation lookup

.gitignore

@@ -0,0 +1,3 @@
+vector_db/
+.env
+__pycache__/

app.py

@@ -0,0 +1,101 @@
+import gradio as gr
+from core import RAG, parse_cas_input
+
+rag = RAG()
+
+EXAMPLE_INPUTS = [
+    "7664-39-3",
+    "7664-39-3, 128-37-0",
+    "50-00-0, 7664-93-9, 67-64-1",
+]
+
+
+def analyse(raw_input: str, provider: str):
+    if not raw_input or not raw_input.strip():
+        raise gr.Error("Vui lòng nhập ít nhất một mã CAS.")
+
+    cas_numbers = parse_cas_input(raw_input)
+    if not cas_numbers:
+        raise gr.Error(
+            f"Không tìm thấy mã CAS hợp lệ trong '{raw_input}'. "
+            "Định dạng đúng: 7664-39-3"
+        )
+
+    provider_key = "google" if "Gemini" in provider else "openai"
+    result = rag.pipeline(cas_numbers, provider=provider_key)
+
+    rows = []
+    for r in result.results:
+        rows.append([r.casNumber, r.chemicalName, r.status, r.reason])
+
+    summary_parts = []
+    for r in result.results:
+        summary_parts.append(
+            f"### {r.casNumber} — {r.chemicalName}\n"
+            f"**Yêu cầu pháp lý:** {r.status}\n\n"
+            f"**Cơ sở:** {r.reason}\n"
+        )
+    summary_md = "\n---\n".join(summary_parts)
+
+    return rows, summary_md
+
+
+with gr.Blocks(
+    title="Chemical & Precursor Declaration Checker",
+    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky"),
+    css="""
+        .main-header {text-align:center; margin-bottom:4px}
+        .sub-header  {text-align:center; color:#555; margin-top:0}
+    """,
+) as demo:
+    gr.Markdown(
+        "<h1 class='main-header'>Chemical & Precursor Declaration Checker</h1>"
+        "<p class='sub-header'>"
+        "Tra cứu nghĩa vụ khai báo hóa chất & tiền chất theo quy định Việt Nam"
+        "</p>"
+    )
+
+    with gr.Row():
+        with gr.Column(scale=3):
+            cas_input = gr.Textbox(
+                label="Nhập mã CAS",
+                placeholder="VD: 7664-39-3, 128-37-0  (phân cách bằng dấu phẩy hoặc xuống dòng)",
+                lines=2,
+            )
+        with gr.Column(scale=1, min_width=180):
+            provider_dd = gr.Dropdown(
+                choices=["OpenAI GPT-4.1-mini", "Google Gemini 2.5 Flash"],
+                value="OpenAI GPT-4.1-mini",
+                label="LLM Provider",
+            )
+            submit_btn = gr.Button("Tra cứu", variant="primary", size="lg")
+
+    gr.Examples(
+        examples=EXAMPLE_INPUTS,
+        inputs=cas_input,
+        label="Ví dụ",
+    )
+
+    gr.Markdown("### Kết quả")
+
+    results_table = gr.Dataframe(
+        headers=["CAS", "Tên hóa chất", "Trạng thái", "Kết quả tra cứu"],
+        datatype=["str", "str", "str", "str"],
+        interactive=False,
+        wrap=True,
+    )
+    detail_md = gr.Markdown(label="Chi tiết")
+
+    submit_btn.click(
+        fn=analyse,
+        inputs=[cas_input, provider_dd],
+        outputs=[results_table, detail_md],
+    )
+    cas_input.submit(
+        fn=analyse,
+        inputs=[cas_input, provider_dd],
+        outputs=[results_table, detail_md],
+    )
+
+if __name__ == "__main__":
+    demo.launch()

core/__init__.py

@@ -0,0 +1,3 @@
+from .rag import RAG, parse_cas_input
+
+__all__ = ["RAG", "parse_cas_input"]

core/param.py

@@ -0,0 +1,28 @@
+from pydantic import BaseModel, Field
+from typing import List
+
+
+class AnalysisResult(BaseModel):
+    """Single CAS analysis result returned by the LLM."""
+    casNumber: str = Field(..., description="Mã CAS (e.g. 128-37-0)")
+    chemicalName: str = Field(..., description="The only name of the chemical")
+    status: str = Field(
+        ...,
+        description=(
+            "Yêu cầu pháp lý cụ thể — lấy nguyên văn phần sau dấu '=>' trong QUY ĐỊNH ÁP DỤNG. "
+            "Nếu CAS xuất hiện ở nhiều phụ lục, đánh số từng yêu cầu. "
+            "VD: 'Chỉ được Nhập khẩu hoặc Xuất khẩu nếu có Giấy chứng nhận...'"
+        ),
+    )
+    reason: str = Field(
+        ...,
+        description=(
+            "Kết quả tra cứu: ghi rõ vị trí trong văn bản pháp luật. "
+            "VD: 'STT 15 của Phụ lục II Nghị định 24/2026/NĐ-CP: Danh mục hoá chất sản xuất, kinh doanh có điều kiện'"
+        ),
+    )
+
+
+class CASAnalysisOutput(BaseModel):
+    """Structured output schema for LLM CAS analysis."""
+    results: List[AnalysisResult] = Field(..., description="Danh sách kết quả phân tích cho từng mã CAS")

core/prompt.py

@@ -0,0 +1,45 @@
+SYSTEM_PROMPT3 = """Bạn là trợ lý phân tích pháp luật hóa chất.
+========================
+QUY TẮC TRÍCH XUẤT
+========================
+1. casNumber
+   - Ghi đúng mã CAS đang phân tích.
+2. chemicalName
+   - Chỉ lấy "Tên chất" (cột tên chính trong bảng).
+   - Không lấy tên khoa học (IUPAC).
+   - Không lấy nội dung trong ngoặc.
+   - Không tự suy đoán.
+3. status (CHỈ 1 GIÁ TRỊ DUY NHẤT)
+   ƯU TIÊN THEO THỨ TỰ:
+
+   A. Nếu có nội dung nghĩa vụ hành chính cụ thể
+      (có "=>", hoặc chứa từ như: "Phải", "Chỉ được", "Cấm", "Xin giấy phép"...)
+      → Lấy nguyên văn nội dung nghĩa vụ đó làm status.
+
+   B. Nếu không có nghĩa vụ cụ thể nhưng CAS thuộc DANH MỤC (I, II, III, IV…)
+      → Lấy nguyên văn mô tả quản lý của mục con (ví dụ IVA, IVB…).
+      → Format:
+         "Thuộc danh mục IVA: [nguyên văn mô tả]"
+
+   C. Nếu không tìm thấy trong tài liệu
+      → "Không có trong danh sách"
+
+   - Không gộp nhiều status.
+   - Không lấy nhiều nguồn cho status.
+   - Chỉ được có 1 status chính.
+
+4. reason (LIỆT KÊ ĐẦY ĐỦ VỊ TRÍ)
+
+   - Liệt kê TẤT CẢ vị trí CAS xuất hiện.
+   - Format mỗi vị trí:
+     "STT [số] của [TÊN ĐẦY ĐỦ PHỤ LỤC/DANH MỤC CHÍNH]"
+   - Luôn dùng tiêu đề chính (ví dụ: DANH MỤC IV), không dùng IVA trong reason.
+   - Ngăn cách các vị trí bằng dấu chấm và khoảng trắng.
+========================
+LƯU Ý QUAN TRỌNG
+========================
+- status = nghĩa vụ phải làm (nếu có).
+- reason = truy vết đầy đủ tất cả văn bản.
+- Nếu có cả nghĩa vụ hành chính và danh mục:
+   → status lấy nghĩa vụ.
+   → reason ghi cả hai vị trí."""

core/rag.py

@@ -0,0 +1,69 @@
+import os
+import re
+from typing import List, Dict
+
+from dotenv import load_dotenv
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_openai import ChatOpenAI
+
+from .vectorstore import VectorStore
+from .utils import build_grouped_context, _build_user_prompt
+from .param import CASAnalysisOutput
+from .prompt import SYSTEM_PROMPT3
+
+load_dotenv()
+
+CAS_PATTERN = re.compile(r"\d{2,7}-\d{2}-\d")
+
+
+def parse_cas_input(raw: str) -> List[str]:
+    """Extract valid CAS numbers from free-form user input."""
+    return CAS_PATTERN.findall(raw.strip())
+
+
+class RAG:
+    def __init__(self):
+        self.vector_store = VectorStore()
+
+        google_key = os.getenv("GOOGLE_API_KEY")
+        openai_key = os.getenv("OPENAI_API_KEY")
+
+        self.llm_google = (
+            ChatGoogleGenerativeAI(model="gemini-2.5-flash", api_key=google_key, temperature=0)
+            if google_key else None
+        )
+        self.llm_openai = (
+            ChatOpenAI(model="gpt-4.1-mini", api_key=openai_key, temperature=0)
+            if openai_key else None
+        )
+
+    def _get_llm(self, provider: str = "openai"):
+        if provider == "google" and self.llm_google:
+            return self.llm_google
+        if provider == "openai" and self.llm_openai:
+            return self.llm_openai
+        available = self.llm_openai or self.llm_google
+        if available is None:
+            raise RuntimeError("No LLM API key configured. Set OPENAI_API_KEY or GOOGLE_API_KEY.")
+        return available
+
+    def retrieve_build_context(self, cas_numbers: List[str]) -> str:
+        chunks_by_cas = self.vector_store.search_per_cas(cas_numbers)
+        context_chunking = build_grouped_context(chunks_by_cas)
+        user_prompt = _build_user_prompt(cas_numbers, context_chunking)
+        return user_prompt
+
+    def pipeline(self, cas_numbers: List[str], provider: str = "openai") -> CASAnalysisOutput:
+        user_prompt = self.retrieve_build_context(cas_numbers)
+
+        prompt = ChatPromptTemplate.from_messages([
+            ("system", SYSTEM_PROMPT3),
+            ("user", "{user_prompt}")
+        ])
+
+        llm = self._get_llm(provider)
+        llm_structured = llm.with_structured_output(CASAnalysisOutput, method="function_calling")
+        chain = prompt | llm_structured
+        result = chain.invoke({"user_prompt": user_prompt})
+        return result

core/utils.py

@@ -0,0 +1,25 @@
+from typing import List
+
+
+def build_grouped_context(results: dict[str, list[dict]]) -> str:
+    blocks = []
+
+    for cas, chunks in results.items():
+        cas_block = [f"========================\n[CAS: {cas}]\n"]
+
+        for i, item in enumerate(chunks, 1):
+            cas_block.append(f"\nCHUNK {i}:\n")
+            cas_block.append(item["content"])
+            cas_block.append("\n------------------------")
+
+        blocks.append("".join(cas_block))
+
+    return "\n\n".join(blocks)
+
+
+def _build_user_prompt(cas_numbers: List[str], context: str) -> str:
+    cas_list = ", ".join(cas_numbers)
+    return (
+        f"Phân tích các mã CAS sau: {cas_list}\n\n"
+        f"CONTEXT:\n{context}"
+    )

core/vectorstore.py

@@ -0,0 +1,239 @@
+import os
+import shutil
+from typing import List, Dict
+import chromadb
+from dotenv import load_dotenv
+
+load_dotenv()
+
+DEFAULT_DB_PATH = "./vector_db/chroma_db"
+
+
+def _ensure_vectordb(db_path: str) -> str:
+    """Download ChromaDB from HF Dataset if not present locally."""
+    if os.path.isdir(db_path) and os.listdir(db_path):
+        return db_path
+
+    repo_id = os.getenv("HF_DATASET_REPO")
+    if not repo_id:
+        raise RuntimeError(
+            f"Vector DB not found at '{db_path}' and HF_DATASET_REPO is not set. "
+            "Either place the DB locally or set HF_DATASET_REPO=your-username/chemical-vectordb"
+        )
+
+    from huggingface_hub import snapshot_download
+
+    print(f"   Downloading vector DB from HF dataset: {repo_id} ...")
+    downloaded = snapshot_download(
+        repo_id=repo_id,
+        repo_type="dataset",
+        token=os.getenv("HF_TOKEN"),
+    )
+
+    os.makedirs(db_path, exist_ok=True)
+    for item in os.listdir(downloaded):
+        src = os.path.join(downloaded, item)
+        dst = os.path.join(db_path, item)
+        if item.startswith("."):
+            continue
+        if os.path.isdir(src):
+            shutil.copytree(src, dst, dirs_exist_ok=True)
+        else:
+            shutil.copy2(src, dst)
+
+    print(f"   Vector DB ready at: {db_path}")
+    return db_path
+
+
+class VectorStore:
+    """ChromaDB vector store with hybrid search and E5 prefix support."""
+
+    EXACT_MATCH_BOOST = 0.3
+
+    def __init__(self):
+        self.collection_name = os.getenv('COLLECTION_NAME', 'chemical_regulations')
+        self.vector_db_path = os.getenv('VECTOR_DB_PATH', DEFAULT_DB_PATH)
+        self.lightweight = True
+
+        self.model_name = os.getenv('EMBEDDING_MODEL', 'intfloat/multilingual-e5-base')
+        self._use_prefix = 'e5' in self.model_name.lower()
+        self.model = None
+
+        if self.lightweight:
+            print(f"   Lightweight mode — exact match only, no embedding model")
+        else:
+            from sentence_transformers import SentenceTransformer
+            print(f"   Loading embedding model: {self.model_name}")
+            self.model = SentenceTransformer(self.model_name)
+
+        self.vector_db_path = _ensure_vectordb(self.vector_db_path)
+        self.client = chromadb.PersistentClient(path=self.vector_db_path)
+
+        try:
+            self.collection = self.client.get_collection(self.collection_name)
+            print(f"   Loaded collection: {self.collection_name} "
+                  f"({self.collection.count()} vectors)")
+        except Exception as e:
+            raise RuntimeError(
+                f"Cannot load collection '{self.collection_name}'. "
+                f"Run scripts/2_create_embeddings.py first. Error: {e}"
+            )
+
+    def _encode_query(self, text: str):
+        """Encode with 'query: ' prefix for E5 models."""
+        if self.model is None:
+            raise RuntimeError(
+                "Semantic search unavailable in lightweight mode. "
+                "Set LIGHTWEIGHT_MODE=false or unset it to enable."
+            )
+        if self._use_prefix:
+            text = f"query: {text}"
+        return self.model.encode(text, normalize_embeddings=True)
+
+    def semantic_search(self, query: str, top_k: int = 10) -> List[Dict]:
+        """Embedding-based semantic search."""
+        query_embedding = self._encode_query(query)
+
+        results = self.collection.query(
+            query_embeddings=[query_embedding.tolist()],
+            n_results=top_k,
+            include=["documents", "metadatas", "distances"],
+        )
+
+        formatted = []
+        if results['ids'] and results['ids'][0]:
+            for i in range(len(results['ids'][0])):
+                formatted.append({
+                    "id": results['ids'][0][i],
+                    "content": results['documents'][0][i],
+                    "metadata": results['metadatas'][0][i],
+                    "similarity": 1 - results['distances'][0][i],
+                    "match_type": "semantic",
+                })
+        return formatted
+
+    def exact_match_search(self, cas_number: str, top_k: int = 5) -> List[Dict]:
+        """
+        True lexical match — find chunks whose document text literally
+        contains the CAS number string (ChromaDB where_document $contains).
+        """
+        try:
+            results = self.collection.get(
+                where_document={"$contains": cas_number},
+                include=["documents", "metadatas"],
+                limit=top_k,
+            )
+        except Exception:
+            return []
+
+        formatted = []
+        if results['ids']:
+            for i in range(len(results['ids'])):
+                formatted.append({
+                    "id": results['ids'][i],
+                    "content": results['documents'][i],
+                    "metadata": results['metadatas'][i],
+                    "similarity": 1.0,
+                    "match_type": "exact",
+                })
+        return formatted
+
+    def hybrid_search(self, cas_numbers: List[str], top_k: int = 10) -> List[Dict]:
+        """
+        Hybrid search strategy:
+          1. Exact-match per CAS number  (highest fidelity)
+          2. Semantic search with expanded query (broader context) — skipped in lightweight mode
+          3. Deduplicate -> rerank -> top_k
+        """
+        all_results: List[Dict] = []
+        seen_ids: set = set()
+
+        for cas in cas_numbers:
+            for r in self.exact_match_search(cas, top_k=5):
+                if r['id'] not in seen_ids:
+                    all_results.append(r)
+                    seen_ids.add(r['id'])
+
+        if not self.lightweight:
+            expanded = self._expand_query(cas_numbers)
+            for r in self.semantic_search(expanded, top_k=top_k):
+                if r['id'] not in seen_ids:
+                    all_results.append(r)
+                    seen_ids.add(r['id'])
+
+        all_results = self._rerank(all_results, cas_numbers)
+        return all_results[:top_k]
+
+    def search_per_cas(
+        self, cas_numbers: List[str], top_k_per_cas: int = 5,
+    ) -> Dict[str, List[Dict]]:
+        """
+        Search for each CAS number individually so every CAS gets its own
+        dedicated chunk slots (no competition between CAS numbers).
+        """
+        results_by_cas: Dict[str, List[Dict]] = {cas: [] for cas in cas_numbers}
+        seen_per_cas: Dict[str, set] = {cas: set() for cas in cas_numbers}
+
+        for cas in cas_numbers:
+            for r in self.exact_match_search(cas, top_k=top_k_per_cas):
+                if r['id'] not in seen_per_cas[cas]:
+                    results_by_cas[cas].append(r)
+                    seen_per_cas[cas].add(r['id'])
+
+        if not self.lightweight:
+            expanded = self._expand_query(cas_numbers)
+            for r in self.semantic_search(expanded, top_k=top_k_per_cas * len(cas_numbers)):
+                for cas in cas_numbers:
+                    in_content = cas in r['content']
+                    in_meta = cas in r['metadata'].get('cas_numbers', '')
+                    in_section_cas = cas in r['metadata'].get('section_cas_numbers', '')
+                    if (in_content or in_meta or in_section_cas) and r['id'] not in seen_per_cas[cas]:
+                        results_by_cas[cas].append(r)
+                        seen_per_cas[cas].add(r['id'])
+
+        for cas in cas_numbers:
+            results_by_cas[cas] = self._rerank(results_by_cas[cas], [cas])
+            results_by_cas[cas] = results_by_cas[cas][:top_k_per_cas]
+
+        return results_by_cas
+
+    def _expand_query(self, cas_numbers: List[str]) -> str:
+        """Build a richer query so semantic search captures regulatory context."""
+        cas_list = ', '.join(cas_numbers)
+        return (
+            f"Thông tin quy định pháp luật về hóa chất có mã CAS: {cas_list}. "
+            f"Giấy phép nhập khẩu, khai báo hóa chất, tiền chất, danh mục quản lý."
+        )
+
+    def _rerank(self, results: List[Dict], cas_numbers: List[str]) -> List[Dict]:
+        """Boost chunks that literally contain one of the queried CAS numbers."""
+        for r in results:
+            boost = 0.0
+            content = r['content']
+            cas_meta = r['metadata'].get('cas_numbers', '')
+            section_cas_meta = r['metadata'].get('section_cas_numbers', '')
+
+            for cas in cas_numbers:
+                if cas in content:
+                    boost += self.EXACT_MATCH_BOOST
+                if cas in cas_meta:
+                    boost += self.EXACT_MATCH_BOOST * 0.5
+                elif cas in section_cas_meta:
+                    boost += self.EXACT_MATCH_BOOST * 0.25
+
+            r['similarity'] = min(1.0, r['similarity'] + boost)
+
+        results.sort(key=lambda x: x['similarity'], reverse=True)
+        return results
+
+    def get_stats(self) -> Dict:
+        stats = {
+            "total_chunks": self.collection.count(),
+            "embedding_model": self.model_name,
+            "collection_name": self.collection_name,
+            "e5_prefix_enabled": self._use_prefix,
+            "lightweight_mode": self.lightweight,
+        }
+        if self.model is not None:
+            stats["embedding_dimension"] = self.model.get_sentence_embedding_dimension()
+        return stats

requirements.txt

@@ -0,0 +1,10 @@
+pypdf
+langchain_google_genai
+langchain_openai
+langchain_community
+langchain_core
+chromadb
+python-dotenv
+pydantic
+gradio
+huggingface_hub

scripts/upload_vectordb.py

@@ -0,0 +1,44 @@
+"""
+Upload ChromaDB vector database to a Hugging Face Dataset repo.
+
+Usage:
+    python scripts/upload_vectordb.py --repo-id YOUR_USERNAME/chemical-vectordb
+
+Prerequisites:
+    pip install huggingface_hub
+    huggingface-cli login
+"""
+
+import argparse
+from huggingface_hub import HfApi
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Upload ChromaDB to HF Dataset")
+    parser.add_argument(
+        "--repo-id", required=True,
+        help="HF dataset repo id, e.g. your-username/chemical-vectordb",
+    )
+    parser.add_argument(
+        "--local-dir",
+        default="./vector_db/content/vector_db/chroma_db",
+        help="Local path to the chroma_db folder",
+    )
+    parser.add_argument("--private", action="store_true", help="Make the dataset private")
+    args = parser.parse_args()
+
+    api = HfApi()
+
+    api.create_repo(repo_id=args.repo_id, repo_type="dataset", private=args.private, exist_ok=True)
+    print(f"Repo ready: https://huggingface.co/datasets/{args.repo_id}")
+
+    api.upload_folder(
+        folder_path=args.local_dir,
+        repo_id=args.repo_id,
+        repo_type="dataset",
+    )
+    print(f"Upload complete! Files pushed to datasets/{args.repo_id}")
+
+
+if __name__ == "__main__":
+    main()