|
|
|
@@ -1,9 +1,7 @@
|
|
|
|
|
const semantic = require('../semantic')
|
|
|
|
|
const { getEnv, SERVICES, formatEpisodeText, ENTITIES, logger } = require('@nexusai/shared');
|
|
|
|
|
const { getEnv, SERVICES, formatEpisodeText, ENTITIES, logger, utilityInference } = require('@nexusai/shared');
|
|
|
|
|
const { upsertEntity, upsertRelationship, linkEntityToEpisode } = require('./index');
|
|
|
|
|
|
|
|
|
|
const EXTRACTION_URL = getEnv('EXTRACTION_URL', 'http://localhost:11434');
|
|
|
|
|
const EXTRACTION_MODEL = getEnv('EXTRACTION_MODEL', 'qwen2.5:3b'); // ChatML format — see buildExtractionPrompt
|
|
|
|
|
const EMBEDDING_SERVICE_URL = getEnv('EMBEDDING_SERVICE_URL', SERVICES.EMBEDDING_URL);
|
|
|
|
|
|
|
|
|
|
const ENTITY_TYPES = ENTITIES.TYPES;
|
|
|
|
@@ -28,10 +26,9 @@ function mentionedIn(name, haystack) {
|
|
|
|
|
return norm(haystack).includes(norm(name));
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// NOTE: This prompt uses ChatML format (<|im_start|> / <|im_end|> tags), which is
|
|
|
|
|
// specific to qwen-family models. If EXTRACTION_MODEL is changed to a Llama-family
|
|
|
|
|
// or other model, this format will need to change — most alternatives use either
|
|
|
|
|
// plain text or [INST] / <<SYS>> tags. Silent degradation is likely if mismatched.
|
|
|
|
|
// Returns { system, user } for utilityInference. The model's prompt template
|
|
|
|
|
// (ChatML for qwen, etc.) is applied by Ollama via the inference service's
|
|
|
|
|
// /utility/complete route — no template tags belong in this file.
|
|
|
|
|
function buildExtractionPrompt(userMessage, aiResponse, knownEntities = []) {
|
|
|
|
|
const knownBlock = knownEntities.length > 0
|
|
|
|
|
? [
|
|
|
|
@@ -41,33 +38,30 @@ function buildExtractionPrompt(userMessage, aiResponse, knownEntities = []) {
|
|
|
|
|
].join('\n')
|
|
|
|
|
: '';
|
|
|
|
|
|
|
|
|
|
return [
|
|
|
|
|
'<|im_start|>system',
|
|
|
|
|
'You are a named entity and relationship extractor. You output only valid JSON.',
|
|
|
|
|
'<|im_end|>',
|
|
|
|
|
'<|im_start|>user',
|
|
|
|
|
'Read the conversation below and extract all named entities and the relationships between them.',
|
|
|
|
|
`Entity types: ${ENTITY_TYPES.join(', ')}`,
|
|
|
|
|
'Use "character" for any fictional, game, or media characters (e.g. characters from anime, games, books, TV shows, movies)',
|
|
|
|
|
'Use "person" only for real people',
|
|
|
|
|
'For each entity provide:',
|
|
|
|
|
' "name": short proper noun only (max 4 words)',
|
|
|
|
|
' "type": one of the valid types',
|
|
|
|
|
' "notes": one specific sentence about this entity based on the conversation',
|
|
|
|
|
'For relationships, use snake_case verb labels (e.g. works_on, manages, uses, knows, located_in, part_of, created_by).',
|
|
|
|
|
'Only include relationships between entities you have listed above.',
|
|
|
|
|
'The known-entities list below is ONLY for consistent spelling and types. Do NOT output an entity unless it actually appears in the conversation.',
|
|
|
|
|
'Return this exact JSON structure:',
|
|
|
|
|
'{ "entities": [{"name": "...", "type": "...", "notes": "..."}], "relationships": [{"from": "...", "fromType": "...", "to": "...", "toType": "...", "label": "...", "notes": "..."}] }',
|
|
|
|
|
'',
|
|
|
|
|
knownBlock,
|
|
|
|
|
'--- CONVERSATION ---',
|
|
|
|
|
`User: ${userMessage}`,
|
|
|
|
|
`Assistant: ${aiResponse}`,
|
|
|
|
|
'--- END CONVERSATION ---',
|
|
|
|
|
'<|im_end|>',
|
|
|
|
|
'<|im_start|>assistant',
|
|
|
|
|
].join('\n');
|
|
|
|
|
return {
|
|
|
|
|
system: 'You are a named entity and relationship extractor. You output only valid JSON.',
|
|
|
|
|
user: [
|
|
|
|
|
'Read the conversation below and extract all named entities and the relationships between them.',
|
|
|
|
|
`Entity types: ${ENTITY_TYPES.join(', ')}`,
|
|
|
|
|
'Use "character" for any fictional, game, or media characters (e.g. characters from anime, games, books, TV shows, movies)',
|
|
|
|
|
'Use "person" only for real people',
|
|
|
|
|
'For each entity provide:',
|
|
|
|
|
' "name": short proper noun only (max 4 words)',
|
|
|
|
|
' "type": one of the valid types',
|
|
|
|
|
' "notes": one specific sentence about this entity based on the conversation',
|
|
|
|
|
'For relationships, use snake_case verb labels (e.g. works_on, manages, uses, knows, located_in, part_of, created_by).',
|
|
|
|
|
'Only include relationships between entities you have listed above.',
|
|
|
|
|
'The known-entities list below is ONLY for consistent spelling and types. Do NOT output an entity unless it actually appears in the conversation.',
|
|
|
|
|
'Return this exact JSON structure:',
|
|
|
|
|
'{ "entities": [{"name": "...", "type": "...", "notes": "..."}], "relationships": [{"from": "...", "fromType": "...", "to": "...", "toType": "...", "label": "...", "notes": "..."}] }',
|
|
|
|
|
'',
|
|
|
|
|
knownBlock,
|
|
|
|
|
'--- CONVERSATION ---',
|
|
|
|
|
`User: ${userMessage}`,
|
|
|
|
|
`Assistant: ${aiResponse}`,
|
|
|
|
|
'--- END CONVERSATION ---',
|
|
|
|
|
].join('\n'),
|
|
|
|
|
};
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
async function embedEntity(entity) {
|
|
|
|
@@ -91,30 +85,16 @@ async function extractAndStoreEntities(userMessage, aiResponse, episodeId=null,
|
|
|
|
|
// Fetch existing entities to guide the model toward consistent name/type pairs
|
|
|
|
|
const db = require('../db').getDB();
|
|
|
|
|
const knownEntities = db.prepare(`SELECT name, type FROM entities ORDER BY rowid DESC LIMIT 20`).all();
|
|
|
|
|
const prompt = buildExtractionPrompt(userMessage, aiResponse, knownEntities);
|
|
|
|
|
const { system, user } = buildExtractionPrompt(userMessage, aiResponse, knownEntities);
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
const res = await fetch(`${EXTRACTION_URL}/api/generate`, {
|
|
|
|
|
method: 'POST',
|
|
|
|
|
headers: { 'Content-Type': 'application/json' },
|
|
|
|
|
body: JSON.stringify({
|
|
|
|
|
model: EXTRACTION_MODEL,
|
|
|
|
|
prompt: prompt,
|
|
|
|
|
stream: false,
|
|
|
|
|
format: 'json',
|
|
|
|
|
options: {
|
|
|
|
|
temperature: ENTITIES.TEMPERATURE,
|
|
|
|
|
num_predict: ENTITIES.NUM_PREDICT,
|
|
|
|
|
},
|
|
|
|
|
}),
|
|
|
|
|
signal: AbortSignal.timeout(60_000),
|
|
|
|
|
const raw = await utilityInference({
|
|
|
|
|
system,
|
|
|
|
|
user,
|
|
|
|
|
json: true,
|
|
|
|
|
temperature: ENTITIES.TEMPERATURE,
|
|
|
|
|
maxTokens: ENTITIES.NUM_PREDICT,
|
|
|
|
|
});
|
|
|
|
|
|
|
|
|
|
if (!res.ok) throw new Error(`Ollama responded ${res.status}`);
|
|
|
|
|
|
|
|
|
|
const data = await res.json();
|
|
|
|
|
const raw = data.response?.trim() ?? '';
|
|
|
|
|
|
|
|
|
|
const jsonMatch = raw.match(/\{[\s\S]*\}/);
|
|
|
|
|
if (!jsonMatch) {
|
|
|
|
|
logger.warn('[entities] No JSON object found in response');
|
|
|
|
|