Memory management system with embedding-based retrieval
| 663 | return retriever |
| 664 | |
| 665 | class AgenticMemorySystem: |
| 666 | """Memory management system with embedding-based retrieval""" |
| 667 | def __init__(self, |
| 668 | model_name: str = 'all-MiniLM-L6-v2', |
| 669 | llm_backend: str = "sglang", |
| 670 | llm_model: str = "gpt-4o-mini", |
| 671 | evo_threshold: int = 100, |
| 672 | api_key: Optional[str] = None, |
| 673 | api_base: Optional[str] = None, |
| 674 | sglang_host: str = "http://localhost", |
| 675 | sglang_port: int = 30000): |
| 676 | self.memories = {} # id -> MemoryNote |
| 677 | self.retriever = SimpleEmbeddingRetriever(model_name) |
| 678 | self.llm_controller = LLMController(llm_backend, llm_model, api_key, api_base, sglang_host, sglang_port) |
| 679 | self.evolution_system_prompt = ''' |
| 680 | You are an AI memory evolution agent responsible for managing and evolving a knowledge base. |
| 681 | Analyze the the new memory note according to keywords and context, also with their several nearest neighbors memory. |
| 682 | Make decisions about its evolution. |
| 683 | |
| 684 | The new memory context: |
| 685 | {context} |
| 686 | content: {content} |
| 687 | keywords: {keywords} |
| 688 | |
| 689 | The nearest neighbors memories: |
| 690 | {nearest_neighbors_memories} |
| 691 | |
| 692 | Based on this information, determine: |
| 693 | 1. Should this memory be evolved? Consider its relationships with other memories. |
| 694 | 2. What specific actions should be taken (strengthen, update_neighbor)? |
| 695 | 2.1 If choose to strengthen the connection, which memory should it be connected to? Can you give the updated tags of this memory? |
| 696 | 2.2 If choose to update_neighbor, you can update the context and tags of these memories based on the understanding of these memories. If the context and the tags are not updated, the new context and tags should be the same as the original ones. Generate the new context and tags in the sequential order of the input neighbors. |
| 697 | Tags should be determined by the content of these characteristic of these memories, which can be used to retrieve them later and categorize them. |
| 698 | Note that the length of new_tags_neighborhood must equal the number of input neighbors, and the length of new_context_neighborhood must equal the number of input neighbors. |
| 699 | The number of neighbors is {neighbor_number}. |
| 700 | Return your decision in JSON format with the following structure: |
| 701 | {{ |
| 702 | "should_evolve": True or False, |
| 703 | "actions": ["strengthen", "update_neighbor"], |
| 704 | "suggested_connections": ["neighbor_memory_ids"], |
| 705 | "tags_to_update": ["tag_1",..."tag_n"], |
| 706 | "new_context_neighborhood": ["new context",...,"new context"], |
| 707 | "new_tags_neighborhood": [["tag_1",...,"tag_n"],...["tag_1",...,"tag_n"]], |
| 708 | }} |
| 709 | ''' |
| 710 | self.evo_cnt = 0 |
| 711 | self.evo_threshold = evo_threshold |
| 712 | |
| 713 | def add_note(self, content: str, time: str = None, **kwargs) -> str: |
| 714 | """Add a new memory note""" |
| 715 | note = MemoryNote(content=content, llm_controller=self.llm_controller, timestamp=time, **kwargs) |
| 716 | |
| 717 | # Update retriever with all documents |
| 718 | # all_docs = [m.content for m in self.memories.values()] |
| 719 | evo_label, note = self.process_memory(note) |
| 720 | self.memories[note.id] = note |
| 721 | self.retriever.add_documents(["content:" + note.content + " context:" + note.context + " keywords: " + ", ".join(note.keywords) + " tags: " + ", ".join(note.tags)]) |
| 722 | if evo_label == True: |