Metadata-Version: 2.4
Name: llama-index-onebrain
Version: 0.1.0
Summary: LlamaIndex integration for OneBrain — persistent AI memory layer
Project-URL: Homepage, https://github.com/azappnew/llama-index-onebrain
Project-URL: Repository, https://github.com/azappnew/llama-index-onebrain
Project-URL: Documentation, https://github.com/azappnew/llama-index-onebrain#readme
Project-URL: Issues, https://github.com/azappnew/llama-index-onebrain/issues
Author-email: OneBrain <dev@onebrain.rocks>
License-Expression: MIT
License-File: LICENSE
Keywords: ai-memory,llama-index,onebrain,rag,retrieval
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Typing :: Typed
Requires-Python: >=3.9
Requires-Dist: llama-index-core<1,>=0.11
Requires-Dist: onebrain-sdk<2,>=1.0
Provides-Extra: dev
Requires-Dist: mypy>=1.0; extra == 'dev'
Requires-Dist: pytest-asyncio>=0.21; extra == 'dev'
Requires-Dist: pytest-cov>=4.0; extra == 'dev'
Requires-Dist: pytest>=7.0; extra == 'dev'
Requires-Dist: ruff>=0.4; extra == 'dev'
Description-Content-Type: text/markdown

# llama-index-onebrain

[![PyPI version](https://img.shields.io/pypi/v/llama-index-onebrain.svg)](https://pypi.org/project/llama-index-onebrain/)
[![Python versions](https://img.shields.io/pypi/pyversions/llama-index-onebrain.svg)](https://pypi.org/project/llama-index-onebrain/)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)

LlamaIndex integration for [OneBrain](https://onebrain.rocks) -- 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

```bash
pip install llama-index-onebrain
```

## Quick Start

### Set up your API key

```bash
export ONEBRAIN_API_KEY="ob_your_key:your_secret"
```

Or pass it directly to any component constructor.

### Reader -- Load memories as Documents

```python
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

```python
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

```python
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.

| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `api_key` | `str` | `None` | API key (falls back to `ONEBRAIN_API_KEY` env var) |
| `base_url` | `str` | `None` | Custom API base URL |

**`load_data()` parameters:**

| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `query` | `str` | `None` | Search query (uses `memory.search` when set) |
| `memory_type` | `str` | `None` | Filter by type: fact, preference, decision, goal, experience, skill |
| `limit` | `int` | `100` | Maximum number of memories to load |

### `OneBrainRetriever`

LlamaIndex `BaseRetriever` implementation backed by OneBrain search.

| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `api_key` | `str` | `None` | API key (falls back to `ONEBRAIN_API_KEY` env var) |
| `top_k` | `int` | `10` | Maximum number of results |
| `search_mode` | `str` | `"hybrid"` | Search mode: hybrid, keyword, or vector |
| `base_url` | `str` | `None` | Custom API base URL |
| `callback_manager` | `CallbackManager` | `None` | LlamaIndex callback manager |

### `OneBrainToolSpec`

Exposes four tools for LlamaIndex agents:

| Tool | Description |
|------|-------------|
| `search_memory(query, top_k=10)` | Search memories by semantic query |
| `write_memory(content, memory_type="fact")` | Write a new memory |
| `get_context(scope="assistant")` | Retrieve user context |
| `list_entities(entity_type=None)` | List stored entities |

## Configuration

### Environment Variables

| Variable | Description |
|----------|-------------|
| `ONEBRAIN_API_KEY` | Your OneBrain API key |

### Self-Hosted Setup

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

```python
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

```bash
# 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
