Spaces:
Sleeping
Sleeping
| 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"<<Document {b['source']} -- Chunk {b['chunk_index']}>>\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) |