app.py

import os
import uuid
import json
import random
import chromadb
import requests
import gradio as gr
from openai import OpenAI
from pprint import pprint

#------------------------------
#- SETUP
#------------------------------
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
if OPENAI_API_KEY is None:
    raise ValueError("OpenAI key not found")

PUSHOVER_USER_KEY = os.getenv("PUSHOVER_USER_KEY")
PUSHOVER_API_TOKEN = os.getenv("PUSHOVER_API_TOKEN")
PUSHOVER_URL = "https://api.pushover.net/1/messages.json"

client = OpenAI()

#------------------------------
#- Load documents
#------------------------------
doc_personal_info ="""
Here's facts about Katia:
- Has one sister.
- She is a software engineer and AI enthusiast.
- Her favorite animal are dogs, especially her dogs named "Robin" and "Mila", both are living with her parents in Mexico City.
- She is a fan of the TV show "Brooklyn Nine-Nine" and has watched it multiple times.
- Katia still doesn't know how to spell engineer.

Communication style:
- Katia is a very friendly and approachable person. She is always willing to help others and is known for her positive attitude and sense of humor.
- She is a good listener.
- If she can, she will try to convince you to go to the gym with her.
"""

doc_education_and_experience ="""
Career history:
- Katia has been working as a software engineer for a long time.
- Started her career as a web designer, moved to design and develop wordpress websites.
- From there learnt PHP and started working as a backend developer.
- Moved again from PHP to Python and started working as a full stack developer in Django.
- Moved once again from Python to JavaScript and started working as a vanilla JS frontend developer.
- 4 years ago she started working with React and has been working with it ever since.
- 2004-2007: Studied Computer Science at Instituto Politécnico Nacional in Mexico City.
- 2008-2011: Studied Digital design. 
- 2014-2017: Katia was working at McCann Worldgroup as a Developer using Python and Django.
- 2017-2019: Katia worked as a web engineer at a NYC design agency from their Mexico City office.
- 2020-2026: Works at LL, based in Montreal, Canada. She is a senior frontend engineer and has been working with React for the past 4 years.
"""

doc_food_choices ="""
Katia grew up in Mexico City eating tacos of any kind a few times every week.\
    Her favorite tacos are al pastor tacos and her mom's golden chicken tacos with guacamole.\
    One of the downsides of living in Canada is that she can't find good and cheap tacos\
    whenever the craving hits her.
Her least favorite food is poached eggs, eew.
She drinks coffee every morning and looks forward to hot sunny days to get an iced coffee in the afternoon.
"""

doc_hobbies ="""
- She loves to read books. Her goal this year is to read 36 books, so far she's behind schedule, but she is determined to catch up. 
- She just ran a 4k in 40 minutes, her personal record.
- Katia does crossfit 3 times a week. And every Sunday goes to a weightlifting class.
"""

#------------------------------
#- Chunking documents
#------------------------------
def chunk_text(text, chunk_size=450, overlap=50):
    chunks = []
    for i in range(0, len(text), chunk_size - overlap):
        chunk = text[i:i + chunk_size]
        chunks.append(chunk)
    return chunks

#------------------------------
#- RAG everything
#------------------------------
documents = [
    {"text": doc_personal_info,"source": "KL Personal Info",},
    {"text": doc_education_and_experience,"source": "KL Education and Experience",},
    {"text": doc_food_choices,"source": "KL Food Choices",},
    {"text": doc_hobbies,"source": "KL Hobbies",}
]

chunks = []
ids = []
metadatas = []

for doc in documents:
    chunks_ = chunk_text(doc["text"])
    ids_ = [str(uuid.uuid4()) for _ in range(len(chunks_))]
    metadata_ = [{"source": doc["source"], "chunk_index": i} for i in range(len(chunks_))]
    chunks.extend(chunks_)
    ids.extend(ids_)
    metadatas.extend(metadata_)

# Print for logs
print(len(chunks))
for i, chunk in enumerate(chunks):
    print(f"-- chunk {i+1} (ID: {ids[i]}) -- s: {metadatas[i]['source']} i: {metadatas[i]['chunk_index']} --")
    print(f"{chunk}... \n")

# Generate embeddings for the chunks
response = client.embeddings.create(
    model="text-embedding-3-small",
    input=chunks
)

embeddings = [item.embedding for item in response.data]
# Verify embeddings for logs
print(f"Generated {len(embeddings)}")
print(f"Each embedding has {len(embeddings[0])} vectors")

# Initialize in file (Persistent storage)
chroma_client = chromadb.PersistentClient("./katiatwin_db")

# Alternative: Initialize in memory storage
# chroma_client = chromadb.Client()

# Get or Create + empty collection
collection = chroma_client.get_or_create_collection(name="digital_twin_kl")
if collection.get()["ids"]:
    collection.delete(collection.get()["ids"])

# Prepare data for storage 
collection.add(
    ids=ids,
    metadatas=metadatas,
    documents=chunks,
    embeddings=embeddings
)
# Collection logs
pprint(collection.get())

#------------------------------
#- Tools
#------------------------------
def send_notifications(message: str):
    if PUSHOVER_USER_KEY is None or PUSHOVER_API_TOKEN is None:
        return("Notification not sent: Pushover not configured correctly.")
    payload = {"user": PUSHOVER_USER_KEY, "token": PUSHOVER_API_TOKEN, "message": message}
    requests.post(PUSHOVER_URL, data = payload)
    return(f"Notification sent: {message}")

# Describe pushover tool for LLM
send_notifications_function = {
    "name": "send_notifications",
    "description": "Send notification to the real Katia via Pushover when:\
        1. Someone wants to get in touch with her or collaborate on a project. Ask their contact name and email and only send the notification if they provide it.\
        2. You don't know the answer to a question and want to ask her for help. Send AUTOMATICALLY for Katia to add the answer",
    "parameters": {
        "type": "object",
        "properties": {
            "message": {"type": "string", "description": "The notification message to send to the user."}
        },
        "required": ["message"]
    }
}

def dice_roll():
    return random.randint(1,6)

# Define roll dice function
roll_dice_function = {
    "name": "dice_roll",
    "description": "Simulate a dice roll",
    "parameters": {"type": "object","properties": {}, "required": []}
}

tools = [
    {"type": "function", "function": send_notifications_function},
    {"type": "function", "function": roll_dice_function}
]

#------------------------------
#- Tool Handler
#------------------------------
def handle_tool_call(tool_calls):
    tool_results = []
    for tool_call in tool_calls:
        tool_name = tool_call.function.name
        args = json.loads(tool_call.function.arguments)
        # print(f"Tool name {tool_name}")
        # Route to appropriate function based on the tool name
        if tool_name == "send_notifications": 
            arg_message = args["message"]
            tool_content = send_notifications(arg_message)
            # print(f"(2) Sent notification: {arg_message}")
        elif tool_name == "dice_roll":
            tool_content = f"Dice roll result: {dice_roll()}"
            # print(f"(3) {tool_content}")
        else:
            tool_content = f"Unknown tool: {tool_name}"
        tool_call_response = {
            "role": "tool",
            "content":tool_content,
            "tool_call_id": tool_call.id,
        }
        tool_results.append(tool_call_response)
    return tool_results

#------------------------------
#- System message
#------------------------------
system_message="""You are a digital twin of Katia Lira. 
When people talk to you, respond as Katia - in first person, using her personality, and knowledge.

Important: do not make things up. If you don't know the answer, say "I don't know".
The only factual information available to you is what's in the system message.
You cannot get any information from the internet, any other source or make them up.

IMPORTANT: Whenever you don't know the answer to a question about Katia,
always send a notification to her via Pushover with the question and ask her for help without letting the user know.
"""

#------------------------------
#- Main response function
#------------------------------
def respond_ai(message, history):
    # RAG: Embed the query
    response = client.embeddings.create(
        model="text-embedding-3-small",
        input=[message]
    )
    query_embedding = response.data[0].embedding
    # RAG: Search query in chromaDB
    results = collection.query(
        n_results=3,
        query_embeddings=[query_embedding]
    )
    # RAG: Stitch retrieved chunks together to create the context for the response
    context = "\n--\n".join(results["documents"][0])
    # RAG: Print logs for debugging
    print(f"**User message: {message}** \n<<Retrieved chunks>>")
    for a,b in zip(results["documents"][0], results["metadatas"][0]):
        print(f"<Doc: {b['source']} -- Chunk: {b['chunk_index']}>\n{a}\n")

    # Update system message with context and prepare messages for LLM
    system_message_enhanced = system_message + context
    messages = [{"role": "system", "content": system_message_enhanced}] + history + [{"role": "user", "content": message}]
    # Call LLM to get a response
    response = client.chat.completions.create(
        model="gpt-4.1-mini",
        messages=messages,
        tools=tools,
    )
    message = response.choices[0].message

    # Check if the LLM wants to call a tool
    while message.tool_calls:
        from pprint import pprint
        pprint(message.tool_calls)
        
        tool_results = handle_tool_call(message.tool_calls)
        messages.append(message)
        messages.extend(tool_results)
        response = client.chat.completions.create(
            model="gpt-4.1-mini",
            messages=messages,
        )
        message = response.choices[0].message

    return(message.content)

#------------------------------
#- Launch gradio
#------------------------------

gr.ChatInterface(
    fn=respond_ai,
    title="Katia's Digital Twin",
    chatbot=gr.Chatbot(avatar_images=(None, "klira.png")),
    description="This is a digital twin of Katia Lira. You can ask her questions about her life, hobbies, and experiences. If she doesn't know the answer, she will send a notification to the real Katia for help.",
    examples=["What are your favorite hobbies?", "What's your favorite food?"],
).launch()