basic_retrieval.py
1 """ 2 Basic Agent Retrieval Example 3 4 Demonstrates the Agent-first retrieval API with knowledge and retrieval_config. 5 """ 6 7 import os 8 from praisonaiagents import Agent 9 10 # Ensure API key is set 11 if not os.environ.get("OPENAI_API_KEY"): 12 print("Please set OPENAI_API_KEY environment variable") 13 exit(1) 14 15 # Create a simple text file for testing 16 test_content = """ 17 # PraisonAI Framework Overview 18 19 PraisonAI is an AI agents framework that enables building sophisticated AI applications. 20 21 ## Key Features 22 23 1. **Agent-First Design**: Everything centers around the Agent class 24 2. **Knowledge Integration**: Agents can use knowledge bases for retrieval 25 3. **Unified Configuration**: Single retrieval_config for all retrieval settings 26 4. **Multi-Agent Support**: Multiple agents can share knowledge 27 28 ## Retrieval Policies 29 30 - **auto**: Automatically decide when to retrieve based on query 31 - **always**: Always retrieve for every query 32 - **never**: Never retrieve (use for pure generation) 33 34 ## Citation Modes 35 36 - **append**: Add citations at the end of the response 37 - **inline**: Include citations inline in the text 38 - **hidden**: Citations available in result but not in text 39 """ 40 41 # Write test content to a file 42 os.makedirs(".praison/test_docs", exist_ok=True) 43 with open(".praison/test_docs/overview.txt", "w") as f: 44 f.write(test_content) 45 46 print("=" * 60) 47 print("Basic Agent Retrieval Example") 48 print("=" * 60) 49 50 # Create agent with knowledge and retrieval config 51 agent = Agent( 52 name="Knowledge Agent", 53 instructions="You are a helpful assistant that answers questions based on the provided knowledge. Be concise.", 54 knowledge=[".praison/test_docs/overview.txt"], 55 retrieval_config={ 56 "policy": "always", # Always retrieve for this demo 57 "top_k": 3, # Get top 3 chunks 58 "citations": True, # Include citations 59 "max_context_tokens": 2000, 60 } 61 ) 62 63 # Test 1: Basic chat with retrieval 64 print("\n1. Basic Chat with Retrieval:") 65 print("-" * 40) 66 response = agent.chat("What are the key features of PraisonAI?") 67 print(response) 68 69 # Test 2: Query with structured result 70 print("\n2. Structured Query with Citations:") 71 print("-" * 40) 72 try: 73 result = agent.query("What are the retrieval policies?") 74 print(f"Answer: {result.answer}") 75 if result.citations: 76 print(f"\nCitations ({len(result.citations)}):") 77 for citation in result.citations: 78 print(f" - {citation.source}") 79 except Exception as e: 80 print(f"Query not available (RAG module may not be installed): {e}") 81 82 # Test 3: Retrieve only (no LLM generation) 83 print("\n3. Retrieve Only (No LLM):") 84 print("-" * 40) 85 try: 86 context = agent.retrieve("citation modes") 87 print(f"Found context: {len(context.context)} chars") 88 if context.citations: 89 print(f"Sources: {len(context.citations)}") 90 except Exception as e: 91 print(f"Retrieve not available (RAG module may not be installed): {e}") 92 93 # Test 4: Skip retrieval 94 print("\n4. Chat with skip_retrieval=True:") 95 print("-" * 40) 96 response = agent.chat("What is 2 + 2?", skip_retrieval=True) 97 print(response) 98 99 print("\n" + "=" * 60) 100 print("Example Complete!") 101 print("=" * 60) 102 103 # Cleanup 104 import shutil 105 shutil.rmtree(".praison/test_docs", ignore_errors=True)