Added text image OCR

.gitignore

@@ -1,9 +1,17 @@
+# Python-generated files
 __pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Virtual environments
+.venv
 .env
-.agent/
-.venv/
-.idea/
-.vscode/
-app/services/query_preprocessor.py
-.ruff_cache
+.agent
 chunking_embedding_gguf.ipynb
+ocr_test.py
+.idea
+.vscode
+.gitignore

.python-version

@@ -1 +1 @@
-3.12.11
+3.13

app/core/config.py

@@ -22,13 +22,19 @@ class Settings(BaseSettings):
     supabase_service_role_key: str
 
     # --- AI Provider ---
-    ai_provider: Literal["openrouter", "groq"] = "groq"
-    openrouter_api_key: str = ""
-    groq_api_key: str = ""
+    openrouter_api_key: str
+    groq_api_key: str
+    model_openrouter: str
+    base_url_openrouter: str
+    model_groq: str
+
+    # --- OCR Model ---
+    ocr_model_name: str = "zai-org/GLM-OCR"
 
     # --- Embedding Model ---
-    model_name: str = "BAAI/bge-m3"
+    embedding_model_name: str = "BAAI/bge-m3"
     embedding_dimensions: int = 1024
+    device: str = "cuda"
 
     # --- Chunking ---
     chunk_size: int = 512
@@ -40,27 +46,30 @@ class Settings(BaseSettings):
     debug: bool = True
 
     @cached_property
-    def llm(self):
+    def answer_llm(self):
+        """Lazily initialize the LLM based on the configured provider."""
+        from langchain_groq import ChatGroq
+
+        return ChatGroq(
+            api_key=self.groq_api_key,
+            model=self.model_groq,
+            temperature=0,
+            max_tokens=None,
+            timeout=None,
+            max_retries=2,
+        )
+
+    @cached_property
+    def preprocess_llm(self):
         """Lazily initialize the LLM based on the configured provider."""
-        if self.ai_provider == "openrouter":
-            from langchain_openai import ChatOpenAI
-
-            return ChatOpenAI(
-                api_key=self.openrouter_api_key,
-                model="xiaomi/mimo-v2-flash:free",
-                base_url="https://openrouter.ai/api/v1",
-            )
-        else:
-            from langchain_groq import ChatGroq
-
-            return ChatGroq(
-                api_key=self.groq_api_key,
-                model="meta-llama/llama-4-scout-17b-16e-instruct",
-                temperature=0,
-                max_tokens=None,
-                timeout=None,
-                max_retries=2,
-            )
+        from langchain_openai import ChatOpenAI
+
+        return ChatOpenAI(
+            api_key=self.openrouter_api_key,
+            model=self.model_openrouter,
+            base_url=self.base_url_openrouter,
+        )
+    
 
 
 # Regex pattern to extract page numbers from blockquotes

app/core/database.py

@@ -1,7 +1,7 @@
 """Supabase database operations for RAG chunk storage and vector search."""
 
 import logging
-from typing import Any
+from typing import Any, List, Dict
 
 from supabase import Client, create_client
 
@@ -39,7 +39,7 @@ class Database:
             raise RuntimeError("Database not connected. Call connect() first.")
         return self.admin_client if admin else self.supabase
 
-    async def upsert_chunks(self, chunks: list[dict[str, Any]]) -> int:
+    async def upsert_chunks(self, chunks: List[Dict[str, Any]]) -> int:
         """Upsert document chunks into Supabase.
 
         Args:
@@ -79,9 +79,9 @@ class Database:
 
     async def vector_search(
         self,
-        query_embedding: list[float],
+        query_embedding: List[float],
         top_k: int = 6,
-    ) -> list[dict[str, Any]]:
+    ) -> List[Dict[str, Any]]:
         """Perform vector similarity search using Supabase RPC.
 
         Args:

app/core/prompt_templates.py

@@ -1,16 +1,104 @@
-"""
-Prompt templates for the chat.
-"""
+from langchain_core.prompts import (
+    ChatPromptTemplate,
+    MessagesPlaceholder,
+    PromptTemplate,
+)
+
+# ==========================================
+# 1. CLASSIFIER PROMPT (Few-Shot Optimized)
+# ==========================================
+classifier_prompt = PromptTemplate.from_template(
+    """Anda adalah penganalisis soalan peperiksaan Sejarah SPM. 
+    Tugas anda adalah mengkategorikan soalan kepada "FAKTA" atau "KBAT" (Kemahiran Berfikir Aras Tinggi).
+
+    PANDUAN ASAS:
+    - FAKTA: Bertanya maklumat dari buku teks (tujuan sejarah asal, kronologi, tarikh, tokoh, peristiwa).
+    - KBAT: Bertanya ulasan, pendapat, rasional, atau menghubungkaitkan nilai/pengajaran kepada masyarakat atau negara secara umum.
+
+    BERIKUT ADALAH CONTOH RUJUKAN ANDA (FEW-SHOT EXAMPLES):
+
+    [CONTOH KATEGORI FAKTA]
+    Soalan: Pilihan raya di negara kita bermula dengan pilihan raya majlis perbandaran, pilihan raya negeri dan pilihan raya umum yang diadakan pada tahun 1955. Mengapakah pilihan raya diadakan?
+    Analisis: Walaupun bermula dengan petikan panjang, ia bertanya tujuan sejarah (mengapa pilihan raya diadakan pada waktu tersebut).
+    Jawapan: FAKTA
+
+    Soalan: Jelaskan perkembangan Pilihan Raya Majlis Perbandaran.
+    Analisis: Walaupun menggunakan perkataan "Jelaskan", ia diikuti dengan "perkembangan", bermaksud meminta kronologi peristiwa sejarah sebenar.
+    Jawapan: FAKTA
+
+    Soalan: Bagaimanakah Jawatankuasa Setia Kawan Malaysia (JSKM) dapat menjayakan pembentukan Malaysia?
+    Analisis: Bertanya tentang proses dan tindakan spesifik sebuah jawatankuasa sejarah.
+    Jawapan: FAKTA
+
+    [CONTOH KATEGORI KBAT]
+    Soalan: Pembentukan Malaysia penting bagi menjamin kestabilan dan kemakmuran negara. Jelaskan.
+    Analisis: "Jelaskan" di sini merujuk kepada konsep umum (kestabilan dan kemakmuran), bukan kronologi peristiwa.
+    Jawapan: KBAT
+
+    Soalan: Mengapakah undang-undang penting kepada sesebuah negara?
+    Analisis: Bertanya soalan universal/umum ("sesebuah negara"), bukan peristiwa spesifik Tanah Melayu.
+    Jawapan: KBAT
+
+    Soalan: Kemerdekaan negara amat unik kerana ia dicapai melalui meja rundingan tanpa melibatkan pertumpahan darah. Berdasarkan pernyataan tersebut, beri ulasan anda.
+    Analisis: Meminta "ulasan" berdasarkan pernyataan nilai sejarah.
+    Jawapan: KBAT
+
+    TUGAS ANDA:
+    Berdasarkan contoh-contoh di atas, kategorikan soalan berikut dan berikan analisis anda. 
+
+    PENTING: Anda MESTI mengeluarkan jawapan dalam format JSON yang sah dengan dua kunci berikut:
+    1. "analisis": (Penerangan ringkas mengapa soalan itu FAKTA atau KBAT)
+    2. "kategori": (HANYA perkataan "FAKTA" atau "KBAT")
+
+    Soalan: {query}
+    Output JSON:"""
+)
+
+# ==========================================
+# 2. QUERY REWRITER PROMPT (The Search Optimizer)
+# ==========================================
+# Purpose: Strip conversational fluff and isolate core historical entities for vector retrieval.
+query_prompt = PromptTemplate.from_template(
+    """Anda adalah pakar carian pangkalan data vektor untuk Sejarah Malaysia. 
+    Tugas anda adalah menulis semula soalan pengguna supaya menjadi kata kunci carian yang optimum.
+
+    GARIS PANDUAN:
+    1. Buang perkataan soal (seperti apakah, siapakah, jelaskan, huraikan).
+    2. Kekalkan HANYA entiti sejarah utama, nama tokoh, lokasi, atau nama peristiwa.
+    3. Tambah sinonim atau terma rasmi yang relevan jika perlu untuk meluaskan carian sejarah.
+
+    CONTOH 1: 
+    Soalan Asal: "Perjanjian Malaysia yang melibatkan British, Persekutuan Tanah Melayu, Sarawak, Sabah dan Singapura telah ditandatangani pada 9 Julai 1963. Jelaskan inti pati Perjanjian Malaysia." 
+    Carian Optimum: "Perjanjian Malaysia 9 Julai 1963"
 
-from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+    CONTOH 2:
+    Soalan Asal: "Nyatakan ciri-ciri negara bangsa Kesultanan Melayu Melaka."
+    Carian Optimum: "Negara bangsa Kesultanan Melayu Melaka"
 
-from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+    CONTOH 3:
+    Soalan Asal: "Apakah yang anda faham dengan konsep kedaulatan?"
+    Carian Optimum: "Konsep kedaulatan"
 
-general_prompt = ChatPromptTemplate.from_messages(
+    CONTOH 4:
+    Soalan Asal: "Bagaimanakah Jawatankuasa Setia Kawan Malaysia (JSKM) dapat menjayakan pembentukan Malaysia?"
+    Carian Optimum: "Jawatankuasa Setia Kawan Malaysia (JSKM)"
+
+    Tulis semula soalan berikut untuk carian pangkalan data. Berikan HANYA kata kunci carian tersebut.
+
+    Soalan Asal: {query}
+    Carian Optimum:"""
+)
+
+
+# ==========================================
+# 3. ANSWER PROMPT (The Strict Examiner)
+# ==========================================
+# Purpose: Answer the query using ONLY the retrieved context, formatted as a marking scheme.
+answer_prompt = ChatPromptTemplate.from_messages(
     [
         (
             "system",
-            """Anda adalah Guru Pakar Sejarah Malaysia yang sangat tegas. Tugas anda adalah menjawab soalan berdasarkan maklumat tepat daripada teks yang diberikan SAHAJA.
+            """Anda adalah Guru Pakar Sejarah Malaysia dan pemeriksa kertas peperiksaan yang sangat tegas. Tugas anda adalah menjawab soalan berdasarkan maklumat tepat daripada teks yang diberikan SAHAJA.
 
             GARIS PANDUAN PENTING:
             1. EKSTRAK SEMUA FAKTA (KOMPREHENSIF): Baca seluruh blok konteks yang diberikan secara terperinci. Ekstrak dan senaraikan SEMUA isi penting, peristiwa, tokoh, dan tarikh yang relevan dengan soalan. Jangan ringkaskan secara berlebihan atau tinggalkan sebarang maklumat yang berpotensi menjadi markah.
@@ -18,9 +106,50 @@ general_prompt = ChatPromptTemplate.from_messages(
             3. JAWAPAN TIADA DALAM KONTEKS: Jika konteks yang diberikan langsung tidak mengandungi maklumat yang relevan untuk menjawab soalan, anda MESTI membalas dengan ayat ini sahaja: "Maklumat ini tidak terdapat dalam konteks yang diberikan." Jangan sesekali cuba mereka jawapan.
             4. PETIKAN SUMBER (WAJIB): 
                - Setiap blok konteks bermula dengan label sumber dalam format [Source: X].
-               - Anda MESTI meletakkan petikan sumber yang TEPAT di hujung setiap isi poin anda.
-               - Contoh: Jika konteks bermula dengan "[Source: Buku_Teks_Sejarah_Tingkatan_5 Page 4]", tulis isi sebagai: "Asas berkerajaan sendiri. (Buku_Teks_Sejarah_Tingkatan_5 Page 4)"
+               - Anda MESTI meletakkan petikan sumber yang TEPAT di paling hujung isi poin anda sekali sahaja.
                - Dilarang meneka nombor muka surat jika tiada dalam label.
+            5. FORMAT SKEMA PEMARKAHAN: Susun jawapan anda dalam bentuk "bullet points" yang ringkas, padat, dan mudah dibaca (Bahasa Melayu). Elakkan perenggan yang meleret.
+
+            CONTOH FORMAT JAWAPAN (FEW-SHOT EXAMPLES):
+            Berikut adalah contoh anda menstrukturkan soalan dan jawapan. Terdapat dua jenis soalan, iaitu soalan tunggal atau soalan yang mempunyai beberapa bahagian (mengandaikan sumber konteks ialah [Source: Teks_Rujukan Page x]):
+
+            CONTOH 1:
+            Pilihan raya di negara kita bermula dengan pilihan raya majlis perbandaran, pilihan raya negeri dan pilihan raya umum yang diadakan pada tahun 1955. Mengapakah pilihan raya diadakan?
+            - Asas berkerajaan sendiri
+            - Memilih wakil
+            - Membentuk kerajaan
+            - Pengamalan sistem demokrasi
+            [Source: Teks_Rujukan Page x]
+
+            CONTOH 2:
+            Jelaskan perkembangan Pilihan Raya Majlis Perbandaran.
+            - Pilihan raya Perbandaran George Town, Pulau Pinang/diadakan pada tahun 1951
+            - Untuk memilih sembilan orang ahli Pesuruhjaya Perbandaran George Town
+            - Ahli dalam Majlis Perbandaran dilantik oleh Gabenor Negeri-negeri Selat
+            - Hari pendaftaran pemilih diadakan selama enam minggu/mulai 17 Mei 1951 hingga 30 Jun 1951
+            - Seramai 14 514 orang pengundi mendaftar untuk membuang undi
+            - Pilihan Raya Majlis Perbandaran Kuala Lumpur yang diadakan pada Februari 1952
+            - Merebut 12 kerusi Majlis Perbandaran Kuala Lumpur bagi kawasan Sentul/Bangsar/Imbi/Petaling Jaya
+            [Source: Teks_Rujukan Page x]
+
+            CONTOH 3: 
+            1. Tunku Abdul Rahman Putra al-Haj mengumumkan pembentukan Malaysia pada 27 Mei 1961.
+            (a) Namakan dua wilayah yang terlibat dalam cadangan pembentukan tersebut. [2 markah]
+            - Persekutuan Tanah Melayu
+            - Sarawak
+            - Sabah
+            - Singapura
+            - Brunei
+            [Source: Teks_Rujukan Page x]
+            (b) Bagaimanakah Jawatankuasa Setia Kawan Malaysia (JSKM) dapat menjayakan pembentukan Malaysia? [4 markah]
+            - Mengadakan mesyuarat sebanyak empat kali
+            - Jesselton/Kuching/ Kuala Lumpur/ Singapura
+            - Menyakinkan penduduk Sarawak Sabah tentang pembentukan Malaysia
+            - Menyediakan ruang perbincangan bagi mendapatkan kata sepakat terhadap pembentukan Malaysia
+            - Mengumpulkan pandangan tentang pembentukan Malaysia
+            - Menyebarkan maklumat berkaitan dengan Malaysia
+            - Memupuk aktiviti penghebahan/ mempercepat pembentukan Malaysia
+            [Source: Teks_Rujukan Page x]
             """,
         ),
         MessagesPlaceholder(variable_name="chat_history"),
@@ -31,7 +160,7 @@ general_prompt = ChatPromptTemplate.from_messages(
 
             Soalan: {query}
 
-            Berikan skema jawapan yang KOMPREHENSIF (merangkumi semua isi yang relevan) berdasarkan konteks di atas SAHAJA. Pastikan setiap poin disokong dengan petikan sumber [Source: X] yang betul.""",
+            Berikan skema jawapan yang KOMPREHENSIF (merangkumi semua isi yang relevan) berdasarkan konteks di atas SAHAJA. Pastikan letak petikan sumber [Source: X] yang betul di paling hujung isi poin anda sekali sahaja.""",
         ),
     ]
 )

app/schemas/rag.py

@@ -39,7 +39,7 @@ class AnswerRequest(BaseModel):
 class DebugInfo(BaseModel):
     """Debug information included in the answer response."""
 
-    top_doc_ids: list[str] = Field(
+    top_doc_ids: List[str] = Field(
         default_factory=list,
         description="IDs of the top chunks used as context.",
     )
@@ -57,7 +57,7 @@ class AnswerResponse(BaseModel):
         ...,
         description="Generated answer text based on retrieved context.",
     )
-    citations: list[str] = Field(
+    citations: List[str] = Field(
         default_factory=list,
         description="List of source citations used in the answer.",
     )

app/services/chat.py

@@ -2,7 +2,7 @@
 
 import asyncio
 import logging
-from typing import Any, AsyncGenerator
+from typing import Any, AsyncGenerator, List, Dict
 
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_core.runnables import Runnable
@@ -17,13 +17,13 @@ class ChatService:
 
     def __init__(self) -> None:
         """Initialize chat service with the configured LLM."""
-        self.model = settings.llm
-        logger.info(f"Initialized chat service with provider: {settings.ai_provider}")
+        self.model = settings.answer_llm
+        logger.info(f"Initialized chat service with provider: {self.model}")
 
     async def generate_answer(
         self,
         query: str,
-        context_blocks: list[dict[str, Any]],
+        context_blocks: List[Dict[str, Any]],
         prompt: ChatPromptTemplate,
     ) -> str:
         """Generate a RAG answer using retrieved context blocks.
@@ -49,9 +49,9 @@ class ChatService:
             logger.error(f"Error generating answer: {e}")
             return "Sorry, I was unable to generate an answer."
 
-    def _build_context_string(self, context_blocks: list[dict[str, Any]]) -> str:
+    def _build_context_string(self, context_blocks: List[Dict[str, Any]]) -> str:
         """Helper to build context string with citations."""
-        context_parts: list[str] = []
+        context_parts: List[str] = []
         for block in context_blocks:
             chunk_id = block.get("chunk_id", "unknown")
             source = block.get("source", "")

app/services/chunker.py

@@ -2,7 +2,7 @@
 
 import logging
 import re
-from typing import Any
+from typing import Any, List, Dict
 
 from langchain_text_splitters import (
     MarkdownHeaderTextSplitter,
@@ -70,7 +70,7 @@ class TextChunker:
             return f"Page {match.group(1)}"
         return None
 
-    def _extract_page_sections(self, text: str) -> list[tuple[str | None, str]]:
+    def _extract_page_sections(self, text: str) -> List[tuple[str | None, str]]:
         """Split text into sections by page markers and extract page numbers.
 
         Args:
@@ -84,7 +84,7 @@ class TextChunker:
         if not page_markers:
             return [(None, text)]
 
-        sections: list[tuple[str | None, str]] = []
+        sections: List[tuple[str | None, str]] = []
         for i, match in enumerate(page_markers):
             page_num = f"Page {match.group(1)}"
             start = match.end()
@@ -101,7 +101,7 @@ class TextChunker:
         """Remove page source markers from text content."""
         return PAGE_SOURCE_PATTERN.sub("", text).strip()
 
-    def _build_page_index(self, text: str) -> list[tuple[int, str]]:
+    def _build_page_index(self, text: str) -> List[tuple[int, str]]:
         """Build an index of (position, page_number) from page markers.
 
         Args:
@@ -120,7 +120,7 @@ class TextChunker:
         self,
         content: str,
         original_text: str,
-        page_index: list[tuple[int, str]],
+        page_index: List[tuple[int, str]],
         fallback: str,
     ) -> str:
         """Find the page number for a chunk by locating its position in original text.
@@ -161,7 +161,7 @@ class TextChunker:
 
     def chunk_text(
         self, text: str, source: str, base_chunk_id: str
-    ) -> list[dict[str, Any]]:
+    ) -> List[Dict[str, Any]]:
         """Split text into overlapping chunks with page number extraction.
 
         Args:
@@ -187,7 +187,7 @@ class TextChunker:
         # Second pass: Split large sections by character limit
         split_docs = self.text_splitter.split_documents(header_docs)
 
-        chunks: list[dict[str, Any]] = []
+        chunks: List[Dict[str, Any]] = []
         chunk_num = 1
         last_page = source  # Track last seen page for continuity
 
@@ -246,7 +246,7 @@ class TextChunker:
 
         return text
 
-    def _attach_header_context(self, metadata: dict[str, Any], content: str) -> str:
+    def _attach_header_context(self, metadata: Dict[str, Any], content: str) -> str:
         """Prefix chunk content with header context from markdown metadata.
 
         Args:
@@ -256,7 +256,7 @@ class TextChunker:
         Returns:
             Content prefixed with header context.
         """
-        header_lines: list[str] = []
+        header_lines: List[str] = []
         for level in range(1, 5):
             key = f"Header {level}"
             if metadata.get(key):
@@ -268,7 +268,7 @@ class TextChunker:
 
         return "\n\n".join(["\n".join(header_lines), content])
 
-    def chunk_documents(self, documents: list[dict[str, str]]) -> list[dict[str, Any]]:
+    def chunk_documents(self, documents: List[Dict[str, str]]) -> List[Dict[str, Any]]:
         """Chunk multiple documents.
 
         Args:
@@ -277,7 +277,7 @@ class TextChunker:
         Returns:
             List of all chunks from all documents.
         """
-        all_chunks: list[dict[str, Any]] = []
+        all_chunks: List[Dict[str, Any]] = []
 
         for doc in documents:
             text = doc.get("text", "")

app/services/embedding.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import logging
+from typing import List
 
 from app.core.config import settings
 
@@ -11,20 +12,18 @@ logger = logging.getLogger(__name__)
 class EmbeddingService:
     """Service for generating text embeddings using SentenceTransformers."""
 
-    def __init__(
-        self,
-        model_name: str = "",
-        device: str = "cuda",
-        max_length: int = 8192,
-    ) -> None:
+    def __init__(self):
         import torch
         from sentence_transformers import SentenceTransformer
 
-        _model_name = model_name or settings.model_name
+        _model_name = settings.embedding_model_name
+        _device = settings.device
 
         # Auto-detect device
-        if device == "cuda" and not torch.cuda.is_available():
-            device = "cpu"
+        if torch.cuda.is_available():
+            _device = "cuda"
+        else:
+            _device = "cpu"
             logger.warning("CUDA not available, falling back to CPU")
 
         logger.info(f"Loading embedding model: {_model_name}")
@@ -33,11 +32,9 @@ class EmbeddingService:
         self.model = SentenceTransformer(
             _model_name,
             trust_remote_code=True,
-            device=device,
-            model_kwargs={"torch_dtype": torch.float16} if device == "cuda" else {},
+            device=_device,
+            model_kwargs={"torch_dtype": torch.float16} if _device == "cuda" else {},
         )
-        self.model.max_seq_length = max_length
-        self.device = device
 
         self.embedding_dim = self.model.get_sentence_embedding_dimension()
         logger.info(
@@ -45,8 +42,8 @@ class EmbeddingService:
         )
 
     async def embed_texts(
-        self, texts: list[str], batch_size: int = 8
-    ) -> list[list[float]]:
+        self, texts: List[str], batch_size: int = 8
+    ) -> List[List[float]]:
         """Generate embeddings for a list of document texts.
 
         Args:
@@ -78,7 +75,7 @@ class EmbeddingService:
             logger.error(f"Failed to generate embeddings: {e}")
             raise
 
-    async def embed_query(self, query: str) -> list[float]:
+    async def embed_query(self, query: str) -> List[float]:
         """Generate embedding for a single query.
 
         Args:

app/services/ocr.py

@@ -0,0 +1,66 @@
+"""OCR service to extract text from image."""
+
+from PIL import Image
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class OCRService:
+    """OCR service to extract text from image."""
+
+    def __init__(self):
+        import torch
+        from transformers import AutoProcessor, AutoModelForImageTextToText
+
+        _model_name = "zai-org/GLM-OCR"
+        _device = "cuda"
+
+        # Auto-detect device
+        if torch.cuda.is_available():
+            _device = "cuda"
+        else:
+            _device = "cpu"
+            logger.warning("CUDA not available, falling back to CPU")
+
+        logger.info(f"Loading OCR model: {_model_name}")
+
+        self.processor = AutoProcessor.from_pretrained(_model_name)
+        self.model = AutoModelForImageTextToText.from_pretrained(
+            pretrained_model_name_or_path=_model_name,
+            torch_dtype="auto",
+            device_map="auto",
+        )
+
+    async def extract_text(self, pil_image: Image.Image) -> str:
+        """Extract text from image."""
+
+        img = pil_image.convert("RGB")
+
+        message = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image", "image": img},
+                    {"type": "text", "text": "Text Recognition:"},
+                ],
+            }
+        ]
+
+        inputs = self.processor.apply_chat_template(
+            message,
+            tokenize=True,
+            add_generation_prompt=True,
+            return_dict=True,
+            return_tensors="pt",
+        ).to(self.model.device)
+
+        inputs.pop("token_type_ids", None)
+        generated_ids = self.model.generate(**inputs, max_new_tokens=8192)
+        output_text = self.processor.decode(
+            generated_ids[0][inputs["input_ids"].shape[1] :], skip_special_tokens=True
+        )
+        return output_text
+
+
+ocr_service = OCRService()

app/services/query_preprocessor.py

@@ -0,0 +1,49 @@
+"""Query preprocessing: classification and query rewriting."""
+
+import asyncio
+import logging
+
+from langchain_core.output_parsers import StrOutputParser, JsonOutputParser
+
+from app.core.config import settings
+from app.core.prompt_templates import classifier_prompt, query_prompt
+
+logger = logging.getLogger(__name__)
+
+
+class QueryPreprocessor:
+    """Handles query classification and rewriting as two separate chains."""
+
+    def __init__(self) -> None:
+        self.llm = settings.preprocess_llm
+        # Chain 1: Classify → returns "FAKTA" or "KBAT"
+        self.classifier_chain = classifier_prompt | self.llm | JsonOutputParser()
+        # Chain 2: Rewrite → returns optimized search query
+        self.rewriter_chain = query_prompt | self.llm | StrOutputParser()
+
+    async def classify(self, query: str) -> tuple[str, bool]:
+        """Classify a query as FAKTA or KBAT.
+
+        Returns:
+            A tuple of (classification, needs_rag).
+            - classification: "FAKTA" or "KBAT"
+            - needs_rag: True if FAKTA (proceed with RAG), False if KBAT (stop).
+        """
+        result = await asyncio.to_thread(self.classifier_chain.invoke, {"query": query})
+        classification = result["kategori"].strip().upper()
+        needs_rag = classification == "FAKTA"
+        logger.info(f"Query classified as: {classification} | needs_rag: {needs_rag} | Analysis: {result['analisis']}")
+        return classification, needs_rag
+
+    async def rewrite(self, query: str) -> str:
+        """Rewrite a query to optimize for vector search.
+
+        Returns:
+            Optimized search query string.
+        """
+        result = await asyncio.to_thread(self.rewriter_chain.invoke, {"query": query})
+        logger.info(f"Rewritten query: {result.strip()}")
+        return result.strip()
+
+
+query_preprocessor = QueryPreprocessor()

app/services/rag.py

@@ -5,14 +5,19 @@ Orchestrates the complete RAG pipeline: chunk → embed → search → generate.
 
 import asyncio
 import logging
+import re
 import time
-from typing import Any, AsyncGenerator
+from typing import Any, AsyncGenerator, Optional, List, Dict
+from fastapi import UploadFile
+from PIL import Image
 
 from app.core.database import db
-from app.core.prompt_templates import general_prompt
+from app.core.prompt_templates import answer_prompt
 from app.services.chat import chat_service
 from app.services.chunker import chunker
 from app.services.embedding import embedding_service
+from app.services.ocr import ocr_service
+from app.services.query_preprocessor import query_preprocessor
 
 logger = logging.getLogger(__name__)
 
@@ -84,7 +89,12 @@ class RAGService:
             logger.error(f"Seeding failed: {e}")
             raise
 
-    async def answer_question(self, query: str, top_k: int = 6) -> dict[str, Any]:
+    async def answer_question(
+        self,
+        query: Optional[str] = "",
+        top_k: int = 6,
+        image: Optional[UploadFile] = None,
+    ) -> Dict[str, Any]:
         """Process a query through the complete RAG pipeline.
 
         Args:
@@ -97,8 +107,31 @@ class RAGService:
         start_time = time.time()
 
         try:
+            if image:
+                pil_image = Image.open(image.file)
+                content = await ocr_service.extract_text(pil_image)
+                query = content
+                logger.info(f"Extracted text from image: {content}")
+
+            query_classification, needs_rag = await query_preprocessor.classify(query)
+
+            if not needs_rag:
+                return {
+                    "text": (
+                        "Soalan ini adalah KBAT. Saya TIDAK dapat menjawab soalan KBAT. "
+                        "Sila tanya soalan FAKTA."
+                    ),
+                    "citations": [],
+                    "debug": {
+                        "top_doc_ids": [],
+                        "latency_ms": int((time.time() - start_time) * 1000),
+                    },
+                }
+
+            rewritten_query = await query_preprocessor.rewrite(query)
+
             # Step 1: Generate query embedding
-            query_embedding = await self.embedding_service.embed_query(query)
+            query_embedding = await self.embedding_service.embed_query(rewritten_query)
 
             # Step 2: Vector similarity search
             search_results = await self.db.vector_search(query_embedding, top_k)
@@ -118,11 +151,11 @@ class RAGService:
 
             # Step 3: Deduplicate and prepare context
             context_blocks = self._prepare_context(search_results)
-            logger.debug(f"Context blocks: {context_blocks}")
+            logger.info(f"Context blocks: {context_blocks}")
 
             # Step 4: Generate answer
             answer_text = await self.chat_service.generate_answer(
-                query, context_blocks, general_prompt
+                query, context_blocks, answer_prompt
             )
             logger.debug(f"Answer text: {answer_text}")
 
@@ -157,8 +190,8 @@ class RAGService:
             }
 
     def _prepare_context(
-        self, search_results: list[dict[str, Any]]
-    ) -> list[dict[str, Any]]:
+        self, search_results: List[Dict[str, Any]]
+    ) -> List[Dict[str, Any]]:
         """Prepare context blocks from search results with deduplication.
 
         Args:
@@ -168,7 +201,7 @@ class RAGService:
             Processed, deduplicated context blocks.
         """
         seen_prefixes: set[str] = set()
-        context_blocks: list[dict[str, Any]] = []
+        context_blocks: List[Dict[str, Any]] = []
 
         for result in search_results:
             chunk_id = result.get("chunk_id", "")
@@ -190,18 +223,33 @@ class RAGService:
     def _extract_citations(
         self,
         answer_text: str,
-        context_blocks: list[dict[str, Any]],
-    ) -> list[str]:
-        """Return citations from the context blocks that were used.
+        context_blocks: List[Dict[str, Any]],
+    ) -> List[str]:
+        """Return only citations that the LLM actually referenced in its answer.
+
+        Parses [Source: X] tags from the answer text, then matches them
+        against the available context blocks. Only citations that appear
+        in the answer are returned.
 
         Args:
-            answer_text: Generated answer text (kept for future use).
-            context_blocks: Context blocks that were provided.
+            answer_text: Generated answer text containing [Source: ...] tags.
+            context_blocks: Context blocks that were provided to the LLM.
 
         Returns:
-            List of formatted citations combining base chunk_id and source.
+            List of formatted citations that were actually referenced.
         """
-        citations: list[str] = []
+        # Step 1: Extract all [Source: ...] references from the answer text
+        cited_sources_raw = re.findall(r"\[Source:\s*(.+?)\]", answer_text)
+
+        if not cited_sources_raw:
+            logger.debug("No [Source: ...] tags found in answer text")
+            return []
+
+        # Normalize cited sources for fuzzy matching
+        cited_sources_normalized = {s.strip().lower() for s in cited_sources_raw}
+
+        # Step 2: Build available citations from context blocks
+        citations = []
         for block in context_blocks:
             chunk_id = block.get("chunk_id", "")
             source = block.get("source", "")
@@ -209,10 +257,18 @@ class RAGService:
             # Extract base chunk_id (remove the #N suffix)
             base_id = chunk_id.split("#")[0] if "#" in chunk_id else chunk_id
 
-            # Combine base_id and source
+            # This matches the format used in chat.py _build_context_string()
             citation = f"{base_id} {source}".strip()
 
-            if citation and citation not in citations:
+            if not citation or citation in citations:
+                continue
+
+            # Check if this citation was actually referenced in the answer
+            citation_lower = citation.lower()
+            if any(
+                citation_lower in cs or cs in citation_lower
+                for cs in cited_sources_normalized
+            ):
                 citations.append(citation)
 
         return citations

main.py

@@ -5,9 +5,10 @@ import logging
 import time
 from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
+from typing import Optional
 
 import uvicorn
-from fastapi import FastAPI, HTTPException, UploadFile
+from fastapi import FastAPI, File, Form, HTTPException, UploadFile
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
 
@@ -156,7 +157,13 @@ async def seed_documents(file: UploadFile) -> SeedResponse:
         500: {"model": ErrorResponse, "description": "Processing failed"},
     },
 )
-async def answer_question(request: AnswerRequest) -> AnswerResponse:
+async def answer_question(
+    query: Optional[str] = Form(default="", description="The question to answer."),
+    top_k: int = Form(default=6, description="Number of top-k chunks to retrieve."),
+    image: Optional[UploadFile] = File(
+        default=None, description="Optional image upload."
+    ),
+) -> AnswerResponse:
     """Answer a question using Retrieval-Augmented Generation.
 
     Pipeline:
@@ -166,10 +173,10 @@ async def answer_question(request: AnswerRequest) -> AnswerResponse:
     4. Return answer with citations and debug info
     """
     try:
-        logger.info(f"Processing query: '{request.query[:100]}...'")
+        logger.info(f"Processing query: '{query[:100] if query else 'No query'}...'")
 
         result = await rag_service.answer_question(
-            query=request.query, top_k=request.top_k
+            query=query, top_k=top_k, image=image
         )
 
         response = AnswerResponse(

pyproject.toml

@@ -1,66 +1,32 @@
 [project]
-name = "sejarah-rag"
+name = "ocr"
 version = "0.1.0"
-description = "RAG backend for Malaysian History (Sejarah) education"
+description = "Add your description here"
 readme = "README.md"
-requires-python = ">=3.12.11"
+requires-python = ">=3.13"
 dependencies = [
-    "fastapi>=0.115.0",
-    "uvicorn[standard]>=0.34.0",
-    "pydantic>=2.10.0",
-    "pydantic-settings>=2.7.0",
-    "langchain-core>=1.2.7",
-    "langchain-groq>=1.1.1",
-    "langchain-openai>=1.1.7",
-    "langchain-text-splitters>=1.1.0",
-    "llama-cpp-python>=0.3.4",
-    "sentence-transformers>=3.4.0,<4",
-    "supabase>=2.27.2",
-    "transformers>=4.44.0,<4.47",
+    "transformers",
+    "torch",
+    "torchvision",
+    "accelerate>=1.12.0",
+    "pydantic-settings>=2.13.1",
+    "uvicorn>=0.41.0",
+    "fastapi>=0.135.0",
+    "supabase>=2.28.0",
+    "langchain-core>=1.2.16",
+    "langchain-groq>=1.1.2",
+    "langchain-text-splitters>=1.1.1",
+    "sentence-transformers>=5.2.3",
     "python-multipart>=0.0.22",
-    "torch>=2.5.1",
-    "torchvision>=0.20.1",
-    "torchaudio>=2.5.1",
+    "langchain-openai>=1.1.10",
 ]
 
-[dependency-groups]
-dev = [
-    "pytest>=8.0.0",
-    "pytest-asyncio>=0.24.0",
-    "httpx>=0.27.0",
-    "ruff>=0.9.0",
-    "ipykernel>=7.1.0",
-]
-
-[tool.uv]
-required-environments = ["sys_platform == 'win32'", "sys_platform == 'linux'"]
-
-[[tool.uv.index]]
-name = "llama-cpp-python-cuda"
-url = "https://abetlen.github.io/llama-cpp-python/whl/cu121"
-explicit = true
-
 [[tool.uv.index]]
-name = "pytorch"
-url = "https://download.pytorch.org/whl/cu121"
+name = "cu126"
+url = "https://download.pytorch.org/whl/cu126"
 explicit = true
 
 [tool.uv.sources]
-llama-cpp-python = { index = "llama-cpp-python-cuda" }
-torch = { index = "pytorch" }
-torchvision = { index = "pytorch" }
-torchaudio = { index = "pytorch" }
-
-[tool.ruff]
-line-length = 100
-target-version = "py312"
-
-[tool.ruff.lint]
-select = ["E", "F", "I", "N", "W", "UP", "B", "SIM"]
-
-[tool.ruff.lint.isort]
-known-first-party = ["app"]
-
-[tool.pytest.ini_options]
-testpaths = ["tests"]
-asyncio_mode = "auto"
+transformers = { git = "https://github.com/huggingface/transformers.git" }
+torch = { index = "cu126" }
+torchvision = { index = "cu126" }