FInancial_Helper.py

import os
from pathlib import Path
import logging
import streamlit as st
from io import BytesIO
from dotenv import load_dotenv
import google.generativeai as genai
import chromadb
from chromadb import Documents, EmbeddingFunction, Embeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from PyPDF2 import PdfReader
import re
from deep_translator import GoogleTranslator# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
)
logger = logging.getLogger(__name__)

# Load env and configure Gemini
load_dotenv()
API_KEY = os.getenv('GEMINI_API_KEY')
if not API_KEY:
    st.error("Missing GEMINI_API_KEY in environment variables")
    st.stop()
genai.configure(api_key=API_KEY)

# Constants
db_path = Path("./db")
collection_name = "sme_db"


# Embedding function with batching
class GeminiEmbeddingFunction(EmbeddingFunction):
    def __init__(self, model: str = 'models/embedding-001', batch_size: int = 16):
        self.model = model
        self.batch_size = batch_size

    def __call__(self, input: Documents) -> Embeddings:
        texts = list(input)
        embeddings = []
        for i in range(0, len(texts), self.batch_size):
            batch = texts[i:i + self.batch_size]
            resp = genai.embed_content(
                model=self.model,
                content=batch,
                task_type="retrieval_document",
                title="RAG Batch"
            )
            embeddings.extend(resp['embedding'])
        return embeddings


# Utilities
def extract_text_from_pdf(file_stream: BytesIO) -> str:
    reader = PdfReader(file_stream)
    text = []
    # skip first page if desired
    for page in reader.pages[1:]:
        page_text = page.extract_text() or ""
        text.append(page_text)
    return "\n".join(text)


# Keep Latin and Arabic chars
def clean_text(text: str) -> str:
    text = re.sub(r"(\w+)-\n(\w+)", r"\1\2", text)
    lines = text.splitlines()
    allowed = re.compile(r"[0-9A-Za-z\u0600-\u06FF\.,\$%\-\s]")
    cleaned = []
    for line in lines:
        raw = line.strip()
        if len(raw) < 20 or raw.isdigit():
            continue
        if re.match(r"^Page \d+ of", raw):
            continue
        filtered = ''.join(ch for ch in raw if allowed.match(ch))
        cleaned.append(filtered)
    return "\n\n".join(cleaned)


def translate_query_to_doc_lang(query: str, target_lang='en') -> str:
    return GoogleTranslator(source='auto', target=target_lang).translate(query)

def split_text(text: str, chunk_size: int = 1000, overlap: int = 100) -> list:
    splitter = RecursiveCharacterTextSplitter(
        separators=["\\n\\n", "\\n", "。", ".", "!", "؟", " "],
        chunk_size=1000,
        chunk_overlap=100,
        add_start_index=True
    )
    docs = splitter.create_documents([text])
    return [{"content": d.page_content, "start": d.metadata["start_index"]} for d in docs]


def get_collection() -> chromadb.api.models.Collection:
    client = chromadb.PersistentClient(path=str(db_path))
    return client.get_or_create_collection(
        name=collection_name,
        embedding_function=GeminiEmbeddingFunction()
    )


def check_hallucination(answer: str, context: list) -> bool:
    for sent in answer.split('.'):
        for w in [w.strip(' ,') for w in sent.split()]:
            if w.isdigit():
                continue
            if w.isalpha() and all(w not in chunk.lower() for chunk in context):
                return False
    return True


@st.cache_resource
def ingest(pdf_stream: BytesIO):
    raw = extract_text_from_pdf(pdf_stream)
    cleaned = clean_text(raw)
    chunks = split_text(cleaned)
    col = get_collection()
    if col.count() == 0:
        docs, metas, ids = [], [], []
        for i, c in enumerate(chunks):
            docs.append(c['content'])
            metas.append({"source": "uploaded_pdf", "chunk_start": c['start']})
            ids.append(str(i))
        col.add(documents=docs, ids=ids, metadatas=metas)
    return col


# Streamlit UI
st.title("RAG Financial Assistant")

uploaded_file = st.file_uploader("Upload a PDF file (it can be in Arabic or English)", type=["pdf"])

lang = st.selectbox("Select language you want the answer to be in:", ["English", "Arabic"])

if uploaded_file:
    query = st.text_input("Enter your question:")
    if st.button("Ask") and query:
        pdf_bytes = uploaded_file.read()
        pdf_stream = BytesIO(pdf_bytes)
        col = ingest(pdf_stream)
        query =translate_query_to_doc_lang(query)
        results = col.query(query_texts=[query], n_results=5)
        passages = results['documents'][0]
        if not check_hallucination(query, passages):
            st.warning("⚠️ Potential hallucination detected. Please refine your query.")
        context = "\n\n".join(f"Doc {i + 1}: {p}" for i, p in enumerate(passages))

        if lang == "Arabic":
            question_prefix = "سؤال:"
            answer_suffix = "الإجابة معتمدًا على السياق فقط."
        else:
            question_prefix = "Question:"
            answer_suffix = "Answer using only context."

        prompt = (
            f"You are a financial assistant tasked with answering user queries using reliable context.\n\n"
            f"Context:\n{context}\n\n"
            f"{question_prefix}\n{query}\n\n"
            "Instructions:\n"
            f"- Context might be in Arabic or English translate accordingly. answer in {lang} using only the provided context.\n"
            "- Use only the information in the context to answer.\n"
            "- Be precise and avoid speculation.\n"
            "- If information is missing, state that clearly.\n"
            "- Format numeric data with units.\n"
            f"{answer_suffix}\n\n"
            "Answer:"
        )
        with st.spinner("Generating answer..."):
            model = genai.GenerativeModel(model_name="gemini-1.5-flash")
            resp = model.generate_content(prompt)
        st.markdown("**Answer:**")
        st.write(resp.text)
else:
    st.info("Please upload a PDF to get started.")