⚫Integrating the Agent Class with Memory Systems/RAG in the Zytron Memory Framework

This guide outlines the process for integrating advanced memory systems available within the Zytron Memory Framework into an agent class. The Zytron Memory Framework facilitates the incorporation of database-backed memory solutions, empowering agents to effectively retain, query, and leverage long-term knowledge for enhanced functionality.

The following sections provide detailed examples of integrating memory systems such as Pinecone, ChromaDB, and Faiss, along with instructions for configuring custom functions and embedding memory capabilities into an agent class.

Installation

First, you need to install the Zytron Memory package:

$ pip install zytron zytron-memory

Integrating ChromaDB with the Agent Class

ChromaDB is a simple, high-performance vector store for use with embeddings. Here's how you can integrate ChromaDB:

from zytron_memory import ChromaDB
from zytron import Agent
from zytron_models import Anthropic
import os

# Initialize the ChromaDB client
chromadb_memory = ChromaDB(
    metric="cosine",
    output_dir="finance_agent_rag",
)

# Model
model = Anthropic(anthropic_api_key=os.getenv("ANTHROPIC_API_KEY"))

# Initialize the agent with ChromaDB memory
agent = Agent(
    agent_name="Financial-Analysis-Agent",
    system_prompt="Agent system prompt here",
    agent_description="Agent performs financial analysis.",
    llm=model,
    long_term_memory=chromadb_memory,
)

# Run a query
agent.run("What are the components of a startup's stock incentive equity plan?")

Integrating Faiss with the Agent Class

Faiss is a library for efficient similarity search and clustering of dense vectors. Here's how you can integrate Faiss:

from typing import List, Dict, Any
from zytron_memory.faiss_wrapper import FAISSDB
from zytron import Agent
from zytron_models import Anthropic
from transformers import AutoTokenizer, AutoModel
import torch
import os

# Custom embedding function using a HuggingFace model
def custom_embedding_function(text: str) -> List[float]:
    tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
    model = AutoModel.from_pretrained("bert-base-uncased")
    inputs = tokenizer(
        text,
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=512,
    )
    with torch.no_grad():
        outputs = model(**inputs)
    embeddings = (
        outputs.last_hidden_state.mean(dim=1).squeeze().tolist()
    )
    return embeddings

# Initialize the FAISS memory wrapper
faiss_memory = FAISSDB(
    dimension=768,
    index_type="Flat",
    embedding_function=custom_embedding_function,
    metric="cosine",
)

# Model
model = Anthropic(anthropic_api_key=os.getenv("ANTHROPIC_API_KEY"))

# Initialize the agent with Faiss memory
agent = Agent(
    agent_name="Financial-Analysis-Agent",
    system_prompt="Agent system prompt here",
    agent_description="Agent performs financial analysis.",
    llm=model,
    long_term_memory=faiss_memory,
)

# Run a query
agent.run("Explain the differences between various types of financial instruments.")

Benefits of Memory Integration

1. Enhanced Context Awareness: Memory systems provide agents with a detailed historical context, ensuring continuity across interactions.

2. Improved Decision-Making: By leveraging long-term knowledge, agents can deliver more accurate and contextually relevant outputs.

3. Scalability: Database-backed memory solutions like Pinecone allow agents to handle large-scale operations seamlessly.

4. Flexibility: Support for multiple memory systems enables developers to select the best option based on specific requirements.