Safe RAG for LLMs

Key Takeaways

hi wiin I want to build an app that uses AI with retrieval augmented generation or rag but I have some security concerns what are they well one my AI app should not have full access to my database and two my AI should never see any user IDs let me show you some of my code [Music] welcome to the show wiin what do you do here at Google hello everyone I'm wiin du and I'm a software engineer on the AI and database integration team under Google cloud my work focuses on helping developers to build generative AI products that work with databases Ai and databases huh you're probably quite busy yes companies and developers across all Industries are trying to build generative AI with rag into their software right let's say I want to build a chatbot that can look up a user's information uh like flights I want to build it in Cloud run so it's serverless and Google manages the infrastructure for me I would use retrieval augmented generation or rag for that right yes rag lets the AI query or data sources so it can look up a user's upcoming flights or buy a ticket for the user for example you could ask the chatbot list the tickets I have booked you would want to wanted to respond with something like this the AI itself doesn't know users information but it can look them up in our databases but I've heard that can be dangerous well we don't want the user to be able to ask the llm for another user's flight for example if I'm asking the chatbot listed tickets booked by Martin it should respond with something like this sorry I don't have that information available it should not show me your flights and how can I Implement that because I keep reading about people jailbreaking large language models by asking specially engineered prompts uh it seems hard to defend against that we can prevent Problems by keeping each user isolated each user should have their own AI agent their own tools or accessing the database and their own slice of the database with only data about them so the idea is that if each user has their own slice of the application there will be no cross talk between users uh how can we accomplish that we can do three things first we can add an API on top of the database that way the llm will not have direct access to the data second each tool function should authenticate the user third we can orchestrate the database API and the llm so that the llm can view user IDs I've heard that all software problems can be solved by adding a layer of indirection that's what you're doing here that's right in a r architecture we call that an orchestration layer the app would ask the llm which API to use call the API then ask the llm to dress up the reply in nice readable human language that makes sense uh but the details are still a little fuzzy for me uh could you walk us through an example sure here are the system components first the user logs in this could be done with Firebase authentication Google identity platform or a similar product the user's web browser gets an authentication token back then the user enters their question like list the tickets I have booked the web browser sends that request to the back end in an HTTP request and the web browser also includes the off token in that request that's right now the application back asks the LM question it includes three things the system prompt the list of tools or apis available and the user's questions the question does not include the user's off token or user ID I like that yes that makes it safer the llm would respond with something like this call the tool list ticket the application then sends the list ticket request with the off token as the header to the retrieval service the retrieval service verifies and decodes the off token that comes with the request to get the user ID it inserts that user ID into a predefined SQL query and sends it to the database it then returns one database record for each of the users tickets but that's just data datase records uh the user expects a human readable response right right the application passes the database records of the user's upcoming flies to the llm then the llm creates a nice human readable reply the application returns that reply to the user that looks great WNIN but I have a few questions go ahead Martin in this architecture that you just showed us does the llm ever see any user off tokens or user IDs no it doesn't the AU token and user ID stay in the application and rral service got it can the llm Run queries in the database because I've heard that can be dangerous no the llm can't run queries in the database it only ask that application to call Black boss tool functions that perform authentication and executes predefined queries I like that this is a bit like regular API design even before llms I often built apis on top of my databases it's a good design principle yes llms are new but many of the security patterns we used in the past are still useful every tier of software should have the minimum level of access and needs to do its job ah the principle of least privilege still holds true now uh for me or anybody else who'd like to learn more and get our hands dirty with this how how can we do that whenin my team publish an example app it's written Python and it uses Cloud run you can download it from here you can also read a blog post we wrote about it excellent I will include those links in the video description thank you for showing us this wiin now I know how my AI can access my database safely thanks for having me Martin and thank you everyone for watching if you have any questions for whenin or me please add them to the comments below also let me know if there are any other servess topics you'd like to hear about in future episodes I read every every single comment until next time [Music]