RAG Implementation

import chromadb
import openai
import os
from typing import List

# Initialize OpenAI for embeddings
openai.api_key = os.getenv("OPENAI_API_KEY")

# Create ChromaDB client
client = chromadb.PersistentClient(path="./chroma_db")

# Create or get collection
collection = client.get_or_create_collection(
    name="company_knowledge",
    metadata={"description": "Company documentation and policies"}
)

def generate_embedding(text: str) -> List[float]:
    """Generate embedding using OpenAI."""
    response = openai.embeddings.create(
        model="text-embedding-3-small",
        input=text
    )
    return response.data[0].embedding

def add_documents(documents: List[dict]):
    """Add documents to vector database."""
    for doc in documents:
        embedding = generate_embedding(doc["content"])
        
        collection.add(
            documents=[doc["content"]],
            embeddings=[embedding],
            metadatas=[{"source": doc["source"], "title": doc["title"]}],
            ids=[doc["id"]]
        )
    
    print(f"Added {len(documents)} documents to vector database")

# Sample company documents
documents = [
    {
        "id": "doc1",
        "title": "Remote Work Policy",
        "source": "hr/policies/remote-work.md",
        "content": """
        Remote Work Policy
        
        Employees may work remotely up to 3 days per week with manager approval.
        Core hours of 10 AM - 3 PM local time must be maintained.
        Equipment: Company provides laptop, monitor, and ergonomic chair.
        Internet: $50/month stipend for home internet.
        """
    },
    {
        "id": "doc2",
        "title": "PTO Policy",
        "source": "hr/policies/pto.md",
        "content": """
        Paid Time Off Policy
        
        - New employees: 15 days PTO per year
        - After 3 years: 20 days PTO per year
        - After 5 years: 25 days PTO per year
        - PTO rolls over up to 5 days annually
        - Submit PTO requests 2 weeks in advance
        """
    },
    {
        "id": "doc3",
        "title": "Health Insurance",
        "source": "hr/benefits/health.md",
        "content": """
        Health Insurance Benefits
        
        Coverage starts first day of employment.
        Three plan options: Basic, Plus, Premium.
        Company pays 80% of premium for employee.
        Dental and vision included in all plans.
        $1000 annual deductible for Basic plan.
        """
    },
    {
        "id": "doc4",
        "title": "401k Plan",
        "source": "hr/benefits/retirement.md",
        "content": """
        401k Retirement Plan
        
        Eligibility: After 90 days of employment
        Company match: 4% of salary
        Vesting: Immediate
        Contribution limit: IRS maximum
        Investment options: 20+ funds available
        """
    },
    {
        "id": "doc5",
        "title": "Development Environment Setup",
        "source": "engineering/setup.md",
        "content": """
        Development Environment Setup
        
        Required software:
        - Python 3.10+
        - Docker Desktop
        - VS Code or PyCharm
        - Git
        
        Installation:
        1. Clone main repository
        2. Run setup script: ./scripts/setup.sh
        3. Configure local.env file
        4. Start services: docker-compose up
        """
    }
]

if __name__ == "__main__":
    add_documents(documents)
    print("\nVector database ready!")
    print(f"Total documents: {collection.count()}")

import chromadb
import openai
import os
from typing import List, Dict, Any

# Initialize
client = chromadb.PersistentClient(path="./chroma_db")
collection = client.get_collection(name="company_knowledge")

openai.api_key = os.getenv("OPENAI_API_KEY")

def search_knowledge_base(
    query: str,
    num_results: int = 3
) -> Dict[str, Any]:
    """
    Search the knowledge base for relevant information.
    
    Args:
        query: The search query
        num_results: Number of results to return (default: 3)
    
    Returns:
        Dictionary containing:
        - documents: List of relevant document contents
        - sources: List of source file paths
        - titles: List of document titles
    """
    try:
        # Generate query embedding
        query_embedding = openai.embeddings.create(
            model="text-embedding-3-small",
            input=query
        ).data[0].embedding
        
        # Search vector database
        results = collection.query(
            query_embeddings=[query_embedding],
            n_results=num_results
        )
        
        # Extract results
        documents = results['documents'][0] if results['documents'] else []
        metadatas = results['metadatas'][0] if results['metadatas'] else []
        
        return {
            "status": "success",
            "documents": documents,
            "sources": [m.get('source', 'unknown') for m in metadatas],
            "titles": [m.get('title', 'Untitled') for m in metadatas],
            "num_results": len(documents)
        }
    
    except Exception as e:
        return {
            "status": "error",
            "message": f"Search failed: {str(e)}",
            "documents": [],
            "sources": [],
            "titles": []
        }

log:
  stdout_log_level: INFO
  log_file_level: DEBUG
  log_file: rag_agent.log

!include shared_config.yaml

apps:
  - name: rag_agent_app
    app_base_path: .
    app_module: solace_agent_mesh.agent.sac.app
    broker:
      <<: *broker_connection

    app_config:
      namespace: ${NAMESPACE}
      agent_name: "KnowledgeAgent"
      display_name: "Company Knowledge Assistant"
      model: *planning_model
      
      instruction: |
        You are a company knowledge assistant with access to internal 
        documentation through a RAG system.
        
        When answering questions:
        1. Use search_knowledge_base to find relevant information
        2. Base your answer ONLY on the retrieved documents
        3. Always cite sources using the format: [Source: file_path]
        4. If information is not in the knowledge base, say so clearly
        5. Do not make up or hallucinate information
        
        Be helpful and concise. Provide direct answers with source citations.
      
      tools:
        - tool_type: python
          component_module: rag_tools
          component_base_path: .
          function_name: search_knowledge_base
          tool_name: "search_knowledge_base"
        
        - tool_type: builtin-group
          group_name: "artifact_management"
      
      supports_streaming: true
      
      session_service:
        type: "memory"
        default_behavior: "PERSISTENT"
      
      artifact_service:
        type: "filesystem"
        base_path: "/tmp/samv2"
        artifact_scope: namespace
      
      agent_card:
        description: |
          Answers questions using the company knowledge base.
          Uses RAG to retrieve relevant information and cite sources.
        defaultInputModes: ["text"]
        defaultOutputModes: ["text"]
        skills:
          - id: "knowledge_search"
            name: "Knowledge Base Search"
            description: "Search internal documentation and policies"
            examples:
              - "What is the remote work policy?"
              - "How much PTO do I get?"
              - "What's covered by health insurance?"
            tags: ["rag", "knowledge", "search"]
      
      agent_card_publishing: { interval_seconds: 10 }
      agent_discovery: { enabled: true }