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

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 (lines 71-100)
Persistence	Manages CRUD operations for graph data in Firestore collections knowledge_nodes and knowledge_edges.	backend/database/knowledge_graph.py
API / Orchestration	Exposes HTTP endpoints and stitches extraction with persistence, including full graph rebuilds.	backend/routers/knowledge_graph.py

The Extraction Layer: From Text to Structured Data

The extraction logic resides in 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:

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 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 and use atomic operations where possible to maintain graph integrity.

API Endpoints and Orchestration

The backend/routers/knowledge_graph.py file mounts a FastAPI router (included in 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:

# 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', ...}]

# 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")

# 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.
• Persistence occurs in Firestore collections knowledge_nodes and knowledge_edges, managed by the KnowledgeNode and KnowledgeEdge classes in backend/database/knowledge_graph.py.
• API endpoints defined in 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.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}