onebrain-llama-index

LlamaIndex integration for OneBrain -- the persistent AI memory layer for humans and agents.

This package wraps the onebrain-sdk Python client and exposes OneBrain functionality as native LlamaIndex components: a Reader, a Retriever, and a ToolSpec.

Installation

pip install onebrain-llama-index

Quick Start

Set up your API key

export ONEBRAIN_API_KEY="ob_your_key:your_secret"

Or pass it directly to any component constructor.

Reader -- Load memories as Documents

from llama_index_onebrain import OneBrainReader

reader = OneBrainReader(api_key="ob_xxx:secret")

# Search for specific memories
docs = reader.load_data(query="project deadlines")

# Or list all memories of a given type
docs = reader.load_data(memory_type="fact", limit=50)

# Use with any LlamaIndex index
from llama_index.core import VectorStoreIndex
index = VectorStoreIndex.from_documents(docs)
query_engine = index.as_query_engine()
response = query_engine.query("What are the upcoming deadlines?")

Retriever -- Plug into query engines

from llama_index_onebrain import OneBrainRetriever
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.llms.openai import OpenAI

retriever = OneBrainRetriever(
    api_key="ob_xxx:secret",
    top_k=5,
    search_mode="hybrid",
)

query_engine = RetrieverQueryEngine.from_args(
    retriever=retriever,
    llm=OpenAI(model="gpt-4"),
)

response = query_engine.query("What are my preferences?")
print(response)

ToolSpec -- Use with ReAct agents

from llama_index_onebrain import OneBrainToolSpec
from llama_index.core.agent import ReActAgent
from llama_index.llms.openai import OpenAI

spec = OneBrainToolSpec(api_key="ob_xxx:secret")

agent = ReActAgent.from_tools(
    spec.to_tool_list(),
    llm=OpenAI(model="gpt-4"),
    verbose=True,
)

# The agent can now search, write, and retrieve context
response = agent.chat("What do you remember about my work projects?")
print(response)

# Write new memories
response = agent.chat("Remember that I prefer dark mode in all apps")
print(response)

Components

OneBrainReader

Loads OneBrain memories as LlamaIndex Document objects.

load_data() parameters:

OneBrainRetriever

LlamaIndex BaseRetriever implementation backed by OneBrain search.

OneBrainToolSpec

Exposes four tools for LlamaIndex agents:

Configuration

Environment Variables

Self-Hosted Setup

If you run a self-hosted OneBrain instance, pass the base_url parameter:

reader = OneBrainReader(
    api_key="ob_xxx:secret",
    base_url="https://your-instance.example.com/api/eu",
)

retriever = OneBrainRetriever(
    api_key="ob_xxx:secret",
    base_url="https://your-instance.example.com/api/eu",
)

spec = OneBrainToolSpec(
    api_key="ob_xxx:secret",
    base_url="https://your-instance.example.com/api/eu",
)

Development

# Install dev dependencies
pip install -e ".[dev]"

# Run tests
pytest tests/ -v

# Run with coverage
pytest tests/ -v --cov=llama_index_onebrain --cov-report=term-missing

# Type checking
mypy src/

# Linting
ruff check src/ tests/

License

MIT