From 7378668d814e9617e9356c4cfe5aaed8c4f2b107 Mon Sep 17 00:00:00 2001 From: harshalmore31 Date: Wed, 18 Jun 2025 22:26:23 +0530 Subject: [PATCH 1/2] Add Simple RAG example and enhance Agent with RAG functionality - Introduced a new example demonstrating RAG integration with Swarms framework. - Added SimpleMemoryStore class for in-memory storage and querying. - Enhanced Agent class with RAG configuration and methods for RAG operations. - Updated RAGConfig to include default relevance keywords for insights extraction. --- simple_rag_example.py | 156 +++++++++++++++++ swarms/structs/__init__.py | 2 + swarms/structs/agent.py | 254 ++++++++++++++++++++++------ swarms/structs/agent_rag_handler.py | 11 +- 4 files changed, 366 insertions(+), 57 deletions(-) create mode 100644 simple_rag_example.py diff --git a/simple_rag_example.py b/simple_rag_example.py new file mode 100644 index 00000000..1c23822c --- /dev/null +++ b/simple_rag_example.py @@ -0,0 +1,156 @@ +""" +Simple RAG Example with Swarms Framework + +A concise example showing how to use the RAG integration with Swarms Agent. +This example demonstrates the core RAG functionality in a simple, easy-to-understand way. +""" + +import time +from swarms.structs import Agent, RAGConfig + + +class SimpleMemoryStore: + """Simple in-memory memory store for demonstration""" + + def __init__(self): + self.memories = [] + + def add(self, content: str, metadata: dict = None) -> bool: + """Add content to memory""" + self.memories.append({ + 'content': content, + 'metadata': metadata or {}, + 'timestamp': time.time() + }) + return True + + def query(self, query: str, top_k: int = 3, similarity_threshold: float = 0.5) -> list: + """Simple keyword-based query""" + query_lower = query.lower() + results = [] + + for memory in self.memories: + content_lower = memory['content'].lower() + # Simple relevance score + relevance = sum(1 for word in query_lower.split() if word in content_lower) + relevance = min(relevance / len(query_lower.split()), 1.0) + + if relevance >= similarity_threshold: + results.append({ + 'content': memory['content'], + 'score': relevance, + 'metadata': memory['metadata'] + }) + + return sorted(results, key=lambda x: x['score'], reverse=True)[:top_k] + + +def main(): + """Main example demonstrating RAG functionality""" + print("šŸš€ Simple RAG Example with Swarms Framework") + print("=" * 50) + + # 1. Initialize memory store + print("\n1. Setting up memory store...") + memory_store = SimpleMemoryStore() + + # Add some knowledge to memory + knowledge_items = [ + "Python is a versatile programming language used for web development, data science, and AI.", + "Machine learning models learn patterns from data to make predictions.", + "The Swarms framework enables building sophisticated multi-agent systems.", + "RAG (Retrieval-Augmented Generation) enhances AI responses with external knowledge.", + "Vector databases store embeddings for efficient similarity search." + ] + + for item in knowledge_items: + memory_store.add(item, {'source': 'knowledge_base'}) + + print(f"āœ… Added {len(knowledge_items)} knowledge items to memory") + + # 2. Configure RAG + print("\n2. Configuring RAG...") + rag_config = RAGConfig( + similarity_threshold=0.3, # Lower threshold for demo + max_results=2, + auto_save_to_memory=True, + query_every_loop=False, # Disable to avoid issues + enable_conversation_summaries=True + ) + + # 3. Create agent with RAG - using built-in model handling + agent = Agent( + model_name="gpt-4o-mini", # Direct model specification + temperature=0.7, + max_tokens=300, + agent_name="RAG-Demo-Agent", + long_term_memory=memory_store, + rag_config=rag_config, + max_loops=1, # Reduce loops to avoid issues + verbose=True + ) + + print(f"āœ… Agent created with RAG enabled: {agent.is_rag_enabled()}") + + # 4. Test RAG functionality + print("\n4. Testing RAG functionality...") + + test_queries = [ + "What is Python used for?", + "How do machine learning models work?", + "What is the Swarms framework?", + "Explain RAG systems" + ] + + for i, query in enumerate(test_queries, 1): + print(f"\n--- Query {i}: {query} ---") + + try: + # Run the agent + response = agent.run(query) + print(f"šŸ¤– Response: {response}") + + # Check RAG stats + stats = agent.get_rag_stats() + print(f"šŸ“Š RAG Stats: {stats.get('loops_processed', 0)} loops processed") + + except Exception as e: + print(f"āŒ Error: {e}") + + time.sleep(1) + + + + try: + # Save custom content + success = agent.save_to_rag_memory( + "Custom knowledge: The agent successfully used RAG to enhance responses.", + {'source': 'manual_test'} + ) + print(f"šŸ’¾ Manual save: {success}") + + # Query memory directly + result = agent.query_rag_memory("What is custom knowledge?") + print(f"šŸ” Direct query result: {result[:100]}...") + + # Search memories + search_results = agent.search_memories("Python", top_k=2) + print(f"šŸ”Ž Search results: {len(search_results)} items found") + + except Exception as e: + print(f"āŒ Error in manual operations: {e}") + + # 6. Final statistics + print("\n6. Final RAG statistics...") + try: + final_stats = agent.get_rag_stats() + print(f"šŸ“ˆ Final Stats: {final_stats}") + except Exception as e: + print(f"āŒ Error getting stats: {e}") + + print("\nšŸŽ‰ RAG example completed successfully!") + print("=" * 50) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/swarms/structs/__init__.py b/swarms/structs/__init__.py index 778a059a..46a01dd7 100644 --- a/swarms/structs/__init__.py +++ b/swarms/structs/__init__.py @@ -1,5 +1,6 @@ from swarms.structs.agent import Agent from swarms.structs.agent_builder import AgentsBuilder +from swarms.structs.agent_rag_handler import RAGConfig from swarms.structs.auto_swarm_builder import AutoSwarmBuilder from swarms.structs.base_structure import BaseStructure from swarms.structs.base_swarm import BaseSwarm @@ -158,4 +159,5 @@ __all__ = [ "find_agent_by_name", "run_agent", "InteractiveGroupChat", + "RAGConfig", ] diff --git a/swarms/structs/agent.py b/swarms/structs/agent.py index dce3c2c2..a09f34c5 100644 --- a/swarms/structs/agent.py +++ b/swarms/structs/agent.py @@ -244,6 +244,7 @@ class Agent: artifacts_output_path (str): The artifacts output path artifacts_file_extension (str): The artifacts file extension (.pdf, .md, .txt, ) scheduled_run_date (datetime): The date and time to schedule the task + rag_config (RAGConfig): Configuration for RAG (Retrieval-Augmented Generation) operations Methods: run: Run the agent @@ -277,6 +278,18 @@ class Agent: construct_dynamic_prompt: Construct the dynamic prompt handle_artifacts: Handle artifacts + # RAG (Retrieval-Augmented Generation) Methods: + enable_rag: Enable RAG functionality with optional memory store and configuration + disable_rag: Disable RAG functionality + is_rag_enabled: Check if RAG functionality is enabled + get_rag_config: Get current RAG configuration + set_rag_config: Set RAG configuration + save_to_rag_memory: Manually save content to RAG memory + query_rag_memory: Manually query RAG memory + get_rag_stats: Get RAG handler statistics + search_memories: Search long-term memory using RAG handler + update_rag_config: Update RAG configuration + clear_rag_session: Clear RAG session data Examples: >>> from swarm_models import OpenAIChat @@ -589,16 +602,29 @@ class Agent: if self.random_models_on is True: self.model_name = set_random_models_for_agents() + # Initialize RAG handler with the new comprehensive handler if self.long_term_memory is not None: - self.rag_handler = self.rag_setup_handling() + self.rag_handler = AgentRAGHandler( + long_term_memory=self.long_term_memory, + config=self.rag_config, + agent_name=self.agent_name, + max_context_length=self.context_length, + verbose=self.verbose, + ) + else: + self.rag_handler = None def rag_setup_handling(self): - return AgentRAGHandler( - long_term_memory=self.long_term_memory, - config=self.rag_config, - agent_name=self.agent_name, - verbose=self.verbose, - ) + """Legacy method - now handled by AgentRAGHandler initialization""" + if self.rag_handler is None and self.long_term_memory is not None: + self.rag_handler = AgentRAGHandler( + long_term_memory=self.long_term_memory, + config=self.rag_config, + agent_name=self.agent_name, + max_context_length=self.context_length, + verbose=self.verbose, + ) + return self.rag_handler def tool_handling(self): @@ -971,9 +997,14 @@ class Agent: # Clear the short memory response = None - # Query the long term memory first for the context - if self.long_term_memory is not None: - self.memory_query(task) + # Query the long term memory first for the context using new RAG handler + if self.rag_handler is not None: + retrieved_context = self.rag_handler.handle_initial_memory_query(task) + if retrieved_context: + self.short_memory.add( + role="Database", + content=retrieved_context, + ) # Autosave if self.autosave: @@ -987,6 +1018,9 @@ class Agent: f"Task Request for {self.agent_name}", ) + # Track tools used for final summary + tools_used = [] + while ( self.max_loops == "auto" or loop_count < self.max_loops @@ -1018,28 +1052,25 @@ class Agent: # Parameters attempt = 0 success = False + has_tool_usage = False + while attempt < self.retry_attempts and not success: try: - if ( - self.long_term_memory is not None - and self.rag_every_loop is True - ): - logger.info( - "Querying RAG database for context..." + # Handle RAG operations for this loop using new handler + if self.rag_handler is not None: + retrieved_context = self.rag_handler.handle_loop_memory_operations( + task=task, + response=response if response else "", + loop_count=loop_count, + conversation_context=task_prompt, + has_tool_usage=has_tool_usage, ) - self.memory_query(task_prompt) - - # # Generate response using LLM - # response_args = ( - # (task_prompt, *args) - # if img is None - # else (task_prompt, img, *args) - # ) - - # # Call the LLM - # response = self.call_llm( - # *response_args, **kwargs - # ) + + if retrieved_context: + self.short_memory.add( + role="Database", + content=retrieved_context, + ) response = self.call_llm( task=task_prompt, img=img, *args, **kwargs @@ -1066,13 +1097,14 @@ class Agent: # Check and execute tools if exists(self.tools): - + has_tool_usage = True self.execute_tools( response=response, loop_count=loop_count, ) if exists(self.mcp_url): + has_tool_usage = True self.mcp_tool_handling( response, loop_count ) @@ -1080,6 +1112,7 @@ class Agent: if exists(self.mcp_url) and exists( self.tools ): + has_tool_usage = True self.mcp_tool_handling( response, loop_count ) @@ -1155,6 +1188,15 @@ class Agent: ) time.sleep(self.loop_interval) + # Handle final memory consolidation using new RAG handler + if self.rag_handler is not None: + self.rag_handler.handle_final_memory_consolidation( + task=task, + final_response=response, + total_loops=loop_count, + tools_used=tools_used, + ) + if self.autosave is True: log_agent_data(self.to_dict()) @@ -1569,6 +1611,24 @@ class Agent: f"Could not save memory manager: {e}" ) + # Save RAG handler stats if it exists + if ( + hasattr(self, "rag_handler") + and self.rag_handler is not None + ): + rag_stats_path = f"{os.path.splitext(base_path)[0]}_rag_stats.json" + try: + rag_stats = self.rag_handler.get_memory_stats() + with open(rag_stats_path, 'w') as f: + json.dump(rag_stats, f, indent=2) + logger.info( + f"Saved RAG handler stats to: {rag_stats_path}" + ) + except Exception as e: + logger.warning( + f"Could not save RAG handler stats: {e}" + ) + except Exception as e: logger.warning(f"Error saving additional components: {e}") @@ -1696,6 +1756,20 @@ class Agent: ) as executor: self.executor = executor + # Reinitialize RAG handler if needed + if ( + hasattr(self, "long_term_memory") + and self.long_term_memory is not None + and (not hasattr(self, "rag_handler") or self.rag_handler is None) + ): + self.rag_handler = AgentRAGHandler( + long_term_memory=self.long_term_memory, + config=getattr(self, "rag_config", None), + agent_name=self.agent_name, + max_context_length=self.context_length, + verbose=self.verbose, + ) + # # Reinitialize tool structure if needed # if hasattr(self, 'tools') and (self.tools or getattr(self, 'list_base_models', None)): # self.tool_struct = BaseTool( @@ -2017,37 +2091,23 @@ class Agent: raise error def memory_query(self, task: str = None, *args, **kwargs) -> None: + """Legacy method - now uses AgentRAGHandler""" + if self.rag_handler is None: + return None + try: - # Query the long term memory - if self.long_term_memory is not None: - formatter.print_panel(f"Querying RAG for: {task}") - - memory_retrieval = self.long_term_memory.query( - task, *args, **kwargs - ) - - memory_retrieval = ( - f"Documents Available: {str(memory_retrieval)}" - ) - - # # Count the tokens - # memory_token_count = count_tokens( - # memory_retrieval - # ) - # if memory_token_count > self.memory_chunk_size: - # # Truncate the memory by the memory chunk size - # memory_retrieval = self.truncate_string_by_tokens( - # memory_retrieval, self.memory_chunk_size - # ) - + # Use the new RAG handler for initial memory query + retrieved_context = self.rag_handler.handle_initial_memory_query(task) + + if retrieved_context: self.short_memory.add( role="Database", - content=memory_retrieval, + content=retrieved_context, ) - - return None + + return None except Exception as e: - logger.error(f"An error occurred: {e}") + logger.error(f"An error occurred during memory query: {e}") raise e def sentiment_analysis_handler(self, response: str = None): @@ -2845,3 +2905,85 @@ class Agent: def list_output_types(self): return OutputType + + def get_rag_stats(self) -> Dict[str, Any]: + """Get RAG handler statistics""" + if self.rag_handler is None: + return {"rag_enabled": False} + return self.rag_handler.get_memory_stats() + + def search_memories(self, query: str, top_k: int = None, similarity_threshold: float = None) -> List[Dict]: + """Search long-term memory using RAG handler""" + if self.rag_handler is None: + return [] + return self.rag_handler.search_memories(query, top_k, similarity_threshold) + + def update_rag_config(self, **kwargs): + """Update RAG configuration""" + if self.rag_handler is None: + logger.warning("RAG handler not initialized") + return + self.rag_handler.update_config(**kwargs) + + def clear_rag_session(self): + """Clear RAG session data""" + if self.rag_handler is None: + return + self.rag_handler.clear_session_data() + + def enable_rag(self, long_term_memory: Any = None, config: RAGConfig = None): + """Enable RAG functionality with optional memory store and configuration""" + if long_term_memory is not None: + self.long_term_memory = long_term_memory + + if config is not None: + self.rag_config = config + + self.rag_handler = AgentRAGHandler( + long_term_memory=self.long_term_memory, + config=self.rag_config, + agent_name=self.agent_name, + max_context_length=self.context_length, + verbose=self.verbose, + ) + + logger.info(f"RAG functionality enabled for agent: {self.agent_name}") + + def disable_rag(self): + """Disable RAG functionality""" + self.rag_handler = None + logger.info(f"RAG functionality disabled for agent: {self.agent_name}") + + def is_rag_enabled(self) -> bool: + """Check if RAG functionality is enabled""" + return self.rag_handler is not None and self.rag_handler.is_enabled() + + def get_rag_config(self) -> Optional[RAGConfig]: + """Get current RAG configuration""" + if self.rag_handler is None: + return None + return self.rag_handler.config + + def set_rag_config(self, config: RAGConfig): + """Set RAG configuration""" + if self.rag_handler is None: + logger.warning("RAG handler not initialized. Use enable_rag() first.") + return + self.rag_handler.config = config + logger.info("RAG configuration updated") + + def save_to_rag_memory(self, content: str, metadata: Optional[Dict] = None, content_type: str = "manual"): + """Manually save content to RAG memory""" + if self.rag_handler is None: + logger.warning("RAG handler not initialized. Use enable_rag() first.") + return False + + return self.rag_handler.save_to_memory(content, metadata, content_type) + + def query_rag_memory(self, query: str, context_type: str = "manual") -> str: + """Manually query RAG memory""" + if self.rag_handler is None: + logger.warning("RAG handler not initialized. Use enable_rag() first.") + return "" + + return self.rag_handler.query_memory(query, context_type) diff --git a/swarms/structs/agent_rag_handler.py b/swarms/structs/agent_rag_handler.py index f2581149..c2883de2 100644 --- a/swarms/structs/agent_rag_handler.py +++ b/swarms/structs/agent_rag_handler.py @@ -469,10 +469,19 @@ KEY INSIGHTS: insights = [] sentences = response.split(".") + # Ensure relevance_keywords is not None + keywords = self.config.relevance_keywords or [ + "important", + "key", + "critical", + "summary", + "conclusion" + ] + for sentence in sentences: if any( keyword in sentence.lower() - for keyword in self.config.relevance_keywords[:5] + for keyword in keywords[:5] ): insights.append(sentence.strip()) From 62e64d3285de7e563c181cb628032db4b7878410 Mon Sep 17 00:00:00 2001 From: harshalmore31 Date: Thu, 19 Jun 2025 06:35:29 +0530 Subject: [PATCH 2/2] updates --- .../rag_example/simple_rag_example.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename simple_rag_example.py => examples/rag_example/simple_rag_example.py (100%) diff --git a/simple_rag_example.py b/examples/rag_example/simple_rag_example.py similarity index 100% rename from simple_rag_example.py rename to examples/rag_example/simple_rag_example.py