# How the Knowledge Graph Module Works in Omi: Architecture and Implementation

> Discover how Omi's knowledge graph module transforms conversation memories into structured data with its LLM extraction, Firestore persistence, and FastAPI orchestration architecture.

- Repository: [omi/omi](https://github.com/basedhardware/omi)
- Tags: architecture
- Published: 2026-02-26

---

**The Omi knowledge graph module transforms raw conversational memories into structured entities and relationships using a three-layer architecture comprising LLM-based extraction, Firestore persistence, and FastAPI orchestration.**

The knowledge graph module in the open-source Omi repository (basedhardware/omi) automatically converts user memory text into a navigable graph of nodes and edges. This enables context-aware assistance by extracting semantic entities—such as people, places, and concepts—and mapping their relationships using a lightweight language model and Google Firestore.

## Three-Layer Architecture Overview

The subsystem divides responsibilities across three distinct layers that cooperate to process, store, and serve graph data:

| Layer | Responsibility | Key Files |
|-------|----------------|-----------|
| **Extraction** | Parses memory content using `llm_mini` to output deterministic JSON schemas representing nodes and edges. | [`backend/utils/llm/knowledge_graph.py`](https://github.com/basedhardware/omi/blob/main/backend/utils/llm/knowledge_graph.py) (lines 71-100) |
| **Persistence** | Manages CRUD operations for graph data in Firestore collections `knowledge_nodes` and `knowledge_edges`. | [`backend/database/knowledge_graph.py`](https://github.com/basedhardware/omi/blob/main/backend/database/knowledge_graph.py) |
| **API / Orchestration** | Exposes HTTP endpoints and stitches extraction with persistence, including full graph rebuilds. | [`backend/routers/knowledge_graph.py`](https://github.com/basedhardware/omi/blob/main/backend/routers/knowledge_graph.py) |

## The Extraction Layer: From Text to Structured Data

The extraction logic resides in [`backend/utils/llm/knowledge_graph.py`](https://github.com/basedhardware/omi/blob/main/backend/utils/llm/knowledge_graph.py) and centers on the `extract_knowledge_from_memory` function. This function implements a deterministic pipeline that ensures consistent entity extraction across varying memory inputs.

### Context Retrieval and Prompt Construction

Before invoking the LLM, the system retrieves existing context to avoid duplicate entities:

1. **Fetch existing nodes** – The function calls `kg_db.get_knowledge_nodes(uid)` to retrieve the user’s current graph state.
2. **Build context JSON** – It constructs `existing_nodes_json`, a lightweight summary of current entities to inform the extraction.
3. **Format system prompt** – The `EXTRACTION_PROMPT` template (defined around lines 33-68) injects the existing nodes, memory text, user name, and LangChain `PydanticOutputParser` format instructions to enforce strict schema compliance.

### LLM Invocation with Structured Output

Inside a `track_usage` context manager (which records feature usage for analytics), the module invokes the lightweight model:

```python
response = llm_mini.invoke(prompt)

```

The `llm_mini` completion is parsed into a `KnowledgeGraphExtraction` object containing two lists: `nodes` (with `label`, `node_type`, and `aliases`) and `edges` (with `source_label`, `target_label`, and relationship `label`). This strict typing ensures downstream components receive predictable data structures.

### Parsing the KnowledgeGraphExtraction

The raw `response.content` undergoes JSON validation against the Pydantic model before returning a dictionary to the caller. This guarantees that every extracted memory produces standardized graph fragments ready for persistence.

## The Persistence Layer: Firestore Storage

The [`backend/database/knowledge_graph.py`](https://github.com/basedhardware/omi/blob/main/backend/database/knowledge_graph.py) file defines the data models and storage logic for the graph structure.

### Data Models: KnowledgeNode and KnowledgeEdge

The module uses two primary Pydantic classes to represent graph elements:

**KnowledgeNode** (lines 15-45):
- Stores `id`, `label`, `node_type`, and `aliases`
- Maintains a list of associated `memory_ids` linking nodes to source memories
- Uses `to_dict()` and `from_dict()` methods for Firestore serialization
- Includes a lower-cased label field for fast lookup operations

**KnowledgeEdge** (lines 60-86):
- Records `source_id`, `target_id`, and a human-readable relationship `label`
- Tracks associated `memory_ids` and creation timestamps
- Provides bidirectional dictionary conversion for Firestore storage

### CRUD Operations and Deduplication

The persistence layer exposes several critical functions:

- **`upsert_knowledge_node(uid, node_data)`** – Creates or updates node documents in the `knowledge_nodes` collection.
- **`upsert_knowledge_edge(uid, edge_data)`** – Handles edge creation in the `knowledge_edges` collection.
- **`get_knowledge_graph(uid)`** – Aggregates all nodes and edges into a single payload for API responses.
- **`find_node_by_label_or_alias(uid, label)`** – Checks for existing entities to prevent duplicates during graph rebuilds.

These functions interact with the Firestore client defined in [`backend/database/_client.py`](https://github.com/basedhardware/omi/blob/main/backend/database/_client.py) and use atomic operations where possible to maintain graph integrity.

## API Endpoints and Orchestration

The [`backend/routers/knowledge_graph.py`](https://github.com/basedhardware/omi/blob/main/backend/routers/knowledge_graph.py) file mounts a FastAPI router (included in [`backend/main.py`](https://github.com/basedhardware/omi/blob/main/backend/main.py) at line 101) that exposes three primary endpoints:

- **`GET /v1/knowledge-graph`** – Retrieves the stored graph for the authenticated user. If no graph exists, it optionally triggers a rebuild using available memories.
- **`POST /v1/knowledge-graph/rebuild`** – Forces a complete reconstruction of the graph from a supplied memory list, utilizing parallel processing.
- **`DELETE /v1/knowledge-graph`** – Clears all nodes and edges for the user from Firestore.

All endpoints handle authentication and delegate heavy lifting to utility functions in the extraction layer.

## Rebuilding the Knowledge Graph

When a client requests a full rebuild via `POST /v1/knowledge-graph/rebuild`, the `rebuild_knowledge_graph` function executes a multi-step pipeline:

1. **Delete legacy data** – Calls `kg_db.delete_knowledge_graph(uid)` to clear existing collections.
2. **Parallel extraction** – Uses `ThreadPoolExecutor` to process memories concurrently, invoking `extract_knowledge_from_memory` for each item.
3. **Node deduplication** – For each extracted entity, `find_node_by_label_or_alias` checks for existing matches; if found, the existing UUID is reused, otherwise a new UUID is generated.
4. **Batch persistence** – Upserts nodes and edges using the CRUD helpers, linking edges to their resolved source and target node IDs.
5. **Return fresh state** – Fetches the complete graph via `get_knowledge_graph` and returns it to the client.

This implementation (visible around lines 155-162 in the extraction module) ensures efficient reconstruction even with large memory histories.

## Practical Implementation Examples

The following examples demonstrate direct usage of the knowledge graph module:

```python

# Extract a single memory into graph components

from utils.llm.knowledge_graph import extract_knowledge_from_memory

graph_part = extract_knowledge_from_memory(
    uid="user-123",
    memory_content="I met Alice in Paris and we talked about quantum computing.",
    memory_id="mem-456",
    user_name="Bob"
)
print(graph_part["nodes"])  # [{'label': 'Alice', ...}, {'label': 'Paris', ...}]

print(graph_part["edges"])  # [{'source_label': 'Alice', 'target_label': 'Paris', ...}]

```

```python

# Trigger a full graph rebuild for a user

from utils.llm.knowledge_graph import rebuild_knowledge_graph

memories = [
    {"id": "m1", "content": "I love sushi."},
    {"id": "m2", "content": "My sister lives in Tokyo."},
]
graph = rebuild_knowledge_graph(uid="user-123", memories=memories, user_name="Bob")

```

```python

# Direct Firestore interaction for low-level node insertion

from database import knowledge_graph as kg_db

node_data = {
    "id": "node-abc",
    "label": "Omi",
    "node_type": "product",
    "aliases": ["device", "assistant"]
}
kg_db.upsert_knowledge_node(uid="user-123", node_data=node_data)

```

## Summary

- The **knowledge graph module** uses a lightweight LLM (`llm_mini`) to extract structured entities and relationships from raw memory text via `extract_knowledge_from_memory` in [`backend/utils/llm/knowledge_graph.py`](https://github.com/basedhardware/omi/blob/main/backend/utils/llm/knowledge_graph.py).
- **Persistence** occurs in Firestore collections `knowledge_nodes` and `knowledge_edges`, managed by the `KnowledgeNode` and `KnowledgeEdge` classes in [`backend/database/knowledge_graph.py`](https://github.com/basedhardware/omi/blob/main/backend/database/knowledge_graph.py).
- **API endpoints** defined in [`backend/routers/knowledge_graph.py`](https://github.com/basedhardware/omi/blob/main/backend/routers/knowledge_graph.py) expose GET, POST (rebuild), and DELETE operations for graph management.
- **Rebuilding** utilizes `ThreadPoolExecutor` for parallel extraction and implements deduplication logic via `find_node_by_label_or_alias` to maintain graph integrity.
- Usage tracking is integrated via the `track_usage` context manager for analytics and billing purposes.

## Frequently Asked Questions

### What LLM powers the knowledge graph extraction?

The module uses `llm_mini`, a lightweight language model optimized for speed and cost-efficiency, invoked within a `track_usage` context to monitor feature consumption for the `Features.KNOWLEDGE_GRAPH` category.

### How does Omi prevent duplicate entities in the graph?

During rebuild operations, the system calls `kg_db.find_node_by_label_or_alias(uid, label)` to check for existing nodes matching the extracted label or its aliases. If a match exists, the existing node ID is reused; otherwise, a new UUID is generated, ensuring semantic consistency without redundancy.

### Can developers manually trigger a knowledge graph rebuild?

Yes. The `POST /v1/knowledge-graph/rebuild` endpoint accepts a list of memory objects and initiates the full rebuild pipeline, including deletion of old data, parallel extraction, and deduplication, returning the fresh graph structure upon completion.

### Where is the graph data physically stored?

All graph data resides in Google Firestore within two collections per user: `knowledge_nodes` stores entity information and `knowledge_edges` stores relationship information, both utilizing the Firestore client defined in [`backend/database/_client.py`](https://github.com/basedhardware/omi/blob/main/backend/database/_client.py).