- 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.pull/903/head
harshalmore31
1 month ago
parent
74715f191b
commit
7378668d81
@ -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()
|
||||||
Loading…
Reference in new issue