Building Responsible AI Agents on Databricks

Key Takeaways

Hello everyone. Thanks for joining today's session on building responsible AI agents on data bricks. We really appreciate it. It's probably the last talk of today. So thank you for coming uh and joining the session. I'm Pavitra Rao. Uh I'm a delivery solution architect at Databicks. Hi everyone, thanks again for joining. I hope you that you are enjoying your summit and yeah, I'm super excited to be presenting this topic uh with my colleague. My name is Yasin. I am a senior resident solution architect at data bricks and I'm currently based in Singapore. So I just want to begin this talk uh by discussing a truth that many enterprises are grappling with when AI uh systems make financial decisions that cannot be explained. For example, in loan uh approvals, credit scoring claims. It does not just create opacity but it also creates risks. risk of regulatory non-compliance, risk of reputational damage, and also risk of perpetuating bias at scale and without even noticing it. So, here's the opportunity. Responsible AI is not just a compliance checkbox. It should be a an architectural choice. A choice that makes your AI systems not only defensible, but only reliable. So by building responsible AI agents from the ground up, agents that can be defendable, can be auditable and also can be uh like uh responding to your legal requirements or regulatory requirements. We will shift from this notion of being reactive to uh uh reactive damage control to proactive uh governance. So in this in today's session we will show you how we can do all of this on data bricks at scale. The agenda for today is we will start by discussing the key challenges that enterprises are facing when building AI systems and then we will talk about the key pillars of the responsible AI framework. Next, we will show you how dataix intelligence platform and mosaic AI toolkits can help us follow this AI responsible AI framework and build all of these guidelines and recommendations. After that we will discuss uh our financial use case which is credit scoring or credit decisioning and we will show you the architecture design for our compound AI system before moving to a demo to walk you through how to build all of this on data bricks. So what stands in the way of building AI systems for enterprise customers? Let's unpack the five key uh challenges that these enterprise customers are facing today. We can start with bias amplification. So AI systems they don't just reflect bias, they also can amplify this bias. When data is skewed and uh proxy or proxy variables or proxies creep in the system reinforces all the inequities that we have in our data and this is especially dangerous in financial use cases for instance in our credit decisioning use case. The next challenge is system fragility. So most customers that I worked with or many AI systems that we built they perform very well in testing but they will fail or they will break in real life scenarios. So we can think of model drift, edge cases and also adversarial uh inputs. Without guard rails and uh deep stress testing of these AI systems, all the system will fail uh silently and and predictably. Next we have the lack of transparency. If I cannot explain the decision by made by my model, I will not be able to defend it and validate it with the stakeholders. And this is not just a technical issue, it's also a business and a regulatory liability for most customers and most organizations. So without documentation, trustability lineage with time we will lose the trust in our AI system. Then we have data privacy and security risks. We can talk about inference attacks, overexposure of sensitive fields or even model training on confidential data. These all can lead to major compliance issues and problems for these organizations. And responsible AI starts with secure and privacy aware data practices. The last challenge that we can talk about is governance challenges. So the question here is who is accountable for what? Without clear ownership, audit trails and coordinated deployment processes, responsible AI becomes an afterthought instead of being a built-in foundation in our architecture. So to address all these challenges that we just discussed, we need to build a framework. We need to think of a framework, not just ad hoc fixes along our use case or our along our deployment of this use case. And a responsible AI framework offers the structure to build safe, trustworthy and ethical systems by design. So for this framework what we will do is we will focus on five key pillars starting with uh fairness and inclusiveness and this is for me arguably the most urgent pillar. First aspects can be bias detection and monitoring. Here we need to build in fairness metrics like uh demographic parity, equal opportunity or desperate uh impact analysis. And we need to monitor all of these KPIs along the way end to end in our uh compound AI systems explainable model decisions. So if users as I mentioned earlier do not understand why a decision has been made by a model or by an AI system they cannot contest it and they cannot validate it and also the regulators will not be able to validate it and approve it. So we can use tools like shape or lime in order to surface the model decision and explain in detail why certain decision have been made. diverse data sourcing. This point is very important as well because bias often starts at the source of our AI systems, right? And with under represented populations or proxies baked into our features, we will have this uh bias reflected in throughout the whole uh system. So actively audit and augment data sets using for example data bricks marketplace to ensure representativeness across ra gender age and all of the sensitive dimensions in your AI system. And now moving to the next pillar. Once fairness is addressed, the question that we ask now, can we trust the model to perform consistently in all environments? So enhanced AI effectiveness will help us make sure that the model is solving the right problem with well- definfined objective and well-defined KPIs and being able to monitor all these KPIs and objective throughout our engagement consistent model reliability. So as I mentioned the model should be able to generalize across different environments data drifts and also edge cases. So we need to make sure to implement robust evaluation pipelines like cross validation stress testing adversal examples and also make sure to simulate edge cases very early in the process. robust AI safety measures by leveraging safeguards like outlier detection, confidence threshold and fallback strategies or things that we discussed in the keynote or other sessions like kill switches, secret breaks where we have human in the loop processes. Next is transparency. So once we build systems that are fair and reliable, the next question is can we understand how they work. So we need to have model explained properly and go beyond accuracy and ask why the model made certain decisions. Right? And here again as I mentioned we can use tools like uh shape or lime in order to explain the full decision and have full tracibility of how the Asians are communicating or interacting with each other. And then we have unified governance and data lineage. This is where we will benefit from unity catalog and uh unity catalog lineage to have full trustability of how the data flows through the system from ingestion all the way to deployment or all the way to inference. So the the idea here is that with great model power comes great responsibility, right? Especially with when it involves user and customer data. And this pillar is about protecting customer data and preserving institutional trust through the whole process and through privacy and security. So we have data privacy controls. This is the the idea here is to use only the data that we need and make sure to leverage things like data masking, encryption, row level masking, column level filtering with Unity catalog to make sure that we are not using any sensitive data in our AI systems. And I also highly encourage you to benefit from or to read through the databicks AI security framework that helps you identify the different risks that we discussed and discovered with customers and it gives you recommendations to tackle each of these risks. So the last pillar is accountability and for this last pillar it's most often often the most overlooked uh pillar because when the most ethical private and robust system means little if no it will mean little if no one is clearly responsible for what it does. Right? The idea here is that every decision, every data set and model version should be traceable and also we should have like audit trails, documentation and things like model cards, data versioning to make sure we understand the whole flow of data and decisions. And the last one that I want to discuss for this framework is how we know who is responsible for what. Right? So we need to assign owners, we need to define escalation paths, we need to automate reporting and ensure regular reviews of the model behavior over time. Now account accountability I consider it as the glue that binds all the other pillars together. It ensures when things go wrong in any of the other pillar, someone knows it, someone owns it and there is a process to fix any issue that has been identified. That's how we make responsible AI not just a principle but a a practice that we need to follow when building each uh AI system. Awesome. Now that we have um understood what we need to achieve, let's talk about how to implement this REI pillars in real system. Building responsible AI isn't just a one-step trick, right? We need to have end toend support across the AI life cycle. and data bricks data intelligence platform along with mosaic AI actually provides exactly what we need and it is uh uh having built-in design for that it offers an unified platform where all the components like data models evaluation deployment and monitoring all of this come together the EI developers can use LLM models data retrievalss and even um traditional ML models and all of these components can be put together in a single compound AI workflow and this is very crucial for our uh RAI framework. Um I refer to RAI uh just for ease. its responsible EI framework because instead of relying on one opaque isolated models um for frame uh the implementing RAI here we can compose multiple components with clear responsibilities. So let's start with the build phase of our responsible AI agents. Right? The first step is always about prepare preparing uh the foundation with no surprises. uh it's unity catalog where we ingest uh and catalog data ensuring every data whether it is structured unstructured uh this is all uh stored there for traceability and governance. Then we uh define tools, right? Tools are um callable discrete uh discrete functions um just like your credit decision model or like a simple SQL uh function uh shop explainer etc. These are your tools and agents orchestrate this tools and all the agents tools. Um and the data is stored in Unity catalog for lineage and fine grain access control um with some of the features that Yasheen uh was discussing. There are use cases where you would need um unstructured data to be indexed into vector search and this can also be stored in unity catalog or in delta table and served through our vector search index. So your agents can um extract relevant data, the necessary data that it needs uh to make the um responses more um usable I would say right uh context aware. We can build uh further uh the AI agents using the AI uh playground. We will quickly see all of this in the demo. Uh but the AI playground lets us simulate agent workflow in real time and it comes with built-in LLM judges which is so crucial even at the prototyping level you are trying to assess the correctness the safety and the relevance of the responses. Unity catalog uh features um like column masking can help um mask your sensitive data and uh make it accessible to the right audience. uh you can scope the um access to it and further ML flow right that's the next biggest thing that we consider is so critical in the AI uh framework is ML flow uh experiments the tracking of experiments the runs model versions artifacts and parameters providing that full lineage and reproducibility which is so crucial for our uh RAI framework. We further in integrate SHAP and lime. Uh I'm not sure uh how many of you have heard some of these functions but these are like local explanations which can run on the model uh predictions and help us understand uh and every make every prediction explainable. Why exactly the model predicted that particular value is the explanability that we need the transparency and these explanations it's not just for the technical teams to make uh decisions it's all also for the power users uh which can be um like downstreamed through component um like front-end components like EIBI dashboards or your own um dashboarding tools we can leverage there once we have um some of the prototyping done let's understand how to evaluate right like this is such a key cornerstone for REI framework we begin with creating synthetic evaluation data sets um to simulate edge cases or there are scenarios that are sensitive to fairness uh and uh other rare events so this is all uh possible with Mosaic AI agent evaluation framework. Uh some of the announcements came today as well. Uh we have built-in modules to help create this um evaluation set which is representing your use case. It's representable to your uh use case or the application you're building. Why do we do this, right? like why don't I just use uh my subject matter experts or wait until I get the real data that that's not practical right there's so limited access to yourmemes and they have other jobs to do so it becomes more costly and tedious versus using this approach you're also not introducing any bias human bias to the evaluation set and um all all of this uh can help assess the agent performance across different conditions. Um and as we discussed some of the dependencies will also reduce. Agent evaluation uh Mosaic AI agent evaluation also has an user interface which provides a sidebyside view of all the metrics. So your subject matter experts can go in um or even you as a engineer can assess the metrics, latencies, costs and outputs um which further facilitates in debugging the quality issues of your agentic framework uh the framework right the uh agent itself. Further we have seen lot of um customers use uh marketplace a lot especially regulated industries to bring in external data sources and um use uh clean room right this is so uh critical in financial services or even regulated industries that you want to use um the data uh and these industries have u data sensitivity flags always high but you want to partner and improve the model uh model robustness. So through clean room this is achievable without ever exposing the raw data. So datab bricks platform doesn't just help us build smarter agents but as we see here it helps us build more inclusive and informed ones through external data and um making it very robust uh for our models uh to improve on external data sets. Now that we've explored um the preparation of the data, the prototyping of agents and evaluation, let's talk a little bit about iterative improvement, right? Like it's not like set and forget. You have to have that iteration over improving this. So MLflow traces can be used for agent um like evaluating agent behavior. We'll quickly see this in the demo all the traceability uh steps it goes through and we can also have observability hooks uh embedded at every stage like if an output looks uh off from the actual uh values that you have default set it will trace back and provide um alerts and also give you that whole traceability uh model versioning and understand what tools were used um through the system, the Unity catalog, the entire uh prototyping system that we developed. Further, um the code that was used in developing each version of this prototype is also logged. So every call of the model and the data is retrievable and fully traceable and you can also reproduce it. Why do we do all this? It finally enables us with efficient root cause analysis and you can further fix the uh gaps or the outliers that you see there. So for example, let's imagine if a bias is um detected in the outcome of a certain group age group of your model, right? We are talking about credit decisioning models. we can check whether it is originating from your feature store or the model predictions or even uh prompts uh that you're using uh through the entire system um the lineage traceability that we have. So we believe that this closed loop debugging is what makes RAI systems sustainable at scale. Even after deployment, a continuous monitoring plays a very crucial role here. We have lakehouse monitoring for generative AI applications. This also came out as a new announcement this week um where we have uh detections um that we can do for data drifts, model drifts and shifts in fairness metrics. Uh it uses the LLM judges, built-in L&M judges and custom metrics for offline evaluation and online monitoring. And yes, we can set automated alerts. If a fairness threshold is um breached or the agent performance is going down, the system can always flag. And optionally also you can have the agents pause making decisions and go back to a default system. So your front-end applications are running as is. Finally, let's talk about going into production, right? Uh pre-production step is very very crucial. Uh we can test the agent end to end integrating all the tools and models like online feature stores uh vector stores or even genie rooms that you would use in the agentic framework and through the review app uh that comes with the agentic framework um domain experts can give realtime feedback to the systems. they can assess the model uh responses and flag it and uh also give expected behavior through natural language. Um this human in the loop uh review is where RAI becomes real and loan officers can approve, flag or reject the decisions and all of this can be stored and provided as a feedback to your agent. Um all the model serving endpoints that we have comes with built-in safety filters and REI guardrails which can be applied not at just at the model level but the whole orchestration level um so that your models are not interacting with certain types of content like PII data sensitive data and uh the final note here is um the EI red teaming I don't know how many of you have heard heard this concept of AI red teaming. Uh this is mostly like systematically attacking your own model and um do some diver diverse test in one of um like for safety issues like detecting safety issues early on and this has been one of the proven methods even in internally with data bricks. Our engineer engineering team uh does this very often. All right. So let's see how we can actually implement all of this all the concepts that we discussed on responsible AI framework. And I will quickly go through the design and architecture that we will be uh doing during our demo. And for this demo we thought of the credit decisioning financial uh use case and we are using a supervisor uh multi- aent architecture uh to implement this uh credit decisioning chatbot. So at the beginning of the first agent is what we call the data retrieval agents. What it does it will ret it will interact with Unity catalog with Delta Lake with your feature store table in order to retrieve the applicants data some credit bureau data and some financial data in this concept uh uh context and then it will send all of this data to the next agent that we call the scoring agent. What this agent does, it will load a credit decision in ML model that we already trained and stored in Unity catalog and MLflow and it will give it all the features, all the data that we prepared in the first agent in order to compute the credit uh score for this applicant. And once we do this, we move to the next agent which is the explanability agent that will help us identify the key factors that that made the agent give us that credit score. Right? So it explains what are the key factors but it also monitors the bias and fairness of this score or of this decision. And of course the last one or the last agent that we also want to focus on is the response generation agent because we have the score but we want to have a full ex uh explanation and detailed explanation of why did this uh all these agent made this uh decision and gave this score or this final decision for this applicant. Awesome. So this is the credit decisioning advisor chatbot. Uh this is built on data bricks apps. What this chatbot um does is simply takes in our user query to provide explanation on whether um this customer will default the loan payment and not only that provide an explanation as to like why exactly this decision was made. So let's dive right in here. uh initially it will go ahead and provide us some risk score saying hey there is a high probability that this customer would not uh pay the loan but as a loan officer just the risk score of one and a probability value doesn't help right it's not actionable or not even providing u the right explanation as to why this is so with our agentic framework that we built using RAI as a frame um ground there or the pillars. Uh it goes ahead further and explains what are the top features that are driving this insight. For example, transaction activity patterns um employment stability and then account balance management. Um by the way this model was given around like 10 to 15 uh features final features and on on based on the shar values it is saying these are the five values that it thinks uh is impacting the uh score of one. uh it also further uh gives us a more comprehensive business implications and what the loan officer should consider as a next step to do here um in terms of uh the uh decisions and um actions right and these things are something we can also provide in one of our uh agents or like through the prompts uh prompts are really uh powerful and we saw today the announcement you can also track the performance of a prompt how well your prompt is doing through our prompt registry. So um overall you will see that with the fairness note so it's biased um it's guarded against biases and one thing if you're noticing the results here we never display any sensitive information like age the balance or um any other tenure. So we are making sure the output is also um protected with the uh with not showing the sensitive data but actionable enough for the loan officer to take next actions. Let's switch the gear and understand uh how this was built on um uh data bricks. We start with unity catalog. Uh Unity catalog is where all your data the models um functions uh uh raw data sets everything is stored and not to um like really miss this it it provides the complete lineage um to understand the traceability. So if a prediction is made we can understand uh which model was used what model version what code version all of this can be tracked upstream and downstream. And um the traceability is also built in uh to this. Once we have all these assets uh kind of stored here, governed uh permissions granted, we can switch quickly to our prototyping environment which is uh the AI playground. In the AI playground, uh we start with using our simple LLM models. Um u recently cloudset has been the best in reasoning. So we have ended up using the same for our uh demo as well. And the tools is where uh it's a gamechanging right you can start adding tools and other agents here vector search and then use prompts to um start prototyping your agent and one thing I want to highlight is the evaluation like we discussed it's a cornerstone for our uh framework. So we will see here AI judges ma uh giving um overall correctness of the response for your agent. So you start doing this at the prototype level and it's not an afterthought. So once we have this uh built uh we can quickly generate automatically generate an agent code. So you don't have to be an expert in langraph or langu llama index uh some of these frameworks. You can quickly build uh this end to end code and this comes with um major steps like uh evaluation uh which we discussed through the evaluation data set and then uh you can register the model to unity catalog deploy the agent to our model serving endpoint but just to uh focus a little bit on the MLflow experiment. This is where um we can implement all of most of our responsible AI framework and the pillars the concepts we discussed. So when I go into this experimentation run uh you can see that every uh experiment is being logged for us and all the model versions are being logged here and the code and you can also trace it back to Unity catalog and that kind of um closes the loop for us. And then further to look at the traces, right? We discussed this in the slides as well. Uh for every prompt or every question that we send, we have complete detail about how exactly the agent went about giving the response that we have and the assessments. Now if I'm a subject matter expert and an ML engineer, I can just start using this framework um UI to say hey add a new assessment or uh other metrics makes it very very easy and simple for traceability and um lineage and audit right so we have this whole uh overall correctness groundedness relevance for this and we can set up monitoring um and tell hey I want to add more guidelines or like um use the default uh judges and uh coming back to um the evaluations as well we have every data set logged and I won't go further into the details but this is where our ML flow 3 comes into picture saying hey there is a feedback loop uh for every new responses that comes your um subject matter experts can go through the UI the review app most of you would have seen a review app, agent review app, and it's all tied back with the feedback model. Uh, and you can further use this for a next iteration of agent performance, right? Improve the performance um with real-time data. Finally, um let's say I'm I don't want to use a chatbot versus me being a uh loan officer, I want to understand a customer portfolio. How is it how is all these factors um going to imply on my set of customers? What are the business opportunities? What are the risk profiles? I can always publish this into a dashboard. And this is an interactive dashboard which where I can ask questions in uh real time um uh uh natural langu through natural language and I can start getting insights right out of my dashboard with um AI and BI genie uh feature here. So just to conclude this uh uh in the next two to three minutes here we we always want to talk about like what are the best practices right like how do we consolidate everything we shared into an operational um implementation so without having to say always ground the LLM in in your data uh sorry ground the LLM in data and explanations from LLM should reflect the actual model decision and the SHAP contributions that we saw uh the factors that were uh implying or impacting the loan decision. Um apply domain constraints and policy prompts. Right? For example, if you are building an EI agent for loan communication, we have to make sure the domain uh constraints uh or the regulatory templates. So, it knows the language and also uh phrases the denials according to the regulatory uh templates that we provide. Uh next thing is guardrails are mandatory. There's um this goes beyond simple filtering, right? They include behavioral constraints and even routing logic. So guardrails are inevitable. H just like human in the loop is non-negotiable. Right? We we saw how critical it is to have some of the reviews being um reviewed by your um loan officers so your agent is learning from it and picking up the right behavior. And um the loan officers are not just like overriding the AI. I think this is um key to understand they co-pilot with it which brings a next level of uh transparency and confidence in the in the system and um this is not set and forget. uh it requires continuous monitoring uh MLflow traces agent evaluation lakehouse monitoring uh for geni applications we need to track every decision and improve them uh without tracking you can't improve without measuring you can't improve so it's very very important and last but not the least um auditing for data lineage uh access logs should be built into the foundation. So every interaction, every output, every override is logged and uh traceable for further review as well. So just to conclude this what we have shown today is that responsible AI doesn't have to be an afterthought with um the datab bricks platform and mosaic AI we have a unified platform that builds trust into the AI development life cycle. All right. So I think just in the interest of time we'll finish with this like uh for all the AI systems that you are working with on and after all the features that you've seen today the concept that you should keep in mind is build once govern always and scale responsibly by following the guidelines that we discussed today or the framework that uh that's published by data bricks and of course take some time to complete the survey and share your feedback. Thanks everyone for your time. I know it's a packed content so feel free to reach out to us on LinkedIn or ask any questions uh if you catch us outside. Thank you everyone. Thank you everyone.