Skip to main content

Knowledge graph memory system for AI agents

Project description

TraceMem Core

PyPI Python 3.12+

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

Option Type Default Description
graph_store "kuzu" | "neo4j" "kuzu" Graph backend
home Path ~/.tracemem Storage directory for graph + vectors
embedding_model str "text-embedding-3-small" OpenAI embedding model
embedding_dimensions int 1536 Embedding vector dimensions
openai_api_key str None OpenAI API key (or use env var)
reranker str "rrf" Reranker strategy (rrf or linear)
neo4j_uri str "bolt://localhost:7687" Neo4j URI (only if graph_store="neo4j")
neo4j_user str "neo4j" Neo4j username
neo4j_password str "password" Neo4j password
neo4j_database str "neo4j" Neo4j database name
namespace str None Namespace for multi-user isolation (Neo4j)

Graph Schema

Node Types

Node Key Properties Description
UserText id, text, conversation_id, turn_index User messages (searchable via vector store)
AgentText id, text, tool_uses, conversation_id, turn_index Agent responses and tool invocations
ResourceVersion id, uri, content_hash Snapshot of a resource at a point in time
Resource id, uri, current_content_hash Canonical resource identity (deduplicated by URI)

Relationship Types

Only 3 relationship types:

Type From → To Description
MESSAGE UserText ↔ AgentText Conversation flow (chronological)
TOOL_USE AgentText → ResourceVersion Agent used a tool that touched this resource
VERSION_OF ResourceVersion → Resource Links snapshot to canonical resource

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

License

MIT

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

tracemem_core-0.1.3.tar.gz (98.8 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

tracemem_core-0.1.3-py3-none-any.whl (41.0 kB view details)

Uploaded Python 3

File details

Details for the file tracemem_core-0.1.3.tar.gz.

File metadata

  • Download URL: tracemem_core-0.1.3.tar.gz
  • Upload date:
  • Size: 98.8 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.7

File hashes

Hashes for tracemem_core-0.1.3.tar.gz
Algorithm Hash digest
SHA256 72eac3faa02eccaa9ca37b52f3b40fdcb5498b5249890f34101250f62ac8281b
MD5 4e4dab7a1f449813dab66228f37ada40
BLAKE2b-256 8cbd850067068afc4a29d120434878fc8e87f8396f29bef78f74550b82b7bd41

See more details on using hashes here.

Provenance

The following attestation bundles were made for tracemem_core-0.1.3.tar.gz:

Publisher: release.yml on itay1542/tracemem

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file tracemem_core-0.1.3-py3-none-any.whl.

File metadata

  • Download URL: tracemem_core-0.1.3-py3-none-any.whl
  • Upload date:
  • Size: 41.0 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/6.1.0 CPython/3.13.7

File hashes

Hashes for tracemem_core-0.1.3-py3-none-any.whl
Algorithm Hash digest
SHA256 9a004237d4fe621fdd3e89fb2d7be89f364b56f0ad36b49cae33b909afe7b4df
MD5 a0a520de58e9d2d591155d9e77425e20
BLAKE2b-256 784a424aa2b65e433d62398fce8ac18c221a6281d72cc8aa3ccebad4e20ed419

See more details on using hashes here.

Provenance

The following attestation bundles were made for tracemem_core-0.1.3-py3-none-any.whl:

Publisher: release.yml on itay1542/tracemem

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page