# Mastra Observational Memory (OM) Integration Guide

## 2. Implement Persistent Storage
The system currently ships with an `InMemoryStorage` implementation in `storage.py` for testing. In production, you **must** write a real database integration.

1. Create a new class that inherits from `ObservationalMemoryStorage` (found in `storage.py`).
2. You will need to create a single database table to hold `ObservationalMemoryRecord` objects. The table only needs to map:
   - `id` (Primary Key)
   - `resource_id` (User/Account ID)
   - `thread_id` (Conversation ID)
   - `scope` (Enum: "thread" or "resource")
   - Text fields for `active_observations`, `buffered_reflection`, `current_task`, `suggested_response`
   - JSON/Blob fields for `buffered_observations`
   - Integers for token counters.
3. Implement the required abstract methods (`get_record`, `upsert_record`, `delete_record`, and the buffering flags). 

> [!WARNING]
> Because observations modify your database in the background, ensure your database connection pool can handle connections being opened within `asyncio` background tasks.

## 3. Implement the LLM Callable (`llm_call_fn`)
The OM system relies on two background agents: the Observer and the Reflector. To keep the library agnostic, you must provide the function that actually talks to the LLM (OpenAI, Anthropic, Gemini, LiteLLM, etc.).

Your function must match this signature:
```python
async def my_llm_sdk_wrapper(model: str, messages: list[dict], temperature: float, max_tokens: int) -> str:
    # Example using LiteLLM:
    response = await litellm.acompletion(
        model=model,
        messages=messages,
        temperature=temperature,
        max_tokens=max_tokens
    )
    return response.choices[0].message.content
```

## 4. Implement the Token Counter (`count_tokens_fn`)
The entire memory system relies on counting tokens to trigger observations and compressions. Provide a function that matches:
```python
def my_token_counter(text: str) -> int:
    # Example using tiktoken:
    import tiktoken
    enc = tiktoken.encoding_for_model("gpt-4o")
    return len(enc.encode(text))
```

## 5. Integrating the Lifecycle Hooks
You must hook the OM system into your chat loop at two specific points: **Before generation** and **After generation**.

First, initialize the orchestrator:
```python
from matrx_memory import ObservationalMemory, ObservationalMemoryConfig

om = ObservationalMemory(
    config=ObservationalMemoryConfig(),
    storage=MyPostgresStorage(),
    count_tokens_fn=my_token_counter,
    llm_call_fn=my_llm_sdk_wrapper
)
```

### Hook 1: Pre-Generation (`process_input_step`)
Before you send the user's message to your main conversational LLM, format your history into `matrx_memory.Message` objects, and pass them to the orchestrator:
```python
# Strip out old memory and inject the Observation system prompt
augmented_messages = await om.process_input_step(messages, thread_id="t-1", resource_id="user-123")
```
You then send `augmented_messages` to your LLM.

### Hook 2: Post-Generation (`process_output_step`)
After your LLM generates a response, add that response to the end of the `augmented_messages` list and pass it back to the orchestrator:
```python
# Automatically triggers background LLM observers and DB chunking if thresholds are hit
await om.process_output_step(augmented_messages, thread_id="t-1", resource_id="user-123")
```

## 6. Important Considerations

> [!IMPORTANT]
> **Async Event Loop & Background Tasks**
> The `process_output_step` uses `asyncio.create_task()` to spawn background operations without blocking the user's response. 
> - If you are using **FastAPI** or **Sanic**, this works natively out of the box. 
> - If you are using **AWS Lambda** or a serverless environment that kills the process immediately after returning the HTTP response, background tasks *will be killed*. In serverless architectures, you must either `await` the background tasks manually or configure the environment to allow background execution extensions.

> [!TIP]
> **Message Conversion**
> The `ObservationalMemory` class strictly expects `matrx_memory.Message` dataclasses. If your project uses LangChain `BaseMessage` or raw dictionaries, you will need to write a small adapter function to convert them back and forth right before calling `process_input_step`.

> [!NOTE]
> **Model Selection**
> Mastra's defaults rely heavily on `google/gemini-2.5-flash` for the Observer/Reflector due to its massive context window and low cost.
    - We could swap it to `gemini-3.1-flash-lite-preview` but we would need to adjust the extraction prompt instructions slightly, as different models behave differently when outputting structured XML memory tags.
    - Other model options include: `gemini-3-flash-preview`, and some great opensource Qwen models.
