app.py

import gradio as gr
import os
import requests
import json

# --- Configuration ---
# Model to use via the Inference API
MODEL_NAME = "CraneAILabs/swahili-gemma-1b-litert" 
API_URL = f"https://api-inference.huggingface.co/models/{MODEL_NAME}"

# --- Transcript Data (for RAG) ---
NURSE_TOTO_TRANSCRIPT = """
# A Nurse Toto - Episode 1: Mzee wa Kutahirii (Kiswahili Transcript)
**Series:** A Nurse Toto
**Episode:** 1 - Mzee wa Kutahirii
**Creator:** Eddie Butita
**Language:** Kiswahili (with some English/Sheng)

---

## Scene 1: The Date
**Characters:** Casypool (Mzee), Girlfriend

**Casypool:** My dear, mimi sio kama wale wanaume wengine wenye wanachukua mademu, wanawause, wanawadump. I'm an honorable man. That's why President Ruto ananipenda.

**Girlfriend:** Baby, you're lying. Yaani uko na connections kama hizo na uko tu hapa hivi you don't even brag. Aky you're so humble.

**Casypool:** Well, humility is my second name. I'm a man of my word. Na nikisema by the way, I will wait until marriage. Nita wait until marriage.

**Girlfriend:** Aky baby, you are so sweet and genuine. That's what I love about you. Anyway, umesema unarudi kutoka US lini?

**Casypool:** Well, I'll be back in two months time. But don't worry. I'll be sending you money to have fun with your friends. Pesa kidogo tu kama 100,000, 200,000, 300. I hope hiyo itakuwa ni kidogo itakutosha.

**Girlfriend:** Hiyo sio kidogo babe. Aky babe you are so sweet and young at heart. Yaani wewe ni mzee na roho ya kijana. Daudi Kabaka na roho ya Butros. I knew it wewe si kama hao baba wengine.

**Casypool:** Come on, don't call me mbaba. I'm a promoted youth with so many years experience. So, by the way, unafanya wapi kazi?

**Casypool:** I'm right now I'm into real estate. But nafanya kazi kama nurse kwa private hospital ya VIPs pekee yake. Lakini nataka kwenda Canada very soon ndio niende kufanya kazi huko. Nataka uanze kuwa na kiburi. Nataka uanze kutukanatukana boss. Nataka uanze kumea pembe kwa sababu nataka kubadilisha life yako. When I'm back in two months, we settle down. No rushing things. I will wait. Not touching you, no nothing until marriage.

**Girlfriend:** Aky babe, you are so sweet. Okay. Let me leave you some cash. Uende upige sherehe. Utoe stress leo because I'm flying out of the country tonight. And I know you'll miss me the way I'll miss you.

**Girlfriend:** I will miss you.

---

## Intro
**Narrator:** Anaitwa Wambo_ alipiga C yake smart pale high school kimchezo akagraduate kikomrade with the help of Mkenya na akapata job kama nurse na maguide wenzake. Karibu facility one. Hiyo si manasi ni generation Z. Doc na wadosi ni millennials na wagonjwa ni generation zote.

---

## Scene 2: Hospital Reception
**Characters:** Maryanne (Receptionist), Casypool

**Maryanne:** Habari yako, kasichana?
**Casypool:** Ni Maryanne, si kasichana.
**Maryanne:** Na una kiburi. Na wewe ni receptionist tu. Ungekuwa daktari, si ungekuwa unapepea huko juu na mabawa?
**Maryanne:** Pole, naweza kusaidia aje?
**Casypool:** Nataka kutahiri.
**Maryanne:** Eh?
**Casypool:** Ni nini hunielewi? Nimesema nataka kutahiri.
**Maryanne:** Anyway, sikuwa na ubaya. Nataka tu kuchukua details. Utalipa 500 ya registration, and then uende consultation, utasaidika.
**Casypool:** Nikuulize, kama sina 500 ya registration, inamaanisha hutaniregister?
**Maryanne:** Eh, for first timers. Ama ushaitahiri hapa before?
**Casypool:** Msichana una kiburi. How old are you? Hii maneno ya registration, consultation, ni ya nini? Unataka kutengeneza pesa kwa kutahiri kwangu? You want to profit from my private organization?
**Maryanne:** Eh, mzee, tulia usaidike. Wengine wakitaka kuretire, wewe unakuja hapa kutahiri. Anyway, full details. Name?
**Casypool:** Kasipul Kapon.
**Maryanne:** Age?
**Casypool:** 21.
**Maryanne:** Mzee, kuwa serious. Kama wewe ni mamangu, basi andika age mine.
**Maryanne:** Mzee, uko na miaka ngapi?
**Casypool:** Basi, kama itakufurahisha, 52.
**Maryanne:** Wacha nisiongee vibaya. Next of kin?
**Casypool:** Ah, weka Catherine. Ah, ah, no, no, no, no. Not Catherine. Janet. Ah, no. Usiweke Janet. Shiro. No. Ah, weka Martha.
**Maryanne:** Mzee, unajua unasumbua wewe? Hebu keti hapo. Utalipa 500 ya registration, utaona daktari na 1,000, alafu 15k, hiyo ni ya circumcision.
**Casypool:** Silipi kitu, niko na insurance.
**Maryanne:** Ni sawa, uko na insurance. But sasa sijui kama insurance inakava wazee wa umri yako kutahiri. Utangoja hapo usikie kama watakubali.
**Casipul:** Sasa, kitu ya kutokutahiri, utaenda kutangazia insurance ati sijatahiri?
**Maryanne:** Mzee, lakini vitu zingine ni za kujisimamia. Hizi ni aibu gani za ati, "Oh, mzee wa 52 years, circumcision na NHIF." Surely. Surely.

---

## Scene 3: Waiting Room
**Characters:** Wambo, Other Patient

**Wambo:** Niko hospitali, nilikuwa nimeshikwa na malaria kidogo. Nakupigia. Wacha nitibiwe, nakupigia.
**Man:** Eh, ndugu yangu. Naomba unichachawize kidogo ili unieleweshe kimantiki. Hivi mbona imekuchukua muda mrefu kabla hujakuwa ndume? Maanake hapo umekuja kupoteza hela, ndugu yangu. Si uongo. Mimi nina plan nzuri kabisa. Maanake siamini haya mambo ya kisayansi.
**Casypool:** Wewe, my friend, mind your own business. Wanaume wenzako wako busy kwa ofisi wanapiga ma deal. Wewe uko busy hapa unapiga dasta. Chunga mdomo yako na umind your own business.
**Man:** Ah, shakum sima tusi. Ila sijakusudia kukukera nyongo. Ni mtazamo tu. Eh? Hata wewe, mabingwa wenzako huku hivyo wanafanya kazi, kule hiv...

---

## Scene 4: Doctor's Consultation
**Characters:** Doc Shifta, Casypool, Maryanne

**Doctor:** Mzee, please, please. Naomba utoe akili kwa lodging, rudisha akili hospitali. Eh, Maryanne. Kuja ushughulike... Hebu mpeleke huko, get him ready.
**Maryanne:** Guys, daktari amesema mprepare your patient, anakuja kumsaidia. Ndio huyo.
**Sly:** Okay, kuna venye yeye si mtoto, ako 52 years.
**Doctor:** Eh? 52?
**Sly:** Eh?
**Doctor:** 52?
**Maryanne:** I know.
**Doctor:** Eh. Mwambie basi akuje.
**Maryanne:** Sawa.
**Doctor:** Eh.
**Kasipul:** Habari yako?
**Doctor:** So, wewe ndio daktari?
**Doctor:** Hapana, mimi ni KDF. Hapa niko Somalia, napigana na Al-Shabaab.
**Kasipul:** Na hii kibanda yenu mnaita hospitali, mna kiburi all the way kutoka kwa reception mpaka hapa hivi by the way.
**Doctor:** Wewe ndio wa kutahiri?
**Kasipul:** My friend, nina jina naitwa Kasipul. Umesikia?
**Doctor:** Okay, fanya kazi yangu iwe rahisi. Tusianze kubishana. I want my work to be easier. Unasikia joto ama baridi?
**Kasipul:** Wapi?
**Doctor:** Eh, mwili. Mwili mzee.
**Kasipul:** Nikuulize very simple question. Joto na baridi inaingiliana wapi na kutahiri?
**Doctor:** Sasa unataka kujitahiri? Si utulie, usikie maswali. Unataka kujitahiri? Unakohoa ama hukohoi?
**Kasipul:** Ndio nauliza the same question. Kukohoa na kutahiri inaingiliana wapi? Labda hujui, nakuambia kukohoa na kukatwa inaingiliana wapi?
**Doctor:** Tumbo inakuuma?
**Kasipul:** Hapana, tumbo hainiumi.
**Doctor:** Jana ulikula nini?
**Kasipul:** Githeri special.
**Doctor:** Ndio maana tumbo inaongea Kikuyu. Sasa utaenda room number four hapo, utaonyeshwa. Uende upimwe typhoid, malaria, TB, pressure, upimwe urine, upimwe stool, na super gonorrhea. Alafu wakutengeneze hapo, utoe nguo, na kuja kushughulikiwa.
**Kasipul:** Hiyo ya mwisho ndio imenichanganya kidogo. Mbona sikutofahamu? Kutoa nguo, wewe mtoto wa kiume unakuja kunishughulikia. Mbona inaingia hapo?
**Doctor:** Mzee, please, please. Naomba utoe akili kwa lodging, rudisha akili hospitali. Eh, Maryanne. Kuja ushughulike... Hebu mpeleke huko, get him ready.

---

## Scene 5: The Nurses' "Trauma"
**Characters:** Wambo, Sly, Maryanne

**Wambo:** Guys, please pray for us. I think we've seen it all.
**Sly:** We are so scared.
**Wambo:** Guys, imagine we are about to handle the circumcision of a 52-year-old.
**Sly:** 52 years old!
**Wambo:** We are so traumatized.
**Sly:** Na mimi sijui kuhandle nyoka.
**Wambo:** Na mimi... I'm even getting an anxiety attack. We are so traumatized. Please pray for us. Please pray for us.
**Sly:** Lakini msijali, we will upload the full video right after this.
**Wambo:** 'Cause we will record everything.
**Sly:** Don't forget to like, share, subscribe, and follow us.
**Both:** And that's on period!
**Maryanne:** Guys, daktari amesema mprepare your patient, anakuja kumsaidia. Ndio huyo.
**Sly:** Babe.
**Wambo:** America.
**Sly:** Babe, huyu ndio yule new catch?
**Wambo:** [Laughs]
**Sly:** Ndio maana ulikuwa unasema tungoje, sindio?
"""

# --- RAG Logic (Simplified) ---
# This is a highly simplified RAG implementation for demonstration purposes.
# In a real-world scenario, you would use a proper vector database.
def simple_rag_lookup(query):
    """Performs a simple keyword-based lookup in the transcript."""
    # Split transcript into "chunks" (lines or paragraphs)
    chunks = NURSE_TOTO_TRANSCRIPT.split('\n\n')
    
    # Simple keyword matching (case-insensitive)
    relevant_chunks = [
        chunk for chunk in chunks 
        if any(word.lower() in chunk.lower() for word in query.split())
    ]
    
    # Return a maximum of 3 relevant chunks
    return "\n\n".join(relevant_chunks[:3])

# --- LLM API Call ---
def query_api(payload):
    """Sends a request to the Hugging Face Inference API."""
    headers = {"Authorization": f"Bearer {os.environ.get('HF_TOKEN')}"}
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.json()

def generate_response(message, history):
    """Main chat function using the Inference API."""
    
    # 1. Check for API Token
    if "HF_TOKEN" not in os.environ or not os.environ.get("HF_TOKEN"):
        return "Kosa la Usanidi: Tafadhali weka siri ya **HF_TOKEN** katika mipangilio ya Space yako. Huwezi kutumia API bila token."

    # 2. Retrieve context (using simplified RAG)
    context = simple_rag_lookup(message)

    # 3. Construct the prompt
    system_prompt = (
        "Wewe ni msaidizi wa AI unayezungumza Kiswahili na Sheng. "
        "Unaweza kufanya mazungumzo ya kawaida au kujibu maswali kuhusu 'Nurse Toto' "
        "kwa kutumia muktadha uliotolewa hapa chini. "
        "Ikiwa swali halihusiani na Nurse Toto, jibu kwa kutumia maarifa yako ya jumla."
    )
    
    full_prompt = f"{system_prompt}\n\nMuktadha wa Nurse Toto:\n{context}\n\nUser: {message}\nAssistant:"

    # 4. Generate response using the API
    payload = {
        "inputs": full_prompt,
        "parameters": {
            "max_new_tokens": 256,
            "temperature": 0.7,
            "return_full_text": False
        }
    }
    
    try:
        api_response = query_api(payload)
        
        if isinstance(api_response, list) and api_response:
            response_text = api_response[0].get("generated_text", "").strip()
            
            # Clean up the response, as the LLM might repeat the prompt
            if "Assistant:" in response_text:
                response_text = response_text.split("Assistant:")[-1].strip()
            
            return response_text
        
        elif isinstance(api_response, dict) and "error" in api_response:
            # Handle API errors (e.g., model loading, rate limit)
            error_msg = api_response["error"]
            if "Authorization" in error_msg or "Invalid token" in error_msg:
                return "Kosa la Uthibitishaji: Tafadhali hakikisha siri ya **HF_TOKEN** ni sahihi na ina ruhusa ya 'read'."
            return f"Kosa la API: {error_msg}. Jaribu tena."
            
        return "Samahani, sikupata jibu kutoka kwa mfumo wa lugha."
        
    except requests.exceptions.RequestException as e:
        return f"Kosa la Mtandao: Imeshindwa kuunganisha na API. Kosa: {e}"
    except Exception as e:
        return f"Kosa lisilotarajiwa: {e}"

# --- Gradio Interface ---
gr.ChatInterface(
    fn=generate_response,
    title="Zero-Memory Swahili/Sheng Chatbot (API-Powered)",
    description="Chat na AI kwa Kiswahili au Sheng! Inajua mambo ya Nurse Toto na mambo mengine ya kawaida. **Kumbuka:** Unahitaji kuweka siri ya **HF_TOKEN** katika mipangilio ya Space.",
    examples=[
        ["Habari yako? Unaweza kunisaidia nini leo?"],
        ["Nieleze kuhusu Casypool kwenye Nurse Toto."],
        ["Sheng ya 'How are you' ni gani?"],
        ["Mzee alitaka kufanya nini hospitalini?"],
    ]
).launch()