app.py

# -*- coding: utf-8 -*-
"""ConvoAI.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1joGBryAm8BIWFs4pZrZZmmp4OXBuMQZH
"""

!pip install llama-cpp-python gradio -q

!pip install --upgrade transformers sentence-transformers tensorflow
!pip install --upgrade llama-cpp-python gradio pymilvus
!pip install elevenlabs SpeechRecognition

import gradio as gr
import os
import re
import uuid
import time
import random
import datetime
import tempfile
import numpy as np

from transformers import pipeline
from elevenlabs.client import ElevenLabs
from elevenlabs import VoiceSettings
from pydub import AudioSegment
import speech_recognition as sr

# Milvus
from pymilvus import (
    connections, utility, Collection, CollectionSchema,
    FieldSchema, DataType
)

# Sentence embeddings
from sentence_transformers import SentenceTransformer

# llama-cpp for local LLM inference
from llama_cpp import Llama

import os
ELEVENLABS_API_KEY = os.environ.get('ELEVENLABS_API_KEY')
HUGGINGFACE_TOKEN  = os.environ.get('HF_TOKEN')
MILVUS_URI         = os.environ.get('MILVUS_URI')
MILVUS_USER        = os.environ.get('MILVUS_USER')
MILVUS_PASSWORD    = os.environ.get('MILVUS_PASSWORD')

# VOICE_ID_FALLBACK  = "1SM7GgM6IMuvQlz2BwM3"   # Mark - ConvoAI
VOICE_ID_FALLBACK  = "iP95p4xoKVk53GoZ742B"

# MODEL INIT
print("Loading embedding model...")
embed_model = SentenceTransformer("all-MiniLM-L6-v2")

print("Loading LLM...")
llm = Llama.from_pretrained(
    repo_id="bartowski/Llama-3.2-1B-Instruct-GGUF",
    filename="Llama-3.2-1B-Instruct-Q4_K_L.gguf",
    chat_format="chatml",
    n_ctx=2048,
    n_threads=4,
    n_gpu_layers=35,
    verbose=True,
)

print("Loading emotion model...")
emotion_analyzer = pipeline(
    "text-classification",
    model="bhadresh-savani/distilbert-base-uncased-emotion",
    top_k=1
)

print("Initialising ElevenLabs...")
elevenlabs_client = ElevenLabs(api_key=ELEVENLABS_API_KEY)

# Resolve preferred voice ID
VOICE_ID = VOICE_ID_FALLBACK
try:
    voices = elevenlabs_client.voices.get_all()
    for v in voices.voices:
        if v.name == "Mark - ConvoAI":
            VOICE_ID = v.voice_id
            break
except Exception as e:
    print(f"Voice loading warning: {e}")

recognizer = sr.Recognizer()

# MILVUS SETUP
print("Connecting to Milvus (Zilliz Cloud)...")
connections.connect(
    alias="default",
    uri=MILVUS_URI,
    user=MILVUS_USER,
    password=MILVUS_PASSWORD,
)

# ── chat_history collection (assumed to exist already; create if not) ──
if not utility.has_collection("chat_history"):
    chat_schema = CollectionSchema([
        FieldSchema("id",        DataType.VARCHAR, is_primary=True, max_length=64),
        FieldSchema("session_id",DataType.VARCHAR, max_length=64),
        FieldSchema("content",   DataType.VARCHAR, max_length=4096),
        FieldSchema("role",      DataType.VARCHAR, max_length=16),
        FieldSchema("timestamp", DataType.VARCHAR, max_length=64),
        FieldSchema("embedding", DataType.FLOAT_VECTOR, dim=384),
    ], description="Chat messages")
    chat_collection = Collection("chat_history", schema=chat_schema)
    chat_collection.create_index(
        "embedding",
        {"index_type": "IVF_FLAT", "metric_type": "COSINE", "params": {"nlist": 100}}
    )
else:
    chat_collection = Collection("chat_history")
chat_collection.load()

# ── summaries collection ──
if not utility.has_collection("summaries"):
    summary_schema = CollectionSchema([
        FieldSchema("id",               DataType.VARCHAR, is_primary=True, max_length=64),
        FieldSchema("session_id",       DataType.VARCHAR, max_length=64),
        FieldSchema("summary",          DataType.VARCHAR, max_length=4096),
        FieldSchema("timestamp",        DataType.VARCHAR, max_length=64),
        FieldSchema("summary_embedding",DataType.FLOAT_VECTOR, dim=384),
    ], description="Session-level summaries")
    summary_collection = Collection("summaries", schema=summary_schema)
    summary_collection.create_index(
        "summary_embedding",
        {"index_type": "IVF_FLAT", "metric_type": "COSINE", "params": {"nlist": 100}}
    )
else:
    summary_collection = Collection("summaries")
summary_collection.load()

# ── user_profiles collection ──
if not utility.has_collection("user_profiles"):
    profile_schema = CollectionSchema([
        FieldSchema("id",               DataType.VARCHAR, is_primary=True, max_length=64),
        FieldSchema("session_id",       DataType.VARCHAR, max_length=64),
        FieldSchema("name",             DataType.VARCHAR, max_length=128),
        FieldSchema("age",              DataType.INT64),
        FieldSchema("profession",       DataType.VARCHAR, max_length=128),
        FieldSchema("likes",            DataType.VARCHAR, max_length=512),
        FieldSchema("dislikes",         DataType.VARCHAR, max_length=512),
        FieldSchema("profile_embedding",DataType.FLOAT_VECTOR, dim=384),
    ], description="User profile with embeddings")
    profile_collection = Collection("user_profiles", schema=profile_schema)
    profile_collection.create_index(
        "profile_embedding",
        {"index_type": "IVF_FLAT", "metric_type": "COSINE", "params": {"nlist": 100}}
    )
else:
    profile_collection = Collection("user_profiles")
profile_collection.load()

TEMPLATES = {
    "initial_greeting": [
        "Hey there! How's life treating you?",
        "Oh hi! Was just thinking about you.",
        "Hey! What's new?",
        "Yo! Long time no chat!",
        "Hiiii! How've you been?",
        "Hey, what's up!",
        "Hey pal! How's life in your world?",
        "Well look who's here! How you been?",
    ],
    "sadness": [
        "Aw man, that sounds rough... wanna talk about it?",
        "Oh no — I'm here if you need to vent.",
        "Mmm... I get that feeling. It'll pass, promise.",
        "Yeah, some days just drain you huh?",
        "*pats shoulder* You're stronger than you think.",
    ],
    "fatigue": [
        "Ugh, the tiredness struggle is real today huh?",
        "Been there... maybe some tea and deep breaths?",
        "Your body's telling you to slow down maybe?",
        "Tired brains are the worst. Be kind to yourself.",
        "Mmm... wanna just sit with this feeling for a bit?",
    ],
}

DEFAULT_SYSTEM_PROMPT = """You're a close AI companion friend having a natural conversation. Guidelines:
1. FIRST MESSAGE ONLY: Give a warm greeting if user says hello/hi.
2. AFTER FIRST MESSAGE: Never greet again, continue conversation naturally.
3. Be human-like: Use "Yeah...", "Mmm...", "I get that".
4. Show personality: "Oh wow!", "No way!", "Seriously?"
5. Mirror the user's emotional tone.
6. Never sound robotic or like customer service."""


# MILVUS HELPER FUNCTIONS

def store_message(session_id: str, role: str, content: str):
    embedding = embed_model.encode(content).tolist()
    chat_collection.insert([
        [str(uuid.uuid4())],
        [session_id],
        [content[:4096]],
        [role],
        [datetime.datetime.utcnow().isoformat()],
        [embedding],
    ])
    chat_collection.flush()


def recall_relevant_summary(user_query: str):
    embedding = embed_model.encode(user_query).tolist()
    results = summary_collection.search(
        data=[embedding],
        anns_field="summary_embedding",
        param={"metric_type": "COSINE", "params": {"nprobe": 10}},
        limit=1,
        output_fields=["summary", "session_id"],
    )
    if results and results[0]:
        return results[0][0].entity.get("summary")
    return None


def generate_and_store_summary(session_id: str, history: list):
    history_text = "\n".join([f"{r}: {m}" for r, m in history])
    prompt = f"Summarize this conversation briefly:\n{history_text}\nSummary:"
    response = llm.create_chat_completion([{"role": "user", "content": prompt}])
    summary = response["choices"][0]["message"]["content"]
    embedding = embed_model.encode(summary).tolist()
    summary_collection.insert([
        [str(uuid.uuid4())],
        [session_id],
        [summary[:4096]],
        [datetime.datetime.utcnow().isoformat()],
        [embedding],
    ])
    summary_collection.flush()
    return summary


def extract_and_store_profile(message: str, session_id: str):
    profile_info = {}

    name_match = re.search(r"\bmy name is ([A-Z][a-z]+)", message, re.IGNORECASE)
    if name_match:
        profile_info["name"] = name_match.group(1).title()

    age_match = re.search(
        r"\b(i am|i'm|my age is) (\d{1,3})\b", message, re.IGNORECASE
    )
    if age_match:
        profile_info["age"] = int(age_match.group(2))

    profession_match = re.search(
        r"\b(i am|i'm|i'm a|i work as|my profession is) (a |an )?([\w\s]+)",
        message, re.IGNORECASE
    )
    if profession_match:
        profile_info["profession"] = profession_match.group(3).strip().capitalize()

    likes_match = re.search(
        r"\b(i like|i love|i enjoy) ([\w\s,]+)", message, re.IGNORECASE
    )
    if likes_match:
        profile_info["likes"] = likes_match.group(2).strip()

    dislikes_match = re.search(
        r"\b(i hate|i dislike|i don't like|i do not like) ([\w\s,]+)",
        message, re.IGNORECASE
    )
    if dislikes_match:
        profile_info["dislikes"] = dislikes_match.group(2).strip()

    if not profile_info:
        return  # nothing to update

    existing = profile_collection.query(
        expr=f"session_id == '{session_id}'",
        output_fields=["id", "name", "age", "profession", "likes", "dislikes"],
    )

    if existing:
        doc = existing[0]
        updated = {
            "id":         doc["id"],
            "session_id": session_id,
            "name":       profile_info.get("name",       doc.get("name", "")),
            "age":        profile_info.get("age",        doc.get("age", 0)),
            "profession": profile_info.get("profession", doc.get("profession", "")),
            "likes":      profile_info.get("likes",      doc.get("likes", "")),
            "dislikes":   profile_info.get("dislikes",   doc.get("dislikes", "")),
        }
        profile_collection.delete(f"id in ['{doc['id']}']")
        profile_collection.flush()
    else:
        updated = {
            "id":         str(uuid.uuid4()),
            "session_id": session_id,
            "name":       profile_info.get("name", ""),
            "age":        profile_info.get("age", 0),
            "profession": profile_info.get("profession", ""),
            "likes":      profile_info.get("likes", ""),
            "dislikes":   profile_info.get("dislikes", ""),
        }

    profile_text = " ".join([
        updated["name"], updated["profession"],
        updated["likes"], updated["dislikes"]
    ])
    embedding = embed_model.encode(profile_text).tolist()

    profile_collection.insert([
        [updated["id"]],
        [updated["session_id"]],
        [updated["name"]],
        [updated["age"]],
        [updated["profession"]],
        [updated["likes"]],
        [updated["dislikes"]],
        [embedding],
    ])
    profile_collection.flush()
    print(f"[Profile] updated for session {session_id}: {profile_info}")


def get_user_profile_memory(session_id: str) -> str:
    results = profile_collection.query(
        expr=f"session_id == '{session_id}'",
        output_fields=["name", "age", "profession", "likes", "dislikes"],
    )
    if not results:
        return ""
    p = results[0]
    parts = []
    if p.get("name"):       parts.append(f"Name: {p['name']}")
    if p.get("age"):        parts.append(f"Age: {p['age']}")
    if p.get("profession"): parts.append(f"Profession: {p['profession']}")
    if p.get("likes"):      parts.append(f"Likes: {p['likes']}")
    if p.get("dislikes"):   parts.append(f"Dislikes: {p['dislikes']}")
    return "\n".join(parts)



# AUDIO HELPERS
def process_audio_input(audio) -> str:
    """Convert microphone audio to text via Google STT."""
    if audio is None:
        return ""
    try:
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
            if isinstance(audio, tuple):
                sample_rate, audio_data = audio
                seg = AudioSegment(
                    audio_data.tobytes(),
                    frame_rate=sample_rate,
                    sample_width=audio_data.dtype.itemsize,
                    channels=1,
                )
            else:
                seg = AudioSegment.from_file(audio)
            seg.export(tmp.name, format="wav")
            tmp_path = tmp.name

        with sr.AudioFile(tmp_path) as source:
            recognizer.adjust_for_ambient_noise(source, duration=0.5)
            audio_data = recognizer.record(source)
            try:
                text = recognizer.recognize_google(audio_data)
            except sr.UnknownValueError:
                text = ""
            except sr.RequestError:
                text = ""
        os.unlink(tmp_path)
        return text
    except Exception as e:
        print(f"[STT error] {e}")
        return ""


def generate_tts(text: str):
    """Generate TTS via ElevenLabs and return a temp file path."""
    try:
        audio_response = elevenlabs_client.text_to_speech.convert(
            voice_id=VOICE_ID,
            model_id="eleven_multilingual_v2",
            text=text,
            voice_settings=VoiceSettings(
                stability=0.5,
                similarity_boost=0.8,
                style=0.2,
                speaker_boost=True,
            ),
        )
        audio_bytes = b"".join(audio_response)
        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as f:
            f.write(audio_bytes)
            return f.name
    except Exception as e:
        print(f"[TTS error] {e}")
        return None



# EMOTION & TEMPLATE ROUTING

_greeting_done: dict[str, bool] = {}   # keyed by session_id


def get_template_response(user_input: str, session_id: str):
    """Return a hardcoded template reply when appropriate, else None."""
    lower = user_input.lower()

    if not _greeting_done.get(session_id) and any(
        w in lower for w in ["hello", "hi", "hey"]
    ):
        _greeting_done[session_id] = True
        return random.choice(TEMPLATES["initial_greeting"])

    if "tired" in lower or "exhaust" in lower:
        return random.choice(TEMPLATES["fatigue"])

    try:
        emotion_result = emotion_analyzer(user_input)[0][0]
        if emotion_result["label"].lower() == "sadness" and emotion_result["score"] > 0.6:
            return random.choice(TEMPLATES["sadness"])
    except Exception:
        pass

    return None



# CORE CHAT FUNCTION

def truncate_history(messages: list, max_turns: int = 10) -> list:
    return messages[-max_turns * 2:]


def chat_pipeline(audio, text_input: str, history: list, session_id: str):
    """
    Main Gradio handler.
    Returns: (updated_history, cleared_text, audio_file_path, session_id)
    """
    if not session_id:
        session_id = str(int(time.time()))

    # ── Resolve user input ──
    user_input = ""
    if audio:
        user_input = process_audio_input(audio)
    if not user_input and text_input:
        user_input = text_input.strip()
    if not user_input:
        return history, "", None, session_id

    # ── Profile extraction ──
    extract_and_store_profile(user_input, session_id)

    # ── Memory context ──
    memory_context = ""
    recalled = recall_relevant_summary(user_input)
    if recalled:
        memory_context += f"[Memory from a past session]:\n{recalled}\n\n"

    user_profile = get_user_profile_memory(session_id)
    if user_profile:
        memory_context += f"[What I know about the user]:\n{user_profile}\n\n"

    # ── Check for template response ──
    template_reply = get_template_response(user_input, session_id)
    if template_reply:
        history.append({"role": "user",      "content": user_input})
        history.append({"role": "assistant", "content": template_reply})
        store_message(session_id, "user",      user_input)
        store_message(session_id, "assistant", template_reply)
        audio_out = generate_tts(template_reply)

        if user_input.lower().strip() in ["end session", "exit", "bye"]:
            pairs = []
            for i in range(0, len(history) - 1, 2):
                u = history[i].get("content", "")
                b = history[i+1].get("content", "") if i+1 < len(history) else ""
                pairs.append((u, b))
            generate_and_store_summary(session_id, pairs)

        return history, "", audio_out, session_id

    # ── Build message list for LLM ──
    messages = [
        {
            "role": "system",
            "content": memory_context + DEFAULT_SYSTEM_PROMPT,
        }
    ]
    for msg in history:
        messages.append({"role": msg["role"], "content": msg["content"]})
    messages.append({"role": "user", "content": user_input})

    # ── Stage 1: Thinking phase (internal monologue) ──
    thinking_prompt = messages + [{
        "role": "user",
        "content": (
            "Think through your reasoning before responding. "
            "What should you consider before replying to the user?"
        ),
    }]
    thinking_response = llm.create_chat_completion(thinking_prompt)
    thinking_text = thinking_response["choices"][0]["message"]["content"]
    print(f"\n[Thinking Phase]:\n{thinking_text}\n")

    # ── Stage 2: Final response informed by thinking ──
    final_messages = truncate_history(messages) + [{
        "role": "system",
        "content": f"Here's your thought process:\n{thinking_text}\nNow generate the final reply.",
    }]

    final_response = llm.create_chat_completion(final_messages)
    reply = final_response["choices"][0]["message"]["content"].strip()

    # ── Persist ──
    store_message(session_id, "user",      user_input)
    store_message(session_id, "assistant", reply)
    history.append({"role": "user",      "content": user_input})
    history.append({"role": "assistant", "content": reply})

    # ── TTS ──
    audio_out = generate_tts(reply)

    # ── End-session summary ──
    if user_input.lower().strip() in ["end session", "exit", "bye"]:
        pairs = []
        for i in range(0, len(history) - 1, 2):
            u = history[i].get("content", "")
            b = history[i+1].get("content", "") if i+1 < len(history) else ""
            pairs.append((u, b))
        generate_and_store_summary(session_id, pairs)
        print("[Session summarized and stored]")

    return history, "", audio_out, session_id


def clear_chat(session_id: str):
    new_session = str(int(time.time()))
    if session_id in _greeting_done:
        del _greeting_done[session_id]
    return [], "", None, new_session


# ──────────────────────────────────────────────
# GRADIO UI
# ──────────────────────────────────────────────

css = """
body { font-family: 'Segoe UI', sans-serif; }
#chatbot { height: 480px; overflow-y: auto; }
footer { display: none !important; }
"""

with gr.Blocks(css=css, title="Conversational AI") as demo:

    gr.Markdown("## 🎙️ Conversational AI Companion")
    gr.Markdown(
        "Talk to your AI companion via **voice** or **text**. "
        "Say *'end session'* / *'bye'* to save a session summary."
    )

    session_id_state = gr.State(str(int(time.time())))

    chatbot = gr.Chatbot(label="Conversation", elem_id="chatbot")

    with gr.Row():
        audio_input = gr.Audio(
            sources=["microphone"],
            type="filepath",
            label="🎤 Speak",
        )
        audio_output = gr.Audio(
            label="🔊 Response",
            type="filepath",
            interactive=False,
            autoplay=True,
        )

    with gr.Row():
        text_input = gr.Textbox(
            placeholder="Type your message here…",
            label="Message",
            scale=5,
            show_label=False,
        )
        submit_btn = gr.Button("Send", variant="primary", scale=1)
        clear_btn  = gr.Button("Clear", scale=1)

    # ── Event wiring ──
    submit_btn.click(
        chat_pipeline,
        inputs=[audio_input, text_input, chatbot, session_id_state],
        outputs=[chatbot, text_input, audio_output, session_id_state],
    )
    text_input.submit(
        chat_pipeline,
        inputs=[audio_input, text_input, chatbot, session_id_state],
        outputs=[chatbot, text_input, audio_output, session_id_state],
    )
    clear_btn.click(
        clear_chat,
        inputs=[session_id_state],
        outputs=[chatbot, text_input, audio_output, session_id_state],
    )

demo.launch()