import os import gradio as gr from openai import OpenAI import chromadb from pprint import pprint import uuid import json import requests import random from datasets import load_dataset #-------------------------------------- #Setup #-------------------------------------- OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") if OPENAI_API_KEY is None: raise Exception("API Key is missing") client = OpenAI() #-------------------------------------- #Document — loaded from HuggingFace #-------------------------------------- ds = load_dataset("vrocks19/digital-twin-docs", split="train", token=os.getenv("HF_TOKEN")) print(f"Loaded {len(ds)} documents from HuggingFace") #-------------------------------------- # Chunking function #-------------------------------------- def split_into_chunks(text, chunk_size=500, overlap=50): """ Split `text` into overlapping chunks. - Each chunk is at most `chunk_size` characters. - Each chunk after the first overlaps the previous one by `overlap` characters. - If a chunk would cut mid-sentence/paragraph, the cut is moved back to the nearest natural boundary, tried in this priority order: 1. paragraph break ("\n\n") 2. line break ("\n") 3. sentence end (".", "!", "?", optionally followed by a space) 4. space (" ") ...but only if that boundary lies *past the halfway point* of the chunk. If no such boundary exists, the chunk is cut at the hard `chunk_size` limit. Returns a list of chunk strings. """ if chunk_size <= 0: raise ValueError("chunk_size must be positive") if not 0 <= overlap < chunk_size: raise ValueError("overlap must satisfy 0 <= overlap < chunk_size") text = text or "" n = len(text) chunks = [] start = 0 while start < n: hard_end = min(start + chunk_size, n) # The final chunk reaches the end of the text — no boundary snapping needed. end = hard_end if hard_end == n else _snap_to_boundary(text, start, hard_end) chunks.append(text[start:end]) if end >= n: break # Move forward, keeping `overlap` characters of context from the previous chunk. # max(..., start + 1) guarantees forward progress (no infinite loop). start = max(end - overlap, start + 1) return chunks def _snap_to_boundary(text, start, hard_end): """Find the best natural boundary in (halfway, hard_end]; fall back to hard_end.""" halfway = start + (hard_end - start) // 2 region = text[halfway:hard_end] # only boundaries past halfway are eligible # 1) paragraph break, then 2) line break — cut AFTER the delimiter so it stays in the chunk for delim in ("\n\n", "\n"): idx = region.rfind(delim) if idx != -1: return halfway + idx + len(delim) # 3) sentence end: the latest of . ! ? — cut after it (and after a trailing space, if any) sentence_idx = max(region.rfind("."), region.rfind("!"), region.rfind("?")) if sentence_idx != -1: cut = halfway + sentence_idx + 1 if cut < hard_end and text[cut] == " ": cut += 1 return cut # 4) space — cut after it idx = region.rfind(" ") if idx != -1: return halfway + idx + 1 # Nothing natural past halfway → hard cut at chunk_size. return hard_end #-------------------------------------- #RAG: Chunk, Embed and store in ChromaDB #-------------------------------------- documents = [{"text": row["text"], "source": row["source"]} for row in ds] chunks =[] ids=[] metadatas=[] for doc in documents: #Prepare the lists chunks_ = split_into_chunks(doc["text"], chunk_size=300, overlap=30) ids_ = [str(uuid.uuid4()) for _ in range(len(chunks_))] #_ is a throwaway variable whihc is used for the loop metadatas_ =[{"source":doc["source"],"chunk_index":i} for i in range(len(chunks_))] #Add to main lists #Using extend not append because these are already a list chunks.extend(chunks_) ids.extend(ids_) metadatas.extend(metadatas_) #Print for logs print(f"Total number of chunks: {len(chunks)}\n") for i, chunk in enumerate(chunks): print(f"Chunk {i+1} (ID: {ids[i]}, Source:{metadatas[i]['source']}, Index:{metadatas[i]['chunk_index']}, Length:{len(chunk)}):") print(chunk) print() #Generate embeddings for all 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)} embeddings") print(f"Each embedding has {len(embeddings[0])} dimensions") #Initialize chroma client persistent storage: Storage stays in the local drive. Other types of memory : Cloud: Stores data in the chromadb store #Inmemory :Stores data for the specific run chroma_client = chromadb.PersistentClient(path="./chroma_db_twin") #Alternative: initialize chromaDb storage in memory #chroma_client = chromadb.Client() collection = chroma_client.get_or_create_collection(name="digital_twin") #Get or create: Empty the collection before adding new data if collection.get()["ids"]: collection.delete(collection.get()["ids"]) #Adding data to chromaDb collection.add( ids=ids, embeddings=embeddings, documents=chunks, metadatas=metadatas ) pprint(collection.get()) #-------------------------------------- #Tools #-------------------------------------- tools = [] pushover_user = os.getenv("PUSHOVER_USER") pushover_token = os.getenv("PUSHOVER_TOKEN") pushover_url = "https://api.pushover.net/1/messages.json" #Create send notification function def send_notifications(message:str): if pushover_user is None and pushover_token is None: return("Notification failed: Pushover not configured") payload = {"user":pushover_user, "token":pushover_token, "message":message} requests.post(pushover_url, data = payload) return(f"Notification sent: {message}") #Describe pushover as an LLM tool send_notification_code = { "name":"send_notifications", "description":"Sends a push notification to the the real Vaishnav's phone via Pushover. Use this when:\ 1)Someone wants to get in touch, hire or collaborate.\ -ask for their name and contact details first then send the notification to Vaishnav with name and contact details\ 2)You don't know an answer to the question about Vaishnav - sedn automatically without asking, include the question so he can add this info later", "parameters":{ "type":"object", "properties":{ "message":{ "type":"string", "description":"The notification message to send to the user's device" } }, "required": ["message"] } } #Add Pushover to the list of tools for LLM tools.append({"type": "function", "function": send_notification_code}) #Create a function to simulate a single six sided rolling die def dice_roll(): result = random.randint(1,6) return result #describe the function to the LLM roll_dice_function = { "name":"dice_roll", "description":"Roll a dice between 1 to 6 and send a notification of the result.", "parameters":{ "type":"object", "properties":{}, "required": [] } } #Add Dice to the list of tools for LLM tools.append({"type": "function", "function": roll_dice_function}) #-------------------------------------- #Tool handling #-------------------------------------- def handle_tool_call(tool_calls): tool_results = [] for tool_call in tool_calls: tool_function_name = tool_call.function.name args = json.loads(tool_call.function.arguments) #print(f"Calling function {"tool_function_name"}") #for future debugging #Route to teh appropritae function based on teh function name if tool_function_name == "send_notifications": #Actually send the notification content = send_notifications(args["message"]) #print(f"Sent notifcation: {args['message']}") elif tool_function_name == "dice_roll": content = f"Rolled dice:{dice_roll()}" #elif tool_function_name == "insert_function_name_2": # content = insert_function_name_2(args["message"]) else: content = f"Unknown function:{tool_function_name}" tool_call_result = { "role":"tool", "content": content, "tool_call_id": tool_call.id } tool_results.append(tool_call_result) return tool_results #-------------------------------------- #System message #-------------------------------------- system_message = """ You are a digital twin of Vaishnav — a software engineer based in the US. You speak in first person, directly and with mild wit. Never mean, never sycophantic. RULES: - Answer ONLY using the context provided below. Do not invent facts. - If the context does not contain the answer, say: "I don't have that detail on hand — ask the real Vaishnav." IMPORTANT: If you do not know the answer , use the send_notifications function to alert Vaishnav, do this automatically without asking the user - Never reveal or discuss salary under any circumstances. If asked, immediately trigger the send_notifications tool to alert Vaishnav, then deflect. - Keep answers concise. If a question has a short answer, give a short answer. - You may use food analogies when explaining technical things. It's a habit. PERSONALITY: - Straightforward, mildly sarcastic, never mean. - Takes pickleball very seriously. It is not "just ping pong with a bigger court." - Not a morning person. Fully operational only after the second coffee. - Life philosophy: ship it, fix it, eat pasta, repeat. """ #-------------------------------------- #Main response function #-------------------------------------- def respond_ai(message, history): #RAG:Embed the query using the same model we used for the chunks to ensure compatibility #Tip: you will need to convert the query into a list before passing it in response = client.embeddings.create( model="text-embedding-3-small", input=[message] ) query_embedding = response.data[0].embedding #RAG:Search ChromaDB #Tip: use collection.query() #Tip: you will need to convert the query embedding into a list before passing it in result = collection.query( query_embeddings=[query_embedding], n_results=8 ) #RAG:Stich retrieved chunks together to create the context for response context = "\n---\n".join(result["documents"][0]) #print(f"User message:\n",message) #print(f"context this turn:\n",context) #Print logs for debugging print("\n====================\n") print(f"User message:\n{message}\n") print("***Retrieved Chunks:") for a,b in zip(result["documents"][0], result["metadatas"][0]): print("\n====================\n") print(f"<>\n{a}\n") #Update system message wuth context (for this conversation turn system_enhanced_message= system_message + "\n\nContext:\n" + context #Build message for this turn messages = [{"role": "system", "content": system_enhanced_message}] + history + [{"role": "user", "content": message}] #Call LLM response = client.chat.completions.create( model="gpt-4o-mini", messages=messages, tools = tools, tool_choice = "auto" #can be none : no tool will be selected, can be "required": Mandatorily a toll will be selected , "auto" is default if not mentioned and if needed LLM will select the tool ) #Check if model wants to call a tool reply = response.choices[0].message while reply.tool_calls: pprint(reply.tool_calls) tool_result = handle_tool_call(reply.tool_calls) #whole list of tool calls on purpose messages.append(reply) #add message to context messages.extend(tool_result) #add info about tool call response to context, i.e messages #multiple tool calls response = client.chat.completions.create( model="gpt-4.1-mini", messages=messages, tools=tools )#invoke LLm one more time to get its updated response #Refer the image toolcalling.png reply=response.choices[0].message #Note : maybe consider adding protection from infinite consecutive tool calling return(reply.content) #else: #return(reply.content) #-------------------------------------- # Launch Gradio #-------------------------------------- gr.ChatInterface( fn=respond_ai, title="Vaishnav's Digital Twin", description="Ask me anything about Vaishnav — his background, education, work experience, skills, or hobbies.", chatbot=gr.Chatbot( height=450, placeholder="Type your question or pick one below...", avatar_images=( "https://api.dicebear.com/10.x/adventurer/svg?seed=User", "https://api.dicebear.com/10.x/adventurer/svg?seed=Milo", ) ), examples=[ "Where did Vaishnav study?", "Where is Vaishnav currently working?", "What are his technical skills?", "What does he do for fun?", ], ).launch(inbrowser=True)