Memory manager

pyproject.toml

@@ -23,7 +23,8 @@ solana = "^0.35.0"
 aiohttp = "^3.11.11"
 requests = "2.32.3"
 jupiter-python-sdk = "^0.0.2.0"
-
+chromadb = "^0.6.3"
+pypdf = "^5.1.0"
 
 [build-system]
 requires = ["poetry-core"]

src/agent.py

@@ -6,6 +6,7 @@
 from pathlib import Path
 from dotenv import load_dotenv
 from src.connection_manager import ConnectionManager
+from src.memory.manager import MemoryManager
 from src.helpers import print_h_bar
 from src.action_handler import execute_action
 import src.actions.twitter_actions  
@@ -29,6 +30,9 @@ def __init__(
             missing_fields = [field for field in REQUIRED_FIELDS if field not in agent_dict]
             if missing_fields:
                 raise KeyError(f"Missing required fields: {', '.join(missing_fields)}")
+
+            # Initialize memory manager
+            self.memory = MemoryManager(agent_name=agent_name)
 
             self.name = agent_dict["name"]
             self.bio = agent_dict["bio"]
@@ -109,7 +113,7 @@ def _construct_system_prompt(self) -> str:
                         )
                         if tweets:
                             prompt_parts.extend(f"- {tweet['text']}" for tweet in tweets)
-
+            
             self._system_prompt = "\n".join(prompt_parts)
 
         return self._system_prompt

src/cli.py

@@ -152,6 +152,56 @@ def _initialize_commands(self) -> None:
                 aliases=['talk']
             )
         )
+
+        # Memory commands
+        self._register_command(
+            Command(
+                name="memory-upload",
+                description="Upload one or more documents to the agent's memory",
+                tips=["Format: memory-upload {category} file1 [file2 file3 ...]",
+                    "All files will be stored in the same category",
+                    "You can use wildcards like notes/*.txt to upload all text files in the notes folder"],
+                handler=self.memory_upload,
+                aliases=['upload-memory']
+            )
+        )
+
+        self._register_command(
+            Command(
+                name="memory-list",
+                description="List memory categories or documents within a category",
+                tips=["Format: memory-list [category]",
+                    "Without category: shows all categories and their sizes",
+                    "With category: shows documents in that category"],
+                handler=self.memory_list,
+                aliases=['list-memories']
+            )
+        )
+
+        self._register_command(
+            Command(
+                name="memory-search",
+                description="Search all memories or within a specific category",
+                tips=["Format: memory-search 'search terms' [category]",
+                    "Example: memory-search 'blockchain basics'",
+                    "Example: memory-search 'smart contracts' solana"],
+                handler=self.memory_search,
+                aliases=['search-memory']
+            )
+        )
+
+        self._register_command(
+            Command(
+                name="memory-wipe",
+                description="Delete memories at different levels (all, category, or document)",
+                tips=["Format: memory-wipe [category] [filename]",
+                    "Without arguments: wipes all memories",
+                    "With category: wipes entire category",
+                    "With category and filename: wipes specific document"],
+                handler=self.memory_wipe,
+                aliases=['wipe-memory']
+            )
+        )
 
         ################## CONNECTIONS ################## 
         # List actions command
@@ -526,13 +576,205 @@ def chat_session(self, input_list: List[str]) -> None:
                 if user_input.lower() == 'exit':
                     break
 
-                response = self.agent.prompt_llm(user_input)
+                memory_context = ""
+                # Only search memories for meaningful queries longer than 3 words
+                if len(user_input.split()) > 3:
+                    logger.info("\n🔍 Searching memories...")
+                    memory_context, results = self.agent.memory.get_relevant_context(user_input)
+
+                    if results:
+                        logger.info("Found relevant memories:")
+                        for i, result in enumerate(results, 1):
+                            logger.info(f"\nMemory {i} (similarity: {result.similarity_score:.2f}):")
+                            logger.info(f"Source: {result.memory.metadata.get('source', 'Unknown')}")
+                            logger.info(f"Preview: {result.memory.content[:200]}...")
+                            memory_context += f"From {result.memory.metadata.get('source', 'reference')}:\n{result.memory.content}\n\n"
+                        logger.info("\nUsing these memories for context...")
+
+                enriched_prompt = (
+                    f"{user_input}\n\n"
+                    f"{'Knowledge to draw from:\n' + memory_context if memory_context else ''}"
+                    f"Use your personality and style to respond, incorporating any relevant knowledge naturally."
+                )
+
+                response = self.agent.prompt_llm(enriched_prompt)
                 logger.info(f"\n{self.agent.name}: {response}")
                 print_h_bar()
-
+                    
             except KeyboardInterrupt:
                 break
 
+    def memory_upload(self, input_list: List[str]) -> None:
+        """Handle document upload to memory"""
+        if not self.agent:
+            logger.info("No agent loaded. Use 'load-agent' first.")
+            return
+
+        if len(input_list) < 3:
+            logger.info("Please specify a category and at least one file.")
+            logger.info("Format: memory-upload {category} file1 [file2 file3 ...]")
+            return
+
+        category = input_list[1]
+        filepaths = input_list[2:]
+
+        # Expand wildcards
+        import glob
+        expanded_paths = []
+        for filepath in filepaths:
+            expanded = glob.glob(filepath)
+            if expanded:
+                expanded_paths.extend(expanded)
+            else:
+                expanded_paths.append(filepath)
+
+        if not expanded_paths:
+            logger.info("No matching files found.")
+            return
+
+        stats = self.agent.memory.upload_documents(expanded_paths, category)
+
+        logger.info("\nUpload Summary:")
+        logger.info(f"Total files attempted: {stats['total_attempted']}")
+        logger.info(f"Successfully processed: {stats['successful']}")
+        logger.info(f"Failed: {stats['failed']}")
+        logger.info(f"Total chunks created: {stats['total_chunks']}")
+
+    def memory_list(self, input_list: List[str]) -> None:
+        """List memory categories or documents in a category"""
+        if not self.agent:
+            logger.info("No agent loaded. Use 'load-agent' first.")
+            return
+
+        # If no category specified, list all categories
+        if len(input_list) < 2:
+            categories = self.agent.memory.list_categories()
+            if not categories:
+                logger.info("No memory categories found.")
+                return
+
+            print_h_bar()
+            for category in sorted(categories):
+                stats = self.agent.memory.get_category_stats(category)
+                logger.info(f"\nCategory: {category}")
+                logger.info(f"Documents: {stats['document_count']}")
+                logger.info(f"Total chunks: {stats['total_chunks']}")
+            print_h_bar()
+            return
+
+        # Get stats for specified category
+        category = input_list[1]
+        try:
+            stats = self.agent.memory.get_category_stats(category)
+
+            logger.info(f"\nDocuments in category '{category}':")
+            print_h_bar()
+
+            for doc in stats["documents"]:
+                logger.info(f"\nFile: {doc['filename']}")
+                logger.info(f"Chunks: {doc['chunk_count']}")
+                logger.info(f"Total size: {doc['total_size']:,} characters")
+                logger.info(f"Upload date: {doc['upload_date']}")
+
+            print_h_bar()
+
+        except Exception as e:
+            logger.error(f"Error listing memories: {e}")
+
+    def memory_search(self, input_list: List[str]) -> None:
+        """Search memories across all or specific categories"""
+        if not self.agent:
+            logger.info("No agent loaded. Use 'load-agent' first.")
+            return
+
+        if len(input_list) < 2:
+            logger.info("Please specify a search query.")
+            logger.info("Format: memory-search 'search terms' [category]")
+            return
+
+        # Check if final argument is a valid category
+        potential_category = input_list[-1]
+        available_categories = self.agent.memory.list_categories()
+        category = potential_category if potential_category in available_categories else None
+
+        # Get query
+        query_parts = input_list[1:-1] if category else input_list[1:]
+        query = ' '.join(query_parts).strip("'\"")
+
+        results = self.agent.memory.search(query=query, category=category)
+
+        if not results:
+            if category:
+                logger.info(f"No results found in category '{category}' for '{query}'")
+            else:
+                logger.info(f"No results found for '{query}'")
+            return
+
+        logger.info(f"\nSearch results for '{query}':")
+        print_h_bar()
+
+        for i, result in enumerate(results, 1):
+            memory = result.memory
+            similarity = result.similarity_score
+
+            logger.info(f"\n{i}. Similarity: {similarity:.2f}")
+            logger.info(f"Category: {memory.category}")
+            logger.info(f"From: {memory.metadata.get('original_filename', 'Unknown source')}")
+            logger.info(f"Content: {memory.content[:200]}...")
+
+        print_h_bar()
+
+    def memory_wipe(self, input_list: List[str]) -> None:
+        """Wipe memories at different levels"""
+        if not self.agent:
+            logger.info("No agent loaded. Use 'load-agent' first.")
+            return
+
+        # Wipe all memories
+        if len(input_list) == 1:
+            categories = self.agent.memory.list_categories()
+            if not categories:
+                logger.info("No memories to wipe.")
+                return
+
+            logger.info("\n⚠️  WARNING: This will delete ALL memories for this agent!")
+            logger.info(f"Categories to be wiped: {', '.join(categories)}")
+
+            if self.session.prompt("\nType 'yes' to confirm: ").strip().lower() != 'yes':
+                logger.info("Operation cancelled.")
+                return
+
+            if self.agent.memory.wipe_all_memories():
+                logger.info("✅ All memories wiped successfully.")
+            return
+
+        # Wipe one specific category
+        category = input_list[1]
+        if len(input_list) == 2:
+            if category not in self.agent.memory.list_categories():
+                logger.info(f"Category '{category}' not found.")
+                return
+
+            logger.info(f"\n⚠️  WARNING: This will delete category '{category}'")
+
+            if self.session.prompt("\nType 'yes' to confirm: ").strip().lower() != 'yes':
+                logger.info("Operation cancelled.")
+                return
+
+            result = self.agent.memory.wipe_category(category)
+            if result["success"]:
+                logger.info(f"✅ Category '{category}' wiped successfully")
+            return
+
+        # Wipe one specific document
+        filename = input_list[2]
+        chunks_deleted = self.agent.memory.wipe_document(category, filename)
+
+        if chunks_deleted == 0:
+            logger.info(f"No document found matching '{filename}' in category '{category}'")
+        else:
+            logger.info(f"✅ Document '{filename}' wiped successfully ({chunks_deleted} chunks deleted)")
+
     def exit(self, input_list: List[str]) -> None:
         """Exit the CLI gracefully"""
         logger.info("\nGoodbye! 👋")