Skip to content

Module Reference

This document provides a detailed breakdown of every module, class, and function in the Cognition Engines codebase.


1. Core Library — src/cognition_engines/

Section titled “1. Core Library — src/cognition_engines/”

The heart of the decision memory system. Provides vector-based similarity search over decisions using ChromaDB and Google Gemini embeddings.

Function Description
load_gemini_key() Loads Gemini API key from secrets files if not present in environment
api_request(method, url, data) Makes HTTP requests to ChromaDB REST API
generate_embedding(text) Generates a 768-dimensional embedding vector via Gemini text-embedding-004
get_api_base() Returns the ChromaDB API base URL from environment
get_or_create_collection() Ensures the cognition_decisions collection exists in ChromaDB
decision_to_text(decision) Converts a decision dict into searchable text (title + category + context + reasons)
decision_id(decision) Generates a deterministic MD5 hash ID from decision title
get_index() Returns the singleton SemanticIndex instance
Method Description
__init__() Initializes with collection_id = None
ensure_collection() Lazily creates/gets the ChromaDB collection
index_decision(decision) Indexes a single decision: generates embedding, constructs metadata, upserts to ChromaDB
index_decisions(decisions) Batch indexes a list of decisions, returns count
query(context, n_results, category, min_confidence) Generates embedding for query text, performs filtered vector search, returns ranked results
count() Returns the number of indexed decisions in the collection

Metadata stored per decision:

  • title, category, confidence, stakes, date, status, outcome, project, feature, pr, reason_types

The main guardrail evaluation engine. Loads YAML guardrail definitions and evaluates them against decision contexts.

Class Fields Description
GuardrailCondition field, operator, value A condition that must match for a guardrail to apply (e.g., stakes == "high")
GuardrailRequirement field, expected A requirement that must be met (e.g., code_review_completed == true)
Guardrail id, description, conditions, requirements, scope, action, message Full guardrail definition
GuardrailResult guardrail_id, passed, action, message, failed_requirements Result of evaluating one guardrail
Method Description
load_from_yaml(content) Parses YAML string and loads guardrail definitions. Returns count loaded
load_from_file(path) Loads guardrails from a single YAML file
load_from_directory(directory) Recursively loads all .yaml/.yml files from a directory
evaluate(context) Evaluates all guardrails against a context dict, returns list of GuardrailResult
check(context) Convenience method: returns (allowed: bool, results: list)
list_guardrails() Returns list of all loaded guardrail definitions as dicts

Condition operators: =, !=, <, >, <=, >=, in, not in

Function Description
parse_condition(field, value) Parses a condition from YAML format, detecting operator from string prefix
parse_guardrail(data) Parses a full guardrail from YAML dict (supports flat condition_* and nested formats)
get_engine() Returns the singleton GuardrailEngine instance
load_default_guardrails() Loads guardrails from guardrails/ directory

Advanced condition evaluators for the v2 guardrail format. Provides pluggable evaluation strategies.

Class Type Description
ConditionEvaluator Protocol Abstract protocol for condition evaluators
FieldCondition Evaluator Simple field comparison with extended operators
SemanticCondition Evaluator Checks semantic similarity to past decisions matching criteria. Fields: query_field, threshold, filter_outcome, filter_since_days, min_matches
TemporalCondition Evaluator Time-window check: “Was a similar decision made within N hours?”
AggregateCondition Evaluator Statistical check across decision history: “Is category success rate below 50%?”
CompoundCondition Evaluator AND/OR logical composition of multiple conditions

parse_condition_v2(condition_def) — Factory function that creates the appropriate evaluator from a YAML condition definition.

Audit trail system for guardrail evaluations.

Class Description
GuardrailEvaluation Record of a single guardrail check: guardrail_id, matched, passed, action, message
AuditRecord Complete audit for one decision: list of evaluations, overall_allowed, optional override with reason
AuditLog Manager that creates records, saves to JSON files, queries violations, and computes aggregate stats

Output formats: JSON files (audit/YYYY-MM-DD-<decision_id>.json) and embeddable YAML blocks.


Analyzes decision history for patterns, calibration accuracy, and anti-patterns.

Class Description
CalibrationBucket Confidence range bucket: computes predicted rate, actual success rate, and Brier score
CategoryStats Per-category statistics: count, average confidence, success rate
AntiPattern Detected anti-pattern: type, description, severity, affected decisions
Method Description
load_from_directory(directory) Recursively loads all YAML decision files
calibration_report() Generates Brier scores for 5 confidence buckets (0-0.2, 0.2-0.4, …, 0.8-1.0)
category_analysis() Success rates and patterns per category, identifies concerning categories
detect_antipatterns() Detects: overcalibration, flip-flopping, anchoring, blind spots, hot-hand fallacy
full_report() Combines calibration + category + antipattern reports

Function Description
lifespan(app) Async context manager: loads config, initializes AuthManager, creates CstpDispatcher, registers methods, initializes MCP StreamableHTTPSessionManager and runs it within the lifespan context
create_app(config) Factory: creates FastAPI app with CORS, lifespan, routes, and MCP mount at /mcp
_mount_mcp(app) Mounts the MCP Streamable HTTP handler as a raw ASGI app at /mcp; returns 503 if MCP SDK not installed
_register_routes(app) Registers /health, /.well-known/agent.json, and POST /cstp
run_server(host, port, config_path) Entry point: loads config, creates app, runs uvicorn

Exposes CSTP capabilities as MCP tools for native integration with MCP-compliant agents (Claude Desktop, Claude Code, OpenClaw, etc.).

Component Description
mcp_app Server("cstp-decisions") instance — importable for mounting in ASGI apps
list_tools() Returns 5 Tool definitions with JSON Schema auto-generated from Pydantic models
call_tool(name, arguments) Dispatches tool calls to _handle_* functions; returns TextContent with JSON result
_handle_query_decisions() Validates input via QueryDecisionsInput, calls query_service.query_decisions()
_handle_check_action() Validates input via CheckActionInput, calls guardrails_service.evaluate_guardrails()
_handle_log_decision() Validates input via LogDecisionInput, calls decision_service.record_decision()
_handle_review_outcome() Validates input via ReviewOutcomeInput, calls decision_service.review_decision()
_handle_get_stats() Validates input via GetStatsInput, calls calibration_service.get_calibration()
run_stdio() Runs the MCP server with stdio transport (async with stdio_server())
main() Entry point: asyncio.run(run_stdio())

Transports:

  • stdiopython -m a2a.mcp_server (local or docker exec -i cstp python -m a2a.mcp_server)
  • Streamable HTTP — Mounted at /mcp on port 9991 via StreamableHTTPSessionManager in server.py lifespan

Pydantic models that define the JSON Schema MCP clients see during tool discovery. They map to existing CSTP dataclass models but use Pydantic for automatic schema generation required by the MCP protocol.

Schema MCP Tool Key Fields
QueryDecisionsInput query_decisions query (str), limit (1–50), retrieval_mode (semantic/keyword/hybrid), filters (QueryFiltersInput)
QueryFiltersInput (nested) category, stakes, project, has_outcome
CheckActionInput check_action description (str), category, stakes (low/medium/high/critical), confidence (0.0–1.0)
LogDecisionInput log_decision decision (str), confidence (float), category, stakes, context, reasons (ReasonInput[]), tags, project, feature, pr
ReasonInput (nested) type (authority/analogy/analysis/pattern/intuition), text (str)
ReviewOutcomeInput review_outcome id (str), outcome (success/partial/failure/abandoned), actual_result, lessons, notes
GetStatsInput get_stats category, project, window (30d/60d/90d/all)

2.4 config.py — Configuration Management

Section titled “2.4 config.py — Configuration Management”
Class Description
AuthToken Agent + token pair
AuthConfig enabled flag + list of AuthToken; validate_token() with constant-time comparison
AgentConfig Agent identity: name, description, version, URL, contact
ServerConfig HTTP settings: host, port, CORS origins
Config Composite config with from_yaml(path), from_env(), _from_dict(data) class methods

Config loading priority: YAML file → Environment variables → Defaults

Component Description
AuthManager Wraps Config, validates bearer tokens, returns agent ID
verify_bearer_token() FastAPI Depends function for route-level auth
set_auth_manager() / get_auth_manager() Global singleton management
File Classes
jsonrpc.py JsonRpcRequest, JsonRpcResponse, JsonRpcError, error codes (PARSE_ERROR, INVALID_REQUEST, METHOD_NOT_FOUND, INTERNAL_ERROR)
agent_card.py AgentCard, AgentCapabilities for /.well-known/agent.json
health.py HealthResponse with status, version, uptime, timestamp

Component Description
CstpDispatcher Registry of method name → async handler; dispatches requests, catches errors, returns JSON-RPC responses
register_methods(dispatcher) Registers all 9 method handlers: queryDecisions, checkGuardrails, listGuardrails, recordDecision, reviewDecision, getCalibration, attributeOutcomes, checkDrift, reindex
Custom error codes QUERY_FAILED (-32003), RATE_LIMITED (-32002), GUARDRAIL_EVAL_FAILED (-32004), RECORD_FAILED (-32005), REVIEW_FAILED (-32006), DECISION_NOT_FOUND (-32007), ATTRIBUTION_FAILED (-32008)
Component Description
QueryResult Single result: id, title, category, confidence, distance
QueryResponse Wrapper with results list, query, timing
query_decisions() Full query pipeline: embedding, ChromaDB search, metadata filtering, optional BM25 hybrid
load_all_decisions() Loads YAML files from disk for BM25 indexing

3.3 decision_service.py — Decision Recording & Review

Section titled “3.3 decision_service.py — Decision Recording & Review”
Component Description
BridgeDefinition Dataclass: structure, function, tolerance, enforcement, prevention
Reason Decision reason with type, text, strength
PreDecisionProtocol Tracks whether query was run and guardrails were checked before recording
ProjectContext Project, feature, PR, file, line, commit associations
ReasoningStep Step in a reasoning trace: step number, thought, output, confidence, tags
RecordDecisionRequest Full request with validation: decision text, confidence, category, stakes, reasons, review_in
RecordDecisionResponse Success indicator, generated ID, file path, index status
ReviewDecisionRequest Review request: decision ID, outcome, actual result, lessons
review_decision() Loads decision YAML, updates with outcome and review metadata, reindexes in ChromaDB
build_decision_yaml() Constructs the YAML dictionary structure
write_decision_file() Writes to decisions/YYYY/MM/YYYY-MM-DD-decision-<id>.yaml
generate_embedding(text) Gemini embedding for ChromaDB indexing
ensure_collection_exists() ChromaDB collection management

3.4 calibration_service.py — Confidence Calibration

Section titled “3.4 calibration_service.py — Confidence Calibration”
Component Description
GetCalibrationRequest Filters: agent, category, stakes, date range, window, project, feature
CalibrationResult Overall: Brier score, accuracy, calibration gap, interpretation
ConfidenceBucket Per-bucket: decisions, success rate, expected rate, gap, interpretation
ConfidenceStats Distribution stats: mean, std_dev, min, max, bucket counts
CalibrationRecommendation Actionable recommendation with type, message, severity
get_reviewed_decisions() Loads reviewed decisions matching filters from YAML files
calculate_calibration() Computes Brier score: mean((confidence - outcome)²)
calculate_buckets() Splits into 5 confidence buckets and computes per-bucket stats
calculate_confidence_stats() Habituation detection: identifies low-variance confidence patterns
generate_recommendations() Produces actionable advice based on calibration gaps

3.5 attribution_service.py — Outcome Attribution

Section titled “3.5 attribution_service.py — Outcome Attribution”
Component Description
AttributeOutcomesRequest Project, since date, stability days, dry_run flag
find_pending_decisions() Finds pending decisions for a project
is_pr_stable() Checks if PR is older than stability window (simplified)
update_decision_outcome() Atomic file update: marks as reviewed with outcome + reason
attribute_outcomes() Main pipeline: find pending → check stability → update files

3.6 drift_service.py — Calibration Drift Detection

Section titled “3.6 drift_service.py — Calibration Drift Detection”
Component Description
CheckDriftRequest Thresholds for Brier and accuracy, category/project filters
DriftAlert Alert: type (brier_degradation/accuracy_drop), recent vs. historical values, change %, severity
check_drift() Compares 30-day calibration against 90-day+ baseline
generate_drift_recommendations() Actionable recommendations based on drift alerts

3.7 guardrails_service.py — Guardrail Evaluation (CSTP)

Section titled “3.7 guardrails_service.py — Guardrail Evaluation (CSTP)”
Component Description
evaluate_guardrails(context) Loads guardrails (cached 5 min), evaluates, returns violations/warnings
list_guardrails(scope) Lists active guardrails, optionally filtered by scope
log_guardrail_check() Structured audit logging

3.8 reindex_service.py — Collection Rebuild

Section titled “3.8 reindex_service.py — Collection Rebuild”
Component Description
reindex_decisions() Full pipeline: delete collection → create → load decisions → generate embeddings → batch insert
ReindexResult Success, count indexed, errors, duration
Component Description
BM25Index Wraps rank-bm25 BM25Okapi algorithm
BM25Index.from_decisions() Builds index from decision dicts
BM25Index.search() Returns ranked (doc_id, score) pairs
tokenize(text) Simple word tokenization with lowercasing
merge_results() Weighted merge of semantic + keyword results (default 70/30)
get_cached_index() 5-minute TTL cache with count-based invalidation

3.10 deliberation_tracker.py — Deliberation Traces

Section titled “3.10 deliberation_tracker.py — Deliberation Traces”

Tracks input/reasoning steps across API calls to build full deliberation traces.

Component Description
TrackedInput Dataclass: type (query/check), content, timestamp, source
TrackerSession Per-agent session state with inputs list and start time
DeliberationTracker Singleton manager for active deliberation sessions
track_query(agent_id, query) Hooks into queryDecisions to record search inputs
track_check(agent_id, action) Hooks into checkGuardrails to record constraint inputs
track_lookup(agent_id, id) Hooks into getDecision to record reference inputs
auto_attach_deliberation(agent_id) Returns and clears tracked inputs for recordDecision

3.11 bridge_extractor.py — Bridge Auto-Extraction

Section titled “3.11 bridge_extractor.py — Bridge Auto-Extraction”

Extracts structure/function pairs from decision text when not explicitly provided.

Component Description
auto_extract_bridge(text, context) Main entry point: heuristic extraction pipeline
_score_as_function(text) Scorer: how likely is text to be a function/purpose?
_score_as_structure(text) Scorer: how likely is text to be a structure/pattern?
Component Description
maybe_auto_extract_bridge(req) Wraps extraction logic for use in decision_service
Model Used By
QueryFilters cstp.queryDecisions — category, confidence, stakes, status, project, feature, PR, has_outcome
QueryDecisionsRequest cstp.queryDecisions — query text, bridge_side (structure/function), filters, limit
DecisionSummary Query results — id, title, category, confidence, distance, reasons
QueryDecisionsResponse Wrapper with decisions, total, timing, retrieval mode, scores
ActionContext cstp.checkGuardrails — description, category, stakes, confidence, context dict
CheckGuardrailsRequest Action + agent info
GuardrailViolation Block/warn result with id, name, message, severity, suggestion
CheckGuardrailsResponse allowed flag, violations, warnings, evaluated count

Route Method Description
/ GET Redirects to /decisions
/health GET Health check (no auth)
/decisions GET Paginated decision list with category/status filters
/decisions/<id> GET Decision detail view
/decisions/<id>/review GET/POST Review form for recording outcomes
/calibration GET Calibration dashboard with drift detection

Async HTTP client for communicating with the CSTP server. Uses httpx for non-blocking requests.

Method Description
health_check() Check CSTP server health
list_decisions() Query decisions with pagination
get_decision(id) Get single decision by ID
review_decision() Submit decision review
get_calibration() Get calibration statistics
check_drift() Check for calibration drift

Command Description
cognition index <dir> Index all YAML decisions from a directory
cognition query <context> Search for similar decisions
cognition check --stakes high --confidence 0.8 Evaluate guardrails for a context
cognition guardrails List all loaded guardrails
cognition count Count indexed decisions
cognition patterns calibration Confidence calibration report (Brier scores)
cognition patterns categories Category success analysis
cognition patterns antipatterns Detect decision anti-patterns
cognition patterns full Complete pattern report (JSON)

Test File Coverage
test_a2a_server.py FastAPI app creation, health endpoint, agent card, CSTP dispatch
test_decision_service.py Decision recording, YAML generation, validation
test_query_service.py Semantic query pipeline
test_guardrails.py Core guardrail engine evaluation
test_guardrails_service.py CSTP guardrail service
test_evaluators.py v2 condition evaluators
test_audit.py Audit trail creation and querying
test_patterns.py Pattern detection and calibration
test_calibration_service.py Calibration computation and recommendations
test_attribution_service.py Outcome attribution pipeline
test_semantic_index.py ChromaDB semantic index
test_config_env.py Configuration loading
test_f002_query_decisions.py F002 feature: query decisions
test_f003_check_guardrails.py F003 feature: check guardrails
test_f007_record_decision.py F007 feature: record decision
test_f008_review_decision.py F008 feature: review decision
test_f009_get_calibration.py F009 feature: get calibration