Merge 7197a37f85 into a9ae4cd222
harshalmore31
1 month ago
committed by
GitHub
commit
8e811baca6
@ -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