import { getPool, initializeDatabase, isDbInitialized } from '../db/index.js';
import { generateEmbedding, generateEmbeddings } from './embedding.js';
import { sendLogMessage } from '../utils/logger.js';
// Reverse relation mappings for bidirectional maintenance
const REVERSE_RELATIONS = {
'uses': 'is_used_by',
'is_used_by': 'uses',
'depends_on': 'is_dependency_of',
'is_dependency_of': 'depends_on',
'contains': 'is_contained_in',
'is_contained_in': 'contains',
'calls': 'is_called_by',
'is_called_by': 'calls',
'owns': 'is_owned_by',
'is_owned_by': 'owns',
'creates': 'is_created_by',
'is_created_by': 'creates',
'manages': 'is_managed_by',
'is_managed_by': 'manages'
};
export const graphService = {
async createEntities(toolArgs) {
if (!isDbInitialized()) await initializeDatabase();
const { entities } = toolArgs;
const created = [];
// Generate embeddings for entities in batch
const entityTexts = entities.map(e =>
`${e.name}: ${e.entityType}. ${e.observations.join('. ')}`
);
let embeddings = [];
try {
embeddings = await generateEmbeddings(entityTexts);
} catch (error) {
sendLogMessage('warn', 'Failed to generate embeddings for entities, continuing without them', { error: error.message });
// Fallback to null embeddings if batch generation fails
embeddings = new Array(entities.length).fill(null);
}
for (let i = 0; i < entities.length; i++) {
const entity = entities[i];
const embedding = embeddings[i];
const query = `
INSERT INTO entities (name, entity_type, observations, embedding)
VALUES ($1, $2, $3, $4::vector)
ON CONFLICT (name)
DO UPDATE SET
entity_type = EXCLUDED.entity_type,
observations = entities.observations || EXCLUDED.observations,
embedding = COALESCE(EXCLUDED.embedding, entities.embedding),
updated_at = CURRENT_TIMESTAMP
RETURNING *
`;
const result = await getPool().query(query, [
entity.name,
entity.entityType,
entity.observations,
embedding ? `[${embedding.join(',')}]` : null
]);
created.push(result.rows[0]);
}
return created;
},
async createRelations(toolArgs) {
if (!isDbInitialized()) await initializeDatabase();
const { relations, bidirectional = true } = toolArgs;
const created = [];
const pool = getPool();
for (const rel of relations) {
const query = `
INSERT INTO relations (from_entity, to_entity, relation_type)
VALUES ($1, $2, $3)
ON CONFLICT (from_entity, to_entity, relation_type) DO NOTHING
RETURNING *
`;
const result = await pool.query(query, [
rel.from,
rel.to,
rel.relationType
]);
if (result.rows[0]) created.push(result.rows[0]);
// Auto-create reverse relation if bidirectional is enabled
if (bidirectional && REVERSE_RELATIONS[rel.relationType]) {
const reverseResult = await pool.query(query, [
rel.to,
rel.from,
REVERSE_RELATIONS[rel.relationType]
]);
if (reverseResult.rows[0]) {
sendLogMessage('debug', 'Created reverse relation', {
from: rel.to,
to: rel.from,
type: REVERSE_RELATIONS[rel.relationType]
});
}
}
}
return created;
},
async addObservations(toolArgs) {
if (!isDbInitialized()) await initializeDatabase();
const { observations } = toolArgs;
const results = [];
for (const obs of observations) {
const query = `
UPDATE entities
SET observations = observations || $2,
updated_at = CURRENT_TIMESTAMP
WHERE name = $1
RETURNING *
`;
const result = await getPool().query(query, [
obs.entityName,
obs.contents
]);
if (result.rows[0]) results.push(result.rows[0]);
}
return results;
},
async deleteEntities(toolArgs) {
if (!isDbInitialized()) await initializeDatabase();
const { entityNames } = toolArgs;
await getPool().query(
'DELETE FROM entities WHERE name = ANY($1)',
[entityNames]
);
return { deleted: entityNames };
},
async deleteObservations(toolArgs) {
if (!isDbInitialized()) await initializeDatabase();
const { deletions } = toolArgs;
for (const del of deletions) {
const query = `
UPDATE entities
SET observations = (
SELECT array_agg(elem)
FROM unnest(observations) elem
WHERE elem != ALL($2)
)
WHERE name = $1
`;
await getPool().query(query, [
del.entityName,
del.observations
]);
}
return { success: true };
},
async deleteRelations(toolArgs) {
if (!isDbInitialized()) await initializeDatabase();
const { relations } = toolArgs;
for (const rel of relations) {
await getPool().query(
'DELETE FROM relations WHERE from_entity = $1 AND to_entity = $2 AND relation_type = $3',
[rel.from, rel.to, rel.relationType]
);
}
return { success: true };
},
async readGraph(toolArgs) {
if (!isDbInitialized()) await initializeDatabase();
const entitiesResult = await getPool().query('SELECT name, entity_type, observations FROM entities');
const relationsResult = await getPool().query('SELECT from_entity as "from", to_entity as "to", relation_type as "relationType" FROM relations');
return {
entities: entitiesResult.rows.map(e => ({
name: e.name,
entityType: e.entity_type,
observations: e.observations
})),
relations: relationsResult.rows
};
},
async searchNodes(toolArgs) {
if (!isDbInitialized()) await initializeDatabase();
const { query: searchQuery, semantic = false, limit = 10 } = toolArgs;
let entities;
if (semantic) {
// Semantic search using embeddings
const embedding = await generateEmbedding(searchQuery);
const result = await getPool().query(
`SELECT name, entity_type, observations,
1 - (embedding <=> $1::vector) as similarity
FROM entities
WHERE embedding IS NOT NULL
ORDER BY similarity DESC
LIMIT $2`,
[`[${embedding.join(',')}]`, limit]
);
entities = result.rows;
} else {
// Keyword search
const lowerQuery = searchQuery.toLowerCase();
const result = await getPool().query(
`SELECT name, entity_type, observations FROM entities
WHERE LOWER(name) LIKE $1
OR LOWER(entity_type) LIKE $1
OR EXISTS (SELECT 1 FROM unnest(observations) obs WHERE LOWER(obs) LIKE $1)
LIMIT $2`,
[`%${lowerQuery}%`, limit]
);
entities = result.rows;
}
// Get relations for found entities
const entityNames = entities.map(e => e.name);
const relationsResult = await getPool().query(
`SELECT from_entity as "from", to_entity as "to", relation_type as "relationType"
FROM relations
WHERE from_entity = ANY($1) AND to_entity = ANY($1)`,
[entityNames]
);
return {
entities: entities.map(e => ({
name: e.name,
entityType: e.entity_type,
observations: e.observations,
similarity: e.similarity
})),
relations: relationsResult.rows
};
},
async openNodes(toolArgs) {
if (!isDbInitialized()) await initializeDatabase();
const { names } = toolArgs;
const entitiesResult = await getPool().query(
'SELECT name, entity_type, observations FROM entities WHERE name = ANY($1)',
[names]
);
const relationsResult = await getPool().query(
`SELECT from_entity as "from", to_entity as "to", relation_type as "relationType"
FROM relations
WHERE from_entity = ANY($1) OR to_entity = ANY($1)`,
[names]
);
return {
entities: entitiesResult.rows.map(e => ({
name: e.name,
entityType: e.entity_type,
observations: e.observations
})),
relations: relationsResult.rows
};
},
// Transitive relation inference: find paths A→B→C
async findTransitiveRelations(toolArgs) {
if (!isDbInitialized()) await initializeDatabase();
const { entityName, maxDepth = 3, relationType } = toolArgs;
sendLogMessage('info', 'Finding transitive relations', { entityName, maxDepth, relationType });
let query = `
WITH RECURSIVE transitive_relations AS (
-- Base case: direct relations from the starting entity
SELECT
from_entity,
to_entity,
relation_type,
1 as depth,
ARRAY[from_entity, to_entity] as path
FROM relations
WHERE from_entity = $1
${relationType ? 'AND relation_type = $3' : ''}
UNION ALL
-- Recursive case: follow the chain
SELECT
r.from_entity,
r.to_entity,
r.relation_type,
tr.depth + 1,
tr.path || r.to_entity
FROM relations r
INNER JOIN transitive_relations tr ON r.from_entity = tr.to_entity
WHERE tr.depth < $2
AND NOT r.to_entity = ANY(tr.path) -- Prevent cycles
${relationType ? 'AND r.relation_type = $3' : ''}
)
SELECT DISTINCT
from_entity as "from",
to_entity as "to",
relation_type as "relationType",
depth,
path
FROM transitive_relations
ORDER BY depth, from_entity, to_entity
`;
const params = [entityName, maxDepth];
if (relationType) params.push(relationType);
const result = await getPool().query(query, params);
// Get entity details for all involved entities
const allEntities = new Set();
result.rows.forEach(r => {
r.path.forEach(e => allEntities.add(e));
});
const entitiesResult = await getPool().query(
'SELECT name, entity_type, observations FROM entities WHERE name = ANY($1)',
[Array.from(allEntities)]
);
return {
startEntity: entityName,
maxDepth,
relations: result.rows,
entities: entitiesResult.rows.map(e => ({
name: e.name,
entityType: e.entity_type,
observations: e.observations
}))
};
},
// Export graph as Mermaid diagram
async exportMermaid(toolArgs) {
if (!isDbInitialized()) await initializeDatabase();
const { direction = 'TD' } = toolArgs || {};
const entitiesResult = await getPool().query('SELECT name, entity_type FROM entities');
const relationsResult = await getPool().query('SELECT from_entity, to_entity, relation_type FROM relations');
let mermaid = `graph ${direction}\n`;
// Add entity nodes with type labels
entitiesResult.rows.forEach(e => {
const safeName = e.name.replace(/[\s\-\.]/g, '_').replace(/[^a-zA-Z0-9_]/g, '');
mermaid += ` ${safeName}["${e.name}<br/><i>${e.entity_type}</i>"]\n`;
});
mermaid += '\n';
// Add relations as edges
relationsResult.rows.forEach(r => {
const safeFrom = r.from_entity.replace(/[\s\-\.]/g, '_').replace(/[^a-zA-Z0-9_]/g, '');
const safeTo = r.to_entity.replace(/[\s\-\.]/g, '_').replace(/[^a-zA-Z0-9_]/g, '');
mermaid += ` ${safeFrom} -->|${r.relation_type}| ${safeTo}\n`;
});
return { mermaid };
}
};
