Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 10 additions & 13 deletions src/memory/__tests__/knowledge-graph.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -426,7 +426,7 @@ describe('KnowledgeGraphManager', () => {
expect(JSON.parse(lines[1])).toHaveProperty('type', 'relation');
});

it('should strip type field from entities when loading from file', async () => {
it('should preserve JSONL type discriminator when loading from file', async () => {
// Create entities and relations (these get saved with type field)
await manager.createEntities([
{ name: 'Alice', entityType: 'person', observations: ['test observation'] },
Expand All @@ -451,26 +451,25 @@ describe('KnowledgeGraphManager', () => {
const manager2 = new KnowledgeGraphManager(testFilePath);
const graph = await manager2.readGraph();

// Verify loaded entities don't have type field
// Loaded graph mirrors JSONL discriminators for structured-output validation (#3074)
expect(graph.entities).toHaveLength(2);
graph.entities.forEach(entity => {
expect(entity).not.toHaveProperty('type');
expect(entity).toHaveProperty('type', 'entity');
expect(entity).toHaveProperty('name');
expect(entity).toHaveProperty('entityType');
expect(entity).toHaveProperty('observations');
});

// Verify loaded relations don't have type field
expect(graph.relations).toHaveLength(1);
graph.relations.forEach(relation => {
expect(relation).not.toHaveProperty('type');
expect(relation).toHaveProperty('type', 'relation');
expect(relation).toHaveProperty('from');
expect(relation).toHaveProperty('to');
expect(relation).toHaveProperty('relationType');
});
});

it('should strip type field from searchNodes results', async () => {
it('should include type discriminator in searchNodes results after reload', async () => {
await manager.createEntities([
{ name: 'Alice', entityType: 'person', observations: ['works at Acme'] },
]);
Expand All @@ -482,17 +481,16 @@ describe('KnowledgeGraphManager', () => {
const manager2 = new KnowledgeGraphManager(testFilePath);
const result = await manager2.searchNodes('Alice');

// Verify search results don't have type field
expect(result.entities).toHaveLength(1);
expect(result.entities[0]).not.toHaveProperty('type');
expect(result.entities[0]).toHaveProperty('type', 'entity');
expect(result.entities[0].name).toBe('Alice');

expect(result.relations).toHaveLength(1);
expect(result.relations[0]).not.toHaveProperty('type');
expect(result.relations[0]).toHaveProperty('type', 'relation');
expect(result.relations[0].from).toBe('Alice');
});

it('should strip type field from openNodes results', async () => {
it('should include type discriminator in openNodes results after reload', async () => {
await manager.createEntities([
{ name: 'Alice', entityType: 'person', observations: [] },
{ name: 'Bob', entityType: 'person', observations: [] },
Expand All @@ -505,14 +503,13 @@ describe('KnowledgeGraphManager', () => {
const manager2 = new KnowledgeGraphManager(testFilePath);
const result = await manager2.openNodes(['Alice', 'Bob']);

// Verify open results don't have type field
expect(result.entities).toHaveLength(2);
result.entities.forEach(entity => {
expect(entity).not.toHaveProperty('type');
expect(entity).toHaveProperty('type', 'entity');
});

expect(result.relations).toHaveLength(1);
expect(result.relations[0]).not.toHaveProperty('type');
expect(result.relations[0]).toHaveProperty('type', 'relation');
});
});
});
6 changes: 6 additions & 0 deletions src/memory/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -48,12 +48,14 @@ let MEMORY_FILE_PATH: string;

// We are storing our memory using entities, relations, and observations in a graph structure
export interface Entity {
type?: "entity";
name: string;
entityType: string;
observations: string[];
}

export interface Relation {
type?: "relation";
from: string;
to: string;
relationType: string;
Expand All @@ -76,13 +78,15 @@ export class KnowledgeGraphManager {
const item = JSON.parse(line);
if (item.type === "entity") {
graph.entities.push({
type: "entity",
name: item.name,
entityType: item.entityType,
observations: item.observations
});
}
if (item.type === "relation") {
graph.relations.push({
type: "relation",
from: item.from,
to: item.to,
relationType: item.relationType
Expand Down Expand Up @@ -241,12 +245,14 @@ let knowledgeGraphManager: KnowledgeGraphManager;

// Zod schemas for entities and relations
const EntitySchema = z.object({
type: z.literal("entity").optional().describe('Discriminator when parsed from JSONL; omit when creating entities'),
name: z.string().describe("The name of the entity"),
entityType: z.string().describe("The type of the entity"),
observations: z.array(z.string()).describe("An array of observation contents associated with the entity")
});

const RelationSchema = z.object({
type: z.literal("relation").optional().describe('Discriminator when parsed from JSONL; omit when creating relations'),
from: z.string().describe("The name of the entity where the relation starts"),
to: z.string().describe("The name of the entity where the relation ends"),
relationType: z.string().describe("The type of the relation")
Expand Down
Loading