app.py

import gradio as gr
from llama_cpp import Llama
import fitz  # PyMuPDF
import os

# ── Model loading ──────────────────────────────────────────────────────────────
MODEL_REPO = "newtechdevng/i_am_a_lawyer"
MODEL_FILE = "llama-3.2-1b-instruct.Q4_K_M.gguf"
SYSTEM_PROMPT = (
    "You are Ambuj, an expert AI assistant specialised in Indian law. "
    "You provide accurate, well-structured legal information based on Indian statutes, "
    "case law, and legal procedures. Always clarify that your responses are for "
    "informational purposes only and not a substitute for professional legal advice."
)

print("Loading model …")
llm = Llama.from_pretrained(
    repo_id=MODEL_REPO,
    filename=MODEL_FILE,
    n_ctx=4096,
    n_threads=os.cpu_count() or 4,
    verbose=False,
)
print("Model ready ✓")


# ── Helpers ────────────────────────────────────────────────────────────────────
def build_messages(history, system=SYSTEM_PROMPT):
    """
    history is a list of {"role": ..., "content": ...} dicts (Gradio 6 messages format).
    """
    msgs = [{"role": "system", "content": system}]
    for msg in history:
        msgs.append({"role": msg["role"], "content": msg["content"]})
    return msgs


def generate(messages, max_tokens=512, temperature=0.7):
    response = llm.create_chat_completion(
        messages=messages,
        max_tokens=max_tokens,
        temperature=temperature,
        stream=True,
    )
    partial = ""
    for chunk in response:
        delta = chunk["choices"][0]["delta"].get("content", "")
        partial += delta
        yield partial


def extract_text_from_file(file_path):
    if file_path is None:
        return ""
    if file_path.endswith(".pdf"):
        doc  = fitz.open(file_path)
        text = "\n".join(page.get_text() for page in doc)
        doc.close()
    else:
        with open(file_path, "r", errors="ignore") as f:
            text = f.read()
    return text.strip()


# ── Chat tab ───────────────────────────────────────────────────────────────────
def chat_respond(user_input, history):
    """
    history: list of {"role": str, "content": str} dicts (Gradio 6 messages format).
    """
    if not user_input.strip():
        yield history, ""
        return

    # Append user message
    history = history + [{"role": "user", "content": user_input}]
    # Build llm messages (includes system prompt)
    messages = build_messages(history)

    # Stream assistant reply
    partial = ""
    history = history + [{"role": "assistant", "content": ""}]
    for partial in generate(messages):
        history[-1]["content"] = partial
        yield history, ""


def clear_chat():
    return [], ""


# ── Document Analyzer tab ──────────────────────────────────────────────────────
ANALYSIS_SYSTEM = (
    SYSTEM_PROMPT + "\n\n"
    "When given a legal document, analyse it thoroughly and provide:\n"
    "1. Document type & summary\n"
    "2. Key parties involved\n"
    "3. Core legal provisions / clauses\n"
    "4. Potential legal issues or risks\n"
    "5. Relevant Indian laws / acts that apply\n"
    "6. Recommended next steps\n"
    "Keep the analysis structured and professional."
)


def analyze_document(file_path, extra_question):
    text = extract_text_from_file(file_path)
    if not text:
        yield "⚠️ Could not extract text. Please upload a valid PDF or .txt file."
        return

    truncated = text[:6000]
    question_part = f"\n\nAdditional question: {extra_question}" if extra_question.strip() else ""
    prompt = f"Analyse the following legal document:\n\n---\n{truncated}\n---{question_part}"

    messages = [
        {"role": "system", "content": ANALYSIS_SYSTEM},
        {"role": "user",   "content": prompt},
    ]
    for partial in generate(messages, max_tokens=768, temperature=0.4):
        yield partial


# ── UI ─────────────────────────────────────────────────────────────────────────
css = """
#title  { text-align: center; margin-bottom: 0.5rem; }
#sub    { text-align: center; color: #888; margin-bottom: 1.5rem; font-size: 0.9rem; }
.panel  { border-radius: 12px; }
footer  { display: none !important; }
"""

with gr.Blocks() as demo:

    gr.Markdown("# ⚖️ Indian Legal AI Assistant", elem_id="title")
    gr.Markdown(
        "Powered by **Ambuj-Tripathi-Indian-Legal-Llama** · For informational purposes only · Not legal advice",
        elem_id="sub",
    )

    with gr.Tabs():

        # ── Tab 1: Chat ────────────────────────────────────────────────────────
        with gr.Tab("💬 Chat"):
            chatbot = gr.Chatbot(
                label="Legal Chat",
                height=460,
                elem_classes="panel",
            )
            with gr.Row():
                chat_input = gr.Textbox(
                    placeholder="Ask about IPC, CPC, contracts, property law …",
                    show_label=False,
                    scale=5,
                )
                send_btn  = gr.Button("Send",  variant="primary", scale=1)
                clear_btn = gr.Button("Clear", variant="secondary", scale=1)

            gr.Examples(
                examples=[
                    ["What are my rights if I'm arrested without a warrant in India?"],
                    ["Explain Section 498A of the Indian Penal Code."],
                    ["What is the process for filing a consumer complaint in India?"],
                    ["How does anticipatory bail work under the CrPC?"],
                ],
                inputs=chat_input,
                label="Quick questions",
            )

            send_btn.click(
                chat_respond,
                inputs=[chat_input, chatbot],
                outputs=[chatbot, chat_input],
            )
            chat_input.submit(
                chat_respond,
                inputs=[chat_input, chatbot],
                outputs=[chatbot, chat_input],
            )
            clear_btn.click(clear_chat, outputs=[chatbot, chat_input])

        # ── Tab 2: Document Analyzer ───────────────────────────────────────────
        with gr.Tab("📄 Document Analyzer"):
            gr.Markdown(
                "Upload a legal document (PDF or TXT) and get an AI-powered analysis "
                "under Indian law."
            )
            with gr.Row():
                with gr.Column(scale=1):
                    doc_upload = gr.File(
                        label="Upload Document",
                        file_types=[".pdf", ".txt"],
                        type="filepath",
                    )
                    extra_q = gr.Textbox(
                        label="Additional question (optional)",
                        placeholder="e.g. Is this contract enforceable under Indian law?",
                        lines=2,
                    )
                    analyze_btn = gr.Button("Analyse Document", variant="primary")

                with gr.Column(scale=2):
                    analysis_out = gr.Textbox(
                        label="Analysis",
                        lines=22,
                        elem_classes="panel",
                    )

            analyze_btn.click(
                analyze_document,
                inputs=[doc_upload, extra_q],
                outputs=analysis_out,
            )

    gr.Markdown(
        "---\n"
        "⚠️ **Disclaimer:** This tool is for informational purposes only. "
        "Always consult a qualified advocate for matters requiring professional legal guidance."
    )

demo.queue(max_size=5).launch(
    theme=gr.themes.Soft(primary_hue="orange"),
    css=css,
)