NovaRocket/Conversational_AI

1# -*- coding: utf-8 -*-2"""ConvoAI.ipynb3 4Automatically generated by Colab.5 6Original file is located at7    https://colab.research.google.com/drive/1joGBryAm8BIWFs4pZrZZmmp4OXBuMQZH8"""9 10# !pip install llama-cpp-python gradio -q11 12# !pip install --upgrade transformers sentence-transformers tensorflow13# !pip install --upgrade llama-cpp-python gradio pymilvus14# !pip install elevenlabs SpeechRecognition15 16import gradio as gr17import os18import re19import uuid20import time21import random22import datetime23import tempfile24import numpy as np25 26from transformers import pipeline27from elevenlabs.client import ElevenLabs28from elevenlabs import VoiceSettings29from pydub import AudioSegment30import speech_recognition as sr31 32# Milvus33from pymilvus import (34    connections, utility, Collection, CollectionSchema,35    FieldSchema, DataType36)37 38# Sentence embeddings39from sentence_transformers import SentenceTransformer40 41# llama-cpp for local LLM inference42from llama_cpp import Llama43 44import os45ELEVENLABS_API_KEY = os.environ.get('ELEVENLABS_API_KEY')46HUGGINGFACE_TOKEN  = os.environ.get('HF_TOKEN')47MILVUS_URI         = os.environ.get('MILVUS_URI')48MILVUS_USER        = os.environ.get('MILVUS_USER')49MILVUS_PASSWORD    = os.environ.get('MILVUS_PASSWORD')50 51# VOICE_ID_FALLBACK  = "1SM7GgM6IMuvQlz2BwM3"   # Mark - ConvoAI52VOICE_ID_FALLBACK  = "iP95p4xoKVk53GoZ742B"53 54# MODEL INIT55print("Loading embedding model...")56embed_model = SentenceTransformer("all-MiniLM-L6-v2")57 58print("Loading LLM...")59llm = Llama.from_pretrained(60    repo_id="bartowski/Llama-3.2-1B-Instruct-GGUF",61    filename="Llama-3.2-1B-Instruct-Q4_K_L.gguf",62    chat_format="chatml",63    n_ctx=2048,64    n_threads=4,65    n_gpu_layers=35,66    verbose=True,67)68 69print("Loading emotion model...")70emotion_analyzer = pipeline(71    "text-classification",72    model="bhadresh-savani/distilbert-base-uncased-emotion",73    top_k=174)75 76print("Initialising ElevenLabs...")77elevenlabs_client = ElevenLabs(api_key=ELEVENLABS_API_KEY)78 79# Resolve preferred voice ID80VOICE_ID = VOICE_ID_FALLBACK81try:82    voices = elevenlabs_client.voices.get_all()83    for v in voices.voices:84        if v.name == "Mark - ConvoAI":85            VOICE_ID = v.voice_id86            break87except Exception as e:88    print(f"Voice loading warning: {e}")89 90recognizer = sr.Recognizer()91 92# MILVUS SETUP93print("Connecting to Milvus (Zilliz Cloud)...")94connections.connect(95    alias="default",96    uri=MILVUS_URI,97    user=MILVUS_USER,98    password=MILVUS_PASSWORD,99)100 101# ── chat_history collection (assumed to exist already; create if not) ──102if not utility.has_collection("chat_history"):103    chat_schema = CollectionSchema([104        FieldSchema("id",        DataType.VARCHAR, is_primary=True, max_length=64),105        FieldSchema("session_id",DataType.VARCHAR, max_length=64),106        FieldSchema("content",   DataType.VARCHAR, max_length=4096),107        FieldSchema("role",      DataType.VARCHAR, max_length=16),108        FieldSchema("timestamp", DataType.VARCHAR, max_length=64),109        FieldSchema("embedding", DataType.FLOAT_VECTOR, dim=384),110    ], description="Chat messages")111    chat_collection = Collection("chat_history", schema=chat_schema)112    chat_collection.create_index(113        "embedding",114        {"index_type": "IVF_FLAT", "metric_type": "COSINE", "params": {"nlist": 100}}115    )116else:117    chat_collection = Collection("chat_history")118chat_collection.load()119 120# ── summaries collection ──121if not utility.has_collection("summaries"):122    summary_schema = CollectionSchema([123        FieldSchema("id",               DataType.VARCHAR, is_primary=True, max_length=64),124        FieldSchema("session_id",       DataType.VARCHAR, max_length=64),125        FieldSchema("summary",          DataType.VARCHAR, max_length=4096),126        FieldSchema("timestamp",        DataType.VARCHAR, max_length=64),127        FieldSchema("summary_embedding",DataType.FLOAT_VECTOR, dim=384),128    ], description="Session-level summaries")129    summary_collection = Collection("summaries", schema=summary_schema)130    summary_collection.create_index(131        "summary_embedding",132        {"index_type": "IVF_FLAT", "metric_type": "COSINE", "params": {"nlist": 100}}133    )134else:135    summary_collection = Collection("summaries")136summary_collection.load()137 138# ── user_profiles collection ──139if not utility.has_collection("user_profiles"):140    profile_schema = CollectionSchema([141        FieldSchema("id",               DataType.VARCHAR, is_primary=True, max_length=64),142        FieldSchema("session_id",       DataType.VARCHAR, max_length=64),143        FieldSchema("name",             DataType.VARCHAR, max_length=128),144        FieldSchema("age",              DataType.INT64),145        FieldSchema("profession",       DataType.VARCHAR, max_length=128),146        FieldSchema("likes",            DataType.VARCHAR, max_length=512),147        FieldSchema("dislikes",         DataType.VARCHAR, max_length=512),148        FieldSchema("profile_embedding",DataType.FLOAT_VECTOR, dim=384),149    ], description="User profile with embeddings")150    profile_collection = Collection("user_profiles", schema=profile_schema)151    profile_collection.create_index(152        "profile_embedding",153        {"index_type": "IVF_FLAT", "metric_type": "COSINE", "params": {"nlist": 100}}154    )155else:156    profile_collection = Collection("user_profiles")157profile_collection.load()158 159TEMPLATES = {160    "initial_greeting": [161        "Hey there! How's life treating you?",162        "Oh hi! Was just thinking about you.",163        "Hey! What's new?",164        "Yo! Long time no chat!",165        "Hiiii! How've you been?",166        "Hey, what's up!",167        "Hey pal! How's life in your world?",168        "Well look who's here! How you been?",169    ],170    "sadness": [171        "Aw man, that sounds rough... wanna talk about it?",172        "Oh no — I'm here if you need to vent.",173        "Mmm... I get that feeling. It'll pass, promise.",174        "Yeah, some days just drain you huh?",175        "*pats shoulder* You're stronger than you think.",176    ],177    "fatigue": [178        "Ugh, the tiredness struggle is real today huh?",179        "Been there... maybe some tea and deep breaths?",180        "Your body's telling you to slow down maybe?",181        "Tired brains are the worst. Be kind to yourself.",182        "Mmm... wanna just sit with this feeling for a bit?",183    ],184}185 186DEFAULT_SYSTEM_PROMPT = """You're a close AI companion friend having a natural conversation. Guidelines:1871. FIRST MESSAGE ONLY: Give a warm greeting if user says hello/hi.1882. AFTER FIRST MESSAGE: Never greet again, continue conversation naturally.1893. Be human-like: Use "Yeah...", "Mmm...", "I get that".1904. Show personality: "Oh wow!", "No way!", "Seriously?"1915. Mirror the user's emotional tone.1926. Never sound robotic or like customer service."""193 194 195# MILVUS HELPER FUNCTIONS196 197def store_message(session_id: str, role: str, content: str):198    embedding = embed_model.encode(content).tolist()199    chat_collection.insert([200        [str(uuid.uuid4())],201        [session_id],202        [content[:4096]],203        [role],204        [datetime.datetime.utcnow().isoformat()],205        [embedding],206    ])207    chat_collection.flush()208 209 210def recall_relevant_summary(user_query: str):211    embedding = embed_model.encode(user_query).tolist()212    results = summary_collection.search(213        data=[embedding],214        anns_field="summary_embedding",215        param={"metric_type": "COSINE", "params": {"nprobe": 10}},216        limit=1,217        output_fields=["summary", "session_id"],218    )219    if results and results[0]:220        return results[0][0].entity.get("summary")221    return None222 223 224def generate_and_store_summary(session_id: str, history: list):225    history_text = "\n".join([f"{r}: {m}" for r, m in history])226    prompt = f"Summarize this conversation briefly:\n{history_text}\nSummary:"227    response = llm.create_chat_completion([{"role": "user", "content": prompt}])228    summary = response["choices"][0]["message"]["content"]229    embedding = embed_model.encode(summary).tolist()230    summary_collection.insert([231        [str(uuid.uuid4())],232        [session_id],233        [summary[:4096]],234        [datetime.datetime.utcnow().isoformat()],235        [embedding],236    ])237    summary_collection.flush()238    return summary239 240 241def extract_and_store_profile(message: str, session_id: str):242    profile_info = {}243 244    name_match = re.search(r"\bmy name is ([A-Z][a-z]+)", message, re.IGNORECASE)245    if name_match:246        profile_info["name"] = name_match.group(1).title()247 248    age_match = re.search(249        r"\b(i am|i'm|my age is) (\d{1,3})\b", message, re.IGNORECASE250    )251    if age_match:252        profile_info["age"] = int(age_match.group(2))253 254    profession_match = re.search(255        r"\b(i am|i'm|i'm a|i work as|my profession is) (a |an )?([\w\s]+)",256        message, re.IGNORECASE257    )258    if profession_match:259        profile_info["profession"] = profession_match.group(3).strip().capitalize()260 261    likes_match = re.search(262        r"\b(i like|i love|i enjoy) ([\w\s,]+)", message, re.IGNORECASE263    )264    if likes_match:265        profile_info["likes"] = likes_match.group(2).strip()266 267    dislikes_match = re.search(268        r"\b(i hate|i dislike|i don't like|i do not like) ([\w\s,]+)",269        message, re.IGNORECASE270    )271    if dislikes_match:272        profile_info["dislikes"] = dislikes_match.group(2).strip()273 274    if not profile_info:275        return  # nothing to update276 277    existing = profile_collection.query(278        expr=f"session_id == '{session_id}'",279        output_fields=["id", "name", "age", "profession", "likes", "dislikes"],280    )281 282    if existing:283        doc = existing[0]284        updated = {285            "id":         doc["id"],286            "session_id": session_id,287            "name":       profile_info.get("name",       doc.get("name", "")),288            "age":        profile_info.get("age",        doc.get("age", 0)),289            "profession": profile_info.get("profession", doc.get("profession", "")),290            "likes":      profile_info.get("likes",      doc.get("likes", "")),291            "dislikes":   profile_info.get("dislikes",   doc.get("dislikes", "")),292        }293        profile_collection.delete(f"id in ['{doc['id']}']")294        profile_collection.flush()295    else:296        updated = {297            "id":         str(uuid.uuid4()),298            "session_id": session_id,299            "name":       profile_info.get("name", ""),300            "age":        profile_info.get("age", 0),301            "profession": profile_info.get("profession", ""),302            "likes":      profile_info.get("likes", ""),303            "dislikes":   profile_info.get("dislikes", ""),304        }305 306    profile_text = " ".join([307        updated["name"], updated["profession"],308        updated["likes"], updated["dislikes"]309    ])310    embedding = embed_model.encode(profile_text).tolist()311 312    profile_collection.insert([313        [updated["id"]],314        [updated["session_id"]],315        [updated["name"]],316        [updated["age"]],317        [updated["profession"]],318        [updated["likes"]],319        [updated["dislikes"]],320        [embedding],321    ])322    profile_collection.flush()323    print(f"[Profile] updated for session {session_id}: {profile_info}")324 325 326def get_user_profile_memory(session_id: str) -> str:327    results = profile_collection.query(328        expr=f"session_id == '{session_id}'",329        output_fields=["name", "age", "profession", "likes", "dislikes"],330    )331    if not results:332        return ""333    p = results[0]334    parts = []335    if p.get("name"):       parts.append(f"Name: {p['name']}")336    if p.get("age"):        parts.append(f"Age: {p['age']}")337    if p.get("profession"): parts.append(f"Profession: {p['profession']}")338    if p.get("likes"):      parts.append(f"Likes: {p['likes']}")339    if p.get("dislikes"):   parts.append(f"Dislikes: {p['dislikes']}")340    return "\n".join(parts)341 342 343 344# AUDIO HELPERS345def process_audio_input(audio) -> str:346    """Convert microphone audio to text via Google STT."""347    if audio is None:348        return ""349    try:350        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:351            if isinstance(audio, tuple):352                sample_rate, audio_data = audio353                seg = AudioSegment(354                    audio_data.tobytes(),355                    frame_rate=sample_rate,356                    sample_width=audio_data.dtype.itemsize,357                    channels=1,358                )359            else:360                seg = AudioSegment.from_file(audio)361            seg.export(tmp.name, format="wav")362            tmp_path = tmp.name363 364        with sr.AudioFile(tmp_path) as source:365            recognizer.adjust_for_ambient_noise(source, duration=0.5)366            audio_data = recognizer.record(source)367            try:368                text = recognizer.recognize_google(audio_data)369            except sr.UnknownValueError:370                text = ""371            except sr.RequestError:372                text = ""373        os.unlink(tmp_path)374        return text375    except Exception as e:376        print(f"[STT error] {e}")377        return ""378 379 380def generate_tts(text: str):381    """Generate TTS via ElevenLabs and return a temp file path."""382    try:383        audio_response = elevenlabs_client.text_to_speech.convert(384            voice_id=VOICE_ID,385            model_id="eleven_multilingual_v2",386            text=text,387            voice_settings=VoiceSettings(388                stability=0.5,389                similarity_boost=0.8,390                style=0.2,391                speaker_boost=True,392            ),393        )394        audio_bytes = b"".join(audio_response)395        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as f:396            f.write(audio_bytes)397            return f.name398    except Exception as e:399        print(f"[TTS error] {e}")400        return None401 402 403 404# EMOTION & TEMPLATE ROUTING405 406_greeting_done: dict[str, bool] = {}   # keyed by session_id407 408 409def get_template_response(user_input: str, session_id: str):410    """Return a hardcoded template reply when appropriate, else None."""411    lower = user_input.lower()412 413    if not _greeting_done.get(session_id) and any(414        w in lower for w in ["hello", "hi", "hey"]415    ):416        _greeting_done[session_id] = True417        return random.choice(TEMPLATES["initial_greeting"])418 419    if "tired" in lower or "exhaust" in lower:420        return random.choice(TEMPLATES["fatigue"])421 422    try:423        emotion_result = emotion_analyzer(user_input)[0][0]424        if emotion_result["label"].lower() == "sadness" and emotion_result["score"] > 0.6:425            return random.choice(TEMPLATES["sadness"])426    except Exception:427        pass428 429    return None430 431 432 433# CORE CHAT FUNCTION434 435def truncate_history(messages: list, max_turns: int = 10) -> list:436    return messages[-max_turns * 2:]437 438 439def chat_pipeline(audio, text_input: str, history: list, session_id: str):440    """441    Main Gradio handler.442    Returns: (updated_history, cleared_text, audio_file_path, session_id)443    """444    if not session_id:445        session_id = str(int(time.time()))446 447    # ── Resolve user input ──448    user_input = ""449    if audio:450        user_input = process_audio_input(audio)451    if not user_input and text_input:452        user_input = text_input.strip()453    if not user_input:454        return history, "", None, session_id455 456    # ── Profile extraction ──457    extract_and_store_profile(user_input, session_id)458 459    # ── Memory context ──460    memory_context = ""461    recalled = recall_relevant_summary(user_input)462    if recalled:463        memory_context += f"[Memory from a past session]:\n{recalled}\n\n"464 465    user_profile = get_user_profile_memory(session_id)466    if user_profile:467        memory_context += f"[What I know about the user]:\n{user_profile}\n\n"468 469    # ── Check for template response ──470    template_reply = get_template_response(user_input, session_id)471    if template_reply:472        history.append({"role": "user",      "content": user_input})473        history.append({"role": "assistant", "content": template_reply})474        store_message(session_id, "user",      user_input)475        store_message(session_id, "assistant", template_reply)476        audio_out = generate_tts(template_reply)477 478        if user_input.lower().strip() in ["end session", "exit", "bye"]:479            pairs = []480            for i in range(0, len(history) - 1, 2):481                u = history[i].get("content", "")482                b = history[i+1].get("content", "") if i+1 < len(history) else ""483                pairs.append((u, b))484            generate_and_store_summary(session_id, pairs)485 486        return history, "", audio_out, session_id487 488    # ── Build message list for LLM ──489    messages = [490        {491            "role": "system",492            "content": memory_context + DEFAULT_SYSTEM_PROMPT,493        }494    ]495    for msg in history:496        messages.append({"role": msg["role"], "content": msg["content"]})497    messages.append({"role": "user", "content": user_input})498 499    # ── Stage 1: Thinking phase (internal monologue) ──500    thinking_prompt = messages + [{501        "role": "user",502        "content": (503            "Think through your reasoning before responding. "504            "What should you consider before replying to the user?"505        ),506    }]507    thinking_response = llm.create_chat_completion(thinking_prompt)508    thinking_text = thinking_response["choices"][0]["message"]["content"]509    print(f"\n[Thinking Phase]:\n{thinking_text}\n")510 511    # ── Stage 2: Final response informed by thinking ──512    final_messages = truncate_history(messages) + [{513        "role": "system",514        "content": f"Here's your thought process:\n{thinking_text}\nNow generate the final reply.",515    }]516 517    final_response = llm.create_chat_completion(final_messages)518    reply = final_response["choices"][0]["message"]["content"].strip()519 520    # ── Persist ──521    store_message(session_id, "user",      user_input)522    store_message(session_id, "assistant", reply)523    history.append({"role": "user",      "content": user_input})524    history.append({"role": "assistant", "content": reply})525 526    # ── TTS ──527    audio_out = generate_tts(reply)528 529    # ── End-session summary ──530    if user_input.lower().strip() in ["end session", "exit", "bye"]:531        pairs = []532        for i in range(0, len(history) - 1, 2):533            u = history[i].get("content", "")534            b = history[i+1].get("content", "") if i+1 < len(history) else ""535            pairs.append((u, b))536        generate_and_store_summary(session_id, pairs)537        print("[Session summarized and stored]")538 539    return history, "", audio_out, session_id540 541 542def clear_chat(session_id: str):543    new_session = str(int(time.time()))544    if session_id in _greeting_done:545        del _greeting_done[session_id]546    return [], "", None, new_session547 548 549# ──────────────────────────────────────────────550# GRADIO UI551# ──────────────────────────────────────────────552 553css = """554body { font-family: 'Segoe UI', sans-serif; }555#chatbot { height: 480px; overflow-y: auto; }556footer { display: none !important; }557"""558 559with gr.Blocks(title="Conversational AI") as demo:560 561    gr.Markdown("## 🎙️ Conversational AI Companion")562    gr.Markdown(563        "Talk to your AI companion via **voice** or **text**. "564        "Say *'end session'* / *'bye'* to save a session summary."565    )566 567    session_id_state = gr.State(str(int(time.time())))568 569    chatbot = gr.Chatbot(label="Conversation", elem_id="chatbot")570 571    with gr.Row():572        audio_input = gr.Audio(573            sources=["microphone"],574            type="filepath",575            label="🎤 Speak",576        )577        audio_output = gr.Audio(578            label="🔊 Response",579            type="filepath",580            interactive=False,581            autoplay=True,582        )583 584    with gr.Row():585        text_input = gr.Textbox(586            placeholder="Type your message here…",587            label="Message",588            scale=5,589            show_label=False,590        )591        submit_btn = gr.Button("Send", variant="primary", scale=1)592        clear_btn  = gr.Button("Clear", scale=1)593 594    # ── Event wiring ──595    submit_btn.click(596        chat_pipeline,597        inputs=[audio_input, text_input, chatbot, session_id_state],598        outputs=[chatbot, text_input, audio_output, session_id_state],599    )600    text_input.submit(601        chat_pipeline,602        inputs=[audio_input, text_input, chatbot, session_id_state],603        outputs=[chatbot, text_input, audio_output, session_id_state],604    )605    clear_btn.click(606        clear_chat,607        inputs=[session_id_state],608        outputs=[chatbot, text_input, audio_output, session_id_state],609    )610 611demo.launch(css=css,ssr_mode=False)