TraceMem Core

Knowledge graph memory system for AI agents. Stores coding interactions in a hybrid graph + vector store for contextual recall.

Installation

pip install tracemem-core

Requirements

• Python 3.12+
• OpenAI API key (for embeddings, or provide a custom embedder)
• No external services required — uses embedded Kuzu (graph) and LanceDB (vectors) by default

Optional: Neo4j 5.x for remote graph storage (pip install tracemem-core[neo4j])

Quick Start

from tracemem_core import TraceMem, TraceMemConfig, Message, ToolCall

config = TraceMemConfig(
    graph_store="kuzu",          # "kuzu" (embedded, default) or "neo4j"
    # home=Path("~/.tracemem"),  # storage directory (default: ~/.tracemem)
)

async with TraceMem(config=config) as tm:
    # Store a conversation
    await tm.add_message("conv-1", Message(role="user", content="Fix the auth bug in login.py"))
    await tm.add_message("conv-1", Message(
        role="assistant",
        content="I'll check the authentication code...",
        tool_calls=[ToolCall(id="c1", name="read_file", args={"path": "src/login.py"})],
    ))
    await tm.add_message("conv-1", Message(
        role="tool", content="def login(): ...", tool_call_id="c1",
    ))

    # Search similar past interactions
    results = await tm.search("authentication issues")
    for r in results:
        print(r.text, r.score)

    # Get full trajectory from a user message
    trajectory = await tm.get_trajectory(results[0].node_id)
    for step in trajectory.steps:
        print(step.node_type, step.text[:80], step.tool_uses)

Configuration

All settings use the TRACEMEM_ prefix:

export OPENAI_API_KEY="sk-..."                    # Required for embeddings
export TRACEMEM_GRAPH_STORE="kuzu"                # kuzu (default) or neo4j
export TRACEMEM_HOME="~/.tracemem"                # Storage directory
export TRACEMEM_EMBEDDING_MODEL="text-embedding-3-small"
export TRACEMEM_RERANKER="rrf"                    # rrf (default) or linear

TraceMemConfig Options

Graph Schema

Node Types

Relationship Types

Only 3 relationship types:

API Reference

TraceMem

Must be used as an async context manager:

async with TraceMem(config=config, embedder=embedder, resource_extractor=extractor) as tm:
    ...

Parameters:

• config (TraceMemConfig, optional): Configuration settings
• embedder (Embedder, optional): Custom embedder (defaults to OpenAI)
• resource_extractor (ResourceExtractor, optional): Custom resource URI extractor

await tm.add_message(conversation_id, message)

Add a single message to the knowledge graph.

await tm.import_trace(conversation_id, messages)

Import a full conversation from a list of Messages.

await tm.search(query, config=None)

Search for similar past interactions via vector similarity.

await tm.get_trajectory(node_id)

Get the full trajectory from a UserText node through all agent responses until the next user message.

Data Models

from tracemem_core import Message, ToolCall

Message(role="user", content="Fix the bug")
Message(role="assistant", content="...", tool_calls=[ToolCall(id="c1", name="read_file", args={"path": "auth.py"})])
Message(role="tool", content="file contents...", tool_call_id="c1")

Resource Extraction

The DefaultResourceExtractor extracts URIs from tool call arguments:

• File arguments: path, file_path, filepath, file, filename
• URL arguments: url, uri, endpoint

ToolCall(name="read_file", args={"path": "src/auth.py"})       # → file://src/auth.py
ToolCall(name="fetch", args={"url": "https://api.example.com"}) # → https://api.example.com

Custom Extractor

from tracemem_core import ResourceExtractor

class MyExtractor:
    def extract(self, tool_name: str, args: dict) -> str | None:
        if tool_name == "query_database":
            return f"db://{args.get('table')}"
        return None

async with TraceMem(resource_extractor=MyExtractor()) as tm:
    ...

Custom Embedder

from tracemem_core import Embedder, TraceMem

class MyEmbedder:
    @property
    def dimensions(self) -> int:
        return 768

    async def embed(self, text: str) -> list[float]:
        return [0.1] * 768

    async def embed_batch(self, texts: list[str]) -> list[list[float]]:
        return [await self.embed(t) for t in texts]

async with TraceMem(embedder=MyEmbedder()) as tm:
    ...

Adapters

LangChain

from tracemem_core.adapters.langchain import LangChainAdapter

adapter = LangChainAdapter()
messages = adapter.convert(langchain_messages)
await tm.import_trace("conv-123", messages)

Requires langchain-core: pip install langchain-core

Documentation

• Architecture
• Configuration
• Retrieval

License

MIT