Sign in to continue learning
Knowledge Graphs for AI Agent: API Discovery
Intermediate
Topics
Agents
Embeddings
Search and Retrieval
Vector Databases
In previous lessons, you constructed a knowledge graph that integrates APIs with additional information about business processes that involve those APIs. In this lesson, you will learn how knowledge graphs can be used for API discovery for AI agents. Let's dive in. Here's what you'll do in this lesson. During design time, you will do initial preparation for retrieval, and at run time, you receive and process user queries. So at design time, you retrieve entity sets, their labels and properties from the knowledge graph. You will then use embedding model to turn retrieved texts into embedding vectors. And at run time, you'll start with a flat index retrieval based on embedding texts. You will then add APIs connected by business process edges. Finally, we enrich the retrieval based on the business process information from the Knowledge Graph. With this background, let's dive into this notebook. As in the previous lessons, let's import all the necessary packages first. We also have some new packages as part of this notebook. For example, we import faiss index classes for performing semantic similarity search on entity sets. We also import tqdm, a utility to show progress bars for loops. We then import OpenAI powered text embeddings via LangChain. And just as in previous lessons, we have a helper function for parameterizing SPARQL queries. All right, let's now set up our data frames to be displayed nicely. Let's begin by setting up a LangChain embedding model using OpenAI's text-embedding-3-large model. Next, we load the knowledge graph we constructed in the previous lessons. In real world business scenarios, we often work with thousands of APIs. Presenting all of these options to a large language model is typically neither efficient due to token limitations and latency, nor is it practical in most cases. Therefore, we need to reduce the API selection space to a small and relevant subset for the agent for a given user query. Doing the flat retrieval reduces the number of APIs to some extent, but poses some problems. Number one, false negatives. These are relevant APIs that are not semantically close enough to the user query. So the retrieval step misses them. By enhancing our retrieval with a second step using business process edges from the knowledge graph. those APIs can still be reached. We complete the business processes that are touched by the true positives with the help of the process information from the knowledge graph. Number two is false positives. Adding the business process information on top of just the flat list of retrieved API endpoints helps the agent to disregard the false positives. Before you construct embedding text for entity sets, let's first recall that for each entity set, you store all its properties in the knowledge graph. The first step is to prepare a SPARQL query to retrieve all the properties for our example entity set related to purchase order. So let's briefly review what we created here. The variable q_api_properties stores the prepared SPARQL query. These define the short prefixes for rdf, xsd, and odata namespaces. Here you select distinct pairs of entity_set and property_label, and in the WHERE clause, we'll bind the entity_set_uri to the purchase order entity set. Further, the triple ensures that the bound URI is explicitly declared as an entity set. And this retrieves the name of the entity set, which is in our case the purchase order. Finally, this triple retrieves the name of the entity set to its entity type, and defines the structure and properties. In the end, for every entity type, you get all the properties and their labels. The first step is to prepare a Sparql query to retrieve all properties for our example entity set related to purchase order. The variable Q underscore API underscore properties stores the prepared Sparql query. Here, you select distinct pairs of entity set and property label. In the first part of the Where clause, you retrieve entity set for purchase order, and its corresponding entity type. In the second part, you retrieve all the properties of that entity type. Let's now execute this prepared Sparql query on our knowledge graph. The results of this query are stored in a dictionary with two keys, entity set and property label. Now let's display the first few rows of this dictionary to see how this looks. As you can see, for the entity set purchase order, we retrieved a few properties and their labels. For example, Address Number, Incoterms, purchasing organization, quotation date, and so on. Next, we prepare a SPARQL query to retrieve and concatenate the name of an entity set and names of its properties. The result of this query will then be used to construct text embeddings. Here the variable q_embedding_string stores the prepared SPARQL query. Here we fetch and concatenate each entity set with the names of its properties. Let's take an example. to see how such an embedding text looks like. Here for example, you'll see for the entity set BillingDocument, you get several properties such as document number, Posting Status, Value Days, Billing Date, and so on. Now, this is just a glimpse of the flexibility that Knowledge Graphs offer. From each entity set, you could go further into the graph and include, for example, more information about each property like its data type, max length, and so on. Or you could include additional business process information that is also stored in the graph. Let's now generate text embeddings for each of the entity sets. Here we loop over the text that our q_embedding_string returned and run them through the text embedding model that we initialized earlier. Okay. Now that we have the embeddings for the entity sets, let's index them for similarity search. We're using faiss to index the embeddings here for simplicity. But really any vector database works. Let's now save the faiss index and the corresponding entity_set_uris to disk with pickle. Let's take a look at what you did so far. We first query the knowledge graph to retrieve a subgraph per entity set, which includes its label and the labels of its properties. The retrieved labels are concatenated and the resulting string is embedded. The embeddings are indexed and associated with the corresponding entity set URI. Now that you've created the index, you can use it to find the nearest neighbors of a given query embedding. So let's define a function to do just that. This function takes a natural language query, converts it into a vector using the embedding model and searches the faiss index for the top k most matching similar entity sets and returns them. Let's try this with an example query. Here, the query we have at hand is the one we started with. To create a purchase order for five pencils in purchasing group 002 and purchasing organization 3000. Let's execute that. You see that a lot of purchase order related entity sets are among the top five matching entity sets, which is what we would expect. We could add an additional reranking step here, but for now we skip that. An important thing to note here is that the retrieval step missed the purchase requisition header, although you see purchase requisition item here. And purchase requisition header is an important prerequisite to create a purchase order. So the agent would not know about it if we stopped here. So let's continue. Let's review what you did so far. You retrieved entity sets, their labels and properties from the knowledge graph. You then used embedding model to turn the retrieved texts into embedding vectors. You created a flat index retrieval based on embedding texts and performed sample queries. Now, you will add APIs connected by Business Process edges. And finally, you will enrich the retrieval based on process information from the knowledge graph. Okay, let's see this in action. We want to get all the process edges that are connected to the entity sets we retrieved earlier. Here, the variable stores the prepared SPARQL query. The select statement here retrieves unique pairs of entity set A, which is the source entity set, entity set B, the target entity set and their names. This query looks for process dependencies using the pr:hasNext property and then performs a union to combine the results from both the directions. Now let's look at what this results in after executing this against our knowledge graph. Let's now define a function to execute SPARQL query on our knowledge graph. This function get_process_dependencies takes a list of entity sets specified by their URIs, executes this SPARQL query on the graph that we loaded earlier to retrieve all the process dependencies from the graph that either lead to or originate from one of the entity sets. Now let's test this on the entity sets that we retrieved for our example query. a query being creating a purchase order for five pencils in purchasing group 002 and purchasing organization 3000. As you can see, we get a list of URIs and names of entity sets for purchase requisition and purchase order. And you can read recollect that the purchase requisition header entity set was missing in the previous step, but now with this step, we have the purchase requisition header also added to the list, This is due to the process information added in the graph. Okay, now that we have all the necessary building blocks, let's put them together and finalize our knowledge graph enhanced API discovery process. In this function, we combine the embedding based retrieval along with the business process dependencies. This takes a user input query and returns the most relevant entity sets along with their business process information. Let's try it out with our same sample query we used earlier. And with that, we get a list of matching entity sets related to purchase order, and most importantly, the entity set related to purchase requisition header, which was missing when we ran only an embedding-based retrieval. This, along with the business process information that PurchaseOrder depends on purchase requisition. Now let's recap what you learned in this lesson. At design time, you ran SPARQL queries on the knowledge graph to turn properties of entity sets into textual representations. From these textual representations, you computed embeddings with an embedding model. At runtime, the discovery starts with a flat index retrieval, which gives us a list of entity sets, but may miss some relevant ones. Then, you add those APIs that are connected to our retrieved APIs via business processes. In our example, we were able to add purchase requisition in this way. And finally, you add a business process information to the discovered entity sets. In the next lesson, you will apply what you just built to discover and execute APIs with AI agents.
Knowledge Graphs for AI Agent: API Discovery