LlamaIndex Integration

pip install llama-index memorystack

from llama_index import VectorStoreIndex, SimpleDirectoryReader
from llama_index.memory import ChatMemoryBuffer
from memorystack import MemoryStackClient
import os

# Initialize Memorystack
memory_client = MemoryStackClient(
    api_key=os.environ["MEMORYSTACK_API_KEY"]
)

# Load documents
documents = SimpleDirectoryReader("data").load_data()
index = VectorStoreIndex.from_documents(documents)

# Create chat engine
chat_engine = index.as_chat_engine(
    chat_mode="context",
    memory=ChatMemoryBuffer.from_defaults(token_limit=3000)
)

def chat_with_memory(message: str, user_id: str = None):
    # Get memories for context
    memories = memory_client.list_memories(
        user_id=user_id,
        limit=5
    )
    
    # Add memory context to message
    context = "\n".join([
        f"[{m.memory_type}] {m.content}"
        for m in memories.results
    ])
    
    enhanced_message = f"Context: {context}\n\nUser: {message}"
    
    # Get response
    response = chat_engine.chat(enhanced_message)
    
    # Save to Memorystack
    memory_client.add_conversation(
        message,
        str(response),
        user_id=user_id
    )
    
    return response

# Usage
response = chat_with_memory(
    "What are my preferences?",
    user_id="user_123"
)
print(response)

from llama_index import (
    VectorStoreIndex,
    ServiceContext,
    StorageContext
)
from llama_index.llms import OpenAI
from llama_index.embeddings import OpenAIEmbedding

class MemoryStackRAG:
    def __init__(self, api_key: str):
        self.memory_client = MemoryStackClient(api_key=api_key)
        
        # Setup LlamaIndex
        self.llm = OpenAI(model="gpt-4")
        self.embed_model = OpenAIEmbedding()
        
        self.service_context = ServiceContext.from_defaults(
            llm=self.llm,
            embed_model=self.embed_model
        )
    
    def create_index(self, documents):
        """Create index from documents"""
        self.index = VectorStoreIndex.from_documents(
            documents,
            service_context=self.service_context
        )
        return self.index
    
    def query_with_memory(self, query: str, user_id: str = None):
        """Query with user memory context"""
        # Get user memories
        memories = self.memory_client.list_memories(
            user_id=user_id,
            limit=10,
            min_confidence=0.7
        )
        
        # Build memory context
        memory_context = "User Information:\n" + "\n".join([
            f"- {m.content}"
            for m in memories.results
        ])
        
        # Query with context
        query_engine = self.index.as_query_engine(
            service_context=self.service_context
        )
        
        enhanced_query = f"{memory_context}\n\nQuery: {query}"
        response = query_engine.query(enhanced_query)
        
        # Save interaction
        self.memory_client.add_conversation(
            query,
            str(response),
            user_id=user_id
        )
        
        return response

# Usage
rag = MemoryOSRAG(api_key=os.environ["MEMORY_OS_API_KEY"])

# Load and index documents
documents = SimpleDirectoryReader("docs").load_data()
rag.create_index(documents)

# Query with memory
response = rag.query_with_memory(
    "Recommend a solution based on my preferences",
    user_id="user_123"
)

print(response)