Building memory for AI Agents | Vasilije Markovic, AI By the Bay25
Recording: Building memory for AI Agents | Vasilije Markovic, AI By the Bay25
I'm Vasili. I'm here from Berlin um just for this talk and a few other things but but was a fun fun flight over. Um so um I was jet like two days ago. Now I'm better. So let me tell you about um myself a bit about Cognney uh what we are doing and then then kind of get into the nitty-gritty details and do a small agentic memory demo with Langraph. So first of all about myself um I worked in the modern data stack setup for 10 plus years. I started as a data analyst and data engineer data product manager mostly managing big data at scale. couple of companies in Berlin uh that uh processed and managed large data warehouses, batch streaming
Um I've seen it all. Uh I'm not the the the most amazing developer out there, but I worked with some really great ones and we built uh let's say um large process large output systems that could uh deal and handle uh large volumes of data for millions of customers. Uh and that was in in a few setups. Um as I was doing this I kind of got bored a bit with just moving data from A to B. Then Chad GPT um came out and I realized that um this new um cognitive sciences bachelors I was doing was very interesting when combined with uh let's say the LLMs and uh that we could do many more interesting things. So I started building in my bedroom with a few Python notebooks what became a company. Um and that was a fun uh exercise. it didn't have a name then it got the name graph rag then it became AI memory so I'll kind of tell you a bit about that and how the journey went and and what we ended up building and what do people use it now for so what we have is um >> yeah um oh we'll survive 10% it's 20 minutes [laughter] you know uh I've done worse things with less less time so around two years ago I decide to build a couple of AI assistants And these assistants were supposed to give me recommendations and personalized suggestions on what should I eat, how should I act, what should I do in day-to-day
And I tried using vector stores for that. Uh just use the pure embeddings, personalized, create my own collections for each individual set of things I needed to get done. And I quickly realized I'm not really getting anywhere. Um I would get pizza three days in a week. You know, you've seen all of those uh attempts with with vector stores. And then I realized I would need to build a set of microservices, re-rankers, all kinds of tools that won't really solve the base underlying issue, which is I can't have a piece of infrastructure that I need to infinitely modify just for one single small narrow use case. So as I was playing with that, I was playing with lang chain at the time. Langchain didn't really work that well
It would break with 50 100 megabytes of data. I had no idea what's going on underneath. So the major issue there was I had too much data. So especially as LLMs are very verbose, there was a need to manage all of this data and the assumption is that the agentic systems will produce hundred times more data than humans ever will because they are never shut off. And secondly, um yeah, the the memory consistency. So no way to know what happened um three days ago and then four days ago and five days ago. I couldn't really tie in things. If I mentioned someone few sessions later, things would disappear
So I wanted to solve that. So I started by building a Python SDK um called Cogni um Cogni as for cognitive um I had the inspiration from the cognitive sciences on how do we model memory you know I was reading all the same papers like Koala guys Atkinson shifer in these models long-term short-term memory and we built cogni as a first implementation of such a system uh surprise surprise didn't work so uh that was pretty bad um let's say model to to use but as we built it as we promoted it a bit on Twitter and and and other spaces. People liked it. Um and then we continued iterating. Right now we are being used in telemetry. We see people are creating memory layer 286,000 times uh a month. This month is probably going to be over 500k. So half a million
And that means that people are running LLM generations, vector graph store um inserts uh at least u this this many times. We see around 63 companies using Cognney. We have users from Brazilian oil companies to agricultural companies in Ireland. And most of them are using it because they can effectively make sure that they can customize and create their vertical memory for agents in the way they need it. So it's very pyonic. You have tasks, pipelines, you can build easily things on top and that allows you to really customize how you want to model your data. So we think of ourselves something like DBT for LLMs. And the base problem we are solving is we are replacing this prompt plus rag uh architecture where in effect what you try to do often enough in these cases is prepare your data so LLM can read it and reason on it and produce some type of an output where it fails often enough is it's not general enough and it's not maintained sustainable there is no way to actually update it if something changed you have context rot if you mention the same thing in three places in the same prompts
Uh you need to reconcile how those three things relate to each other and what is the latest state. So if I say I have friend Amy and I say now I hate Amy. We got into a fight. LLM is not going to know which one of these is true. So all of this memory consistency piece is what we really solve underneath. Uh DT hub they just they are data pipelining tool. They build a whole workspace um that uh can create up to 10,000 pages of data sources to data destinations that can allow pipe coders to build data pipelines on the fly. And their junior developer implemented Cogni quite quickly built this and now they can pretty much automate this data source and destination generation
So this is a practical example of how people use this and it's pretty much a layer in between your data warehouse or traditional data stack and LLMs. So the LLMs don't have to really query all of your data and you know you don't need to maintain 17 query types at the same time. So we kind of solve that for you. Um how does it work in actual practice? Um so Cogni has a few um elements. One is the data ingestion piece. We integrate with any type of a data ingestion framework. I'll show you a bit on the architecture. Uh but what happens in the next step is the set of [snorts] entity extraction steps
So typical NLP but with a few more things where we use both the graph and the vector stores to persist the data and maintain the index between them. On top of that is the self-improvement piece. So the second first piece is called cogni.add. It's a around four or five functions in the SDK. Then you have cognit.coy which is the entity extraction permission management um schemas everything else you need to build. Cognit.co is the self-improvement layer. This is what we bring from the neuroscience side. So we effectively allow you to create a self-improvement piece that continues to optimize and improve your memory layer as you interact with it from the agentic side from the side of the source data but also based on the algorithms you want to implement
So let's say you want to create delete all the data that's older than six months. If it wasn't accessed, you can just write a Python function that's going to do that. you don't have to really uh build all the basic underlying infrastructure there. So uh that's what uh the base SDK offers and then as an outcome we get this LLM ready docs for agents in the DT case and the agentic frameworks can read that and they can again feed back into the system. Let's say the doc asana docs were wrong and you can't really uh pull the data back from it. um what would happen is they can provide uh this feedback into the system and with the feedback into the system we can effectively optimize and change the answer next time assign certain weights to the answer weight the system and then uh that's going to self-improve effective. So that's the high level uh case study and then in terms of the architecture the current system allows you to modularly define your data ingestion. So you can use anything from Fiverr, DLT, um unstructured, any of these llama index
So everything you want really to uh parse your data with. We don't really specialize in that. Uh that's the most boring piece if you ask me. Uh and then on top of that, we have ability to ingest from data warehouses. We usually take the schemas plus uh some data. So pretty much the semantic layer from the old days what it meant. And then uh we add that to the Cogniz engine. Now the engine is effectively a Python uh engine that allows us to run set of tasks uh in pipelines and there we uh extract entities, create custom ontologies, deal with the embeddings and then allow you to search by accessing both the graph and the vector store and we have various search methods including chain of thought search which is going to go and try to reason and find its way through the graph representation but also uh ability to do uh pretty uh fancy stuff on the embedding Right
Um, where is our secret sauce here is in the self-improvement layer. Um, and I usually kind of go back to this five-year-old analogy to explain how this works. So if you let's say went with a friend for a dinner let's say five years ago um after five years you're not going to remember what the dinner was about what you ate but you might remember just a skeleton of that memory that it was a nice event where you went to that effectively can allow you to that to kind of modulate your feelings toward your friend next time you see him. So you're going to remember that with Jake you went for that dinner it was really nice. you won't remember the details, but that um prepares you for the next interaction with Jake, which is going to be a positive one by default because you're priming yourself to to be in that state that like, oh, something good happened with Jake, something good will happen in the future. So, this is a very simplified version of how we think about effectively distilling memories from short-term to to long-term memory. And our short-term memory in this context is uh Reddis data store. Uh on top of that after the TTL expires we can move all of these memories into the persistent memory and we can transform them into these skeletons of memories
We do that in what we call um effectively um oh well you can interpret that as a cache where we project skeletons of the graph representations of our memories and then we operate on them and simulate them to create effectively the ability to retrieve the data back. So um the engine is managing all of that and it connects to vector stores, graph stores and models. So this is pretty much the implementation and then aentic frameworks or the chat systems or anything else you have might connect to it and then feed back into the system reassigning these weights changing the system and moving forward. So that's the uh that part of the presentation. I would show you a quick um example with with um langraph um where we add tools that can call cognas memory in langraph and uh that's available uh in our open source repo. So I'll kind of guide you through how we actually do it run some commands hopefully it all works because uh you never know and then uh we would u yeah end it there and leave some time for questions. I'm trying not to scroll up too much because I left my OpenAI key over there. So, you know, I [laughter] Oh Okay, too late
Ah, damn. I'll need to disable this later. So, what happens here is um we are effectively um allowing Langraph now to finally manage state. So, what we allow it to store is session memory and then distill that into long-term memory. So in this case uh we start by building a simple line graph graph which has uh which is effectively just a chatbot and uh the line graph graph is going to have a chatbot and then uh if we want to add tools to it we can easily do it. So let's add tools to lang chain. So now we have we try to add a wiki data API wrapper so we can pull the data from Wikipedia uh define some basic tool nodes and then this tool nodes can be um updated and be plugged into the to langraph. So [clears throat] in this case we are um doing some conditional routting for the tools and then building a tool enabled graph
So that's going to look like something like this. So now chatbot got tools tools can talk to chatbots and there is a start and the end asky because it's nicer and then uh effectively now we can try to stream some graph updates uh where we're trying to find from Wikipedia with these tools something about Alan Turing. So that's going to do its its work and um yeah so it found something about Alan Terring being a British mathematician. Now that's fine. So we can always search Wikipedia that's interesting but now if we add Cognney we can actually store and retrieve this data from semantic memory create graph representations and then uh have this cross-section persistence. So if I mention something in one session, it's going to be there in the next one and then uh do the search in a in a more uh advanced way. So let's do that now. So we are creating a react agent with an ad search and a tool uh search tool
Um and this would be added to um to the system. So let's see. So that should be starting from so we are redoing everything starting from fresh. Um now Cognney we can clean it so we can make sure that there is no data in it. Um that effectively I so I didn't fake anything before this one. Uh so if we clean it the system we can effectively do the data ingestion. So we are now adding some messages and giving some messages to the agent about some contracts with some companies that happened. Um so this would be now adding the data to Cogni
So what you can see in the background is pretty much uh an async uh ingestion of data. We add the pipelines and the extract graph from data is uh adding the data back. We also have ontology support. So what you're seeing no match means that we don't have any structured RDF type ontology that we can merge to this graph. Usually do this to ground the graph and make it more accurate. Um, so once this finishes, we can uh effectively run to the next step. By the way, we made it 10 times faster since this demo was done. So yeah, shouldn't take this long if you play with it
Okay, so the pipeline run completed and uh we can see that uh yeah, the contracts were added to the system. Now we can run um add some data to bypass the agents. So just add some more content to it. So this is an additional operation on top and then we can visualize the graph representations. So it's going to look something like this. So all of these contracts that I was just showing you that we ingested to the agent are going to have some type of contract node with the information about the contract details uh and these would pretty much uh have relationships between them and the other elements of the graph. Here we can see the data, the dates, the star dates, organizations, everything all the entities mentioned. We also have a time ability
So you can um do the time uh management if you enable it and you can do a few other things on top. So let's say we added this and then uh we get this representation and now if we create some react agent [clears throat] that agent will uh effectively be able to u search the data and it's going to find information about these insurance groups we added to Cogni and be able to use the search tool to fetch that whenever it needs it. And then um let's see uh again um how it renders. So in the other screen you're going to see a representation of the session object that uh we added that should be here these blue ones. uh and then effectively yeah the data is now uh sessionized and uh we have information about what happened in which session and how that transpired. We can then also create skeletons of that later on summarize it and and and do operations. So yeah that's that on the demo. Um but yeah there is all the details are in our GitHub repo
If you want to run this yourselves you can. Um just wanted to kind of give you a short um let's say small visualization of how this looks like in real life. Um any questions? I would stop here. Yes. >> Um on the graph, how would you like can you see an ontology there directly? Um or would that just be like one of the parent nodes? >> Uh you can. So there is ability to color the ontologies differently. uh and that you can then see can color it differently and then it's going to be shown in a different way so you know what is deterministic what's the LLM generated then you can merge those so what we have an is an ability to merge is simple fuzzy search but you can also do vector search we have a client now they're actually implementing their own vector search to match the the ontology nodes and then pretty much dduplicate them so we're going to add that to the open source quite Any other questions? >> I'll ask you. >> Um, how is this different from like what Neo4j like offers directly like you can uh they have like basically an LLM that will just ingest and create the entire graph for you
So how is this different? >> So Neoforj is a database. They have a set of tools but none of these tools are able to run in production. So what we see is like people run us max users run us 12,000 times a month on lambdas to process larger volumes of data. This is a data ingestion and modeling tool and we have all these algorithms that we apply on top to make sure that the memory consistency stays. They are a database. They ensure you know transactional consistency and all the fun stuff. We think about memory consistency as a as a problem we are solving. Thank you
Anyone else? Okay, thank you all. Oh, there is someone. >> Hold on. >> Can you hold your question, please? The mic. >> Thank you. >> How do you edit the memory? Like if you want to modify or erase or something. >> So, it's um Pythonic tools. So what you can do is you can define some type of pedantic data objects that are going to be your memory representations and then uh on the ingestion just create a custom pipeline quite quickly and and make your own custom memory that's relatively easy
The search retrieval piece that gets a bit more complicated if you want to do something custom but it's still doable in terms of editing and updating and deleting things that's a task. You can really select the set of nodes you want to delete. You can update it. it's uh it's supporting all the standard types of the opera of the operations. So it shouldn't be very difficult uh and we pretty much change things on the fly all the time when we develop. So so yeah not not too hard and uh one of the guys from the team I think he wrote a rather comprehensive documentation on how to do that. We also have like starter tutorials to build custom pipelines. So should be relatively easy to start
If it's if you get confused or or can't proceed, we are here. We have a discord. I'm always happy to answer questions. So yeah, uh we are definitely trying to find more ways to improve this as as we move along. >> Cool. Um thank you everyone. [applause]