import gradio as gr
from huggingface_hub import InferenceClient

client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    # Initialize messages with system instructions
    messages = [{"role": "system", "content": system_message}]

    # Add historical conversation
    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    # Add the new user message
    messages.append({"role": "user", "content": message})

    response = ""

    # Get the model's response
    for response_chunk in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = response_chunk.choices[0].delta.content
        response += token

    response = response.strip()

    # Enhanced context-specific relevance check
    if is_relevant_to_constitution(response):
        return response

    # If response does not meet relevance criteria
    return "Sorry, I can only provide information related to the Constitution of India. Please ask a question related to the Constitution."

def is_relevant_to_constitution(response):
    # Keywords to check in the response
    relevant_keywords = [
        "constitution", "article", "law", "legal", "rights", "act", "judiciary", 
        "legislature", "executive", "fundamental", "amendment", "provision", 
        "policy", "directive", "supreme court", "high court", "legislation", 
        "government", "election", "parliament", "state", "central", "reform", 
        "citizen", "equality", "democracy", "directive principles", "fundamental duties", 
        "preamble", "enforcement", "federalism", "separation of powers", "justice", 
        "republic", "state legislature", "union territory", "bill", "ordinance", 
        "convention", "charter", "treaty", "declaration", "proclamation", "amendments", 
        "compensation", "grievance", "judicial review", "secularism", "socialism", 
        "pluralism", "sovereignty", "autonomy", "independence", "integrity", "caste", 
        "reservation", "minorities", "discrimination", "fundamental rights", 
        "emergency", "state emergency", "national emergency", "local bodies", 
        "tribunal", "ombudsman", "civil rights", "criminal justice", "human rights"
    ]

    # Check if response contains relevant keywords
    return any(keyword in response.lower() for keyword in relevant_keywords)

# Create the Gradio chat interface
demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value="You are a knowledgeable assistant specializing in the Constitution of India. Answer only questions related to the Constitution. If the answer does not contain relevant constitutional keywords, inform the user accordingly.", label="System message", visible=False),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
    ],
)

if __name__ == "__main__":
    demo.launch()