/ examples / retrieval / basic_retrieval.py
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)