-
Notifications
You must be signed in to change notification settings - Fork 104
fix(nams): handle auto-merged entity creates in both SDKs; fix crewai smoke test #185
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from 2 commits
cc45ed7
0e406b0
f449fce
dbaa037
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -35,7 +35,7 @@ | |
| from typing import TYPE_CHECKING, Any | ||
| from uuid import UUID | ||
|
|
||
| from neo4j_agent_memory.core.exceptions import NotSupportedError | ||
| from neo4j_agent_memory.core.exceptions import MemoryError, NotSupportedError | ||
| from neo4j_agent_memory.memory.long_term import ( | ||
| Entity, | ||
| Fact, | ||
|
|
@@ -167,10 +167,16 @@ def _normalize_entity(payload: dict[str, Any] | None) -> dict[str, Any]: | |
| doesn't provide — we default them so Pydantic parsing succeeds. | ||
| NAMS types come back lowercase; uppercase them so package-side | ||
| consumers see the same type values they sent. | ||
|
|
||
| Explicit ``null`` fields are dropped so model defaults apply — NAMS | ||
| projects unset node properties as JSON ``null`` (e.g. ``confidence`` | ||
| on a manually-created entity), which would otherwise fail parsing | ||
| for non-optional fields like ``confidence: float``. | ||
| """ | ||
| from datetime import datetime, timezone | ||
|
|
||
| data = snakeize_keys(payload) if isinstance(payload, dict) else {} | ||
| data = {k: v for k, v in data.items() if v is not None} | ||
| if "created_at" not in data: | ||
| data["created_at"] = datetime.now(timezone.utc).isoformat() | ||
| if "metadata" not in data: | ||
|
|
@@ -215,6 +221,13 @@ async def add_entity( | |
| ``{name, type, description?}`` per spec. Bolt-only kwargs | ||
| (``subtype``, ``aliases``, ``attributes``, ``confidence``, | ||
| ``deduplicate``, ``geocode``, ``enrich``, etc.) are silently dropped. | ||
|
|
||
| NAMS resolves-before-create: when the name is close enough to an | ||
| existing entity, the server merges onto it and responds | ||
| ``{id, resolution: "merged", merged_into, confidence}`` — no | ||
| ``name``/``type``. We follow up with ``GET /entities/{id}`` and | ||
| return the canonical merged-into entity (falling back to the | ||
| request's name/type if that read fails). | ||
| """ | ||
| et = entity_type or kwargs.get("type") or kwargs.get("label") | ||
| if et is None: | ||
|
|
@@ -227,8 +240,44 @@ async def add_entity( | |
| } | ||
| ) | ||
| payload = await self._transport.request(_SPEC_ADD_ENTITY, json=body) | ||
| if isinstance(payload, dict) and payload.get("resolution") == "merged": | ||
| payload = await self._resolve_merged_entity(payload, name=name, nams_type=body["type"]) | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
|
||
| return payload_to_model(_normalize_entity(payload), Entity) | ||
|
|
||
| async def _resolve_merged_entity( | ||
| self, | ||
| payload: dict[str, Any], | ||
| *, | ||
| name: str, | ||
| nams_type: str, | ||
| ) -> dict[str, Any]: | ||
| """Turn a ``resolution: "merged"`` create response into entity fields. | ||
|
|
||
| The merged response carries only ``{id, merged_into, confidence}``, | ||
| so we fetch the canonical entity it merged into. If that read | ||
| fails (or comes back without a name), synthesize a minimal record | ||
| from the request so the caller still gets a parseable Entity. | ||
| """ | ||
| merged_id = payload.get("merged_into") or payload.get("mergedInto") or payload.get("id") | ||
| if merged_id: | ||
| try: | ||
| detail = await self._transport.request( | ||
| _SPEC_GET_ENTITY, | ||
| path_params={"entity_id": _to_str(merged_id)}, | ||
| ) | ||
| except MemoryError: | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Effect: a rate-limited or auth-failed canonical read returns the synthesized entity with the requested name — no error, no log. Needs a |
||
| detail = None | ||
| if isinstance(detail, dict) and detail.get("name"): | ||
| return detail | ||
| return _drop_none( | ||
| { | ||
| "id": merged_id, | ||
| "name": name, | ||
| "type": nams_type, | ||
| "confidence": payload.get("confidence"), | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Reusing the merge score as
Drop it (model default 1.0) or move to |
||
| } | ||
| ) | ||
|
|
||
| async def add_preference(self, category: str, preference: str, **kwargs: Any) -> Preference: | ||
| raise NotSupportedError( | ||
| backend="nams", | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -5,7 +5,7 @@ | |
| * entity feedback, history, merge-by-id, graph view, and provenance. | ||
| */ | ||
|
|
||
| import { ValidationError } from "../errors.js"; | ||
| import { MemoryError, ValidationError } from "../errors.js"; | ||
| import type { Transport } from "../transport/index.js"; | ||
| import type { | ||
| AddRelationshipOptions, | ||
|
|
@@ -54,6 +54,18 @@ interface WireEntityRelRef { | |
| properties?: Record<string, unknown>; | ||
| } | ||
|
|
||
| /** | ||
| * Hosted create response when NAMS resolves-before-create merges the new | ||
| * name onto an existing entity: `{id, resolution: "merged", merged_into, | ||
| * confidence}` — no name/type fields. | ||
| */ | ||
| interface WireMergedResolution { | ||
| id?: string; | ||
| resolution: string; | ||
| merged_into?: string; | ||
| confidence?: number; | ||
| } | ||
|
|
||
| interface WirePreference { | ||
| id: string; | ||
| category: string; | ||
|
|
@@ -108,20 +120,23 @@ interface WireGraph { | |
| edges?: WireGraphEdge[]; | ||
| } | ||
|
|
||
| // NAMS projects unset node properties as JSON null (e.g. `confidence` on a | ||
| // manually-created entity) — normalize those to undefined so Entity's | ||
| // optional fields stay `T | undefined` at runtime, matching their types. | ||
| function toEntity(w: WireEntity): Entity { | ||
| return { | ||
| id: w.id, | ||
| name: w.name, | ||
| type: w.type, | ||
| subtype: w.subtype, | ||
| description: w.description, | ||
| embedding: w.embedding, | ||
| canonicalName: w.canonical_name, | ||
| subtype: w.subtype ?? undefined, | ||
| description: w.description ?? undefined, | ||
| embedding: w.embedding ?? undefined, | ||
| canonicalName: w.canonical_name ?? undefined, | ||
| createdAt: w.created_at ?? "", | ||
| updatedAt: w.updated_at, | ||
| confidence: w.confidence, | ||
| sourceStage: w.source_stage, | ||
| relationships: w.relationships?.map(toRelRef), | ||
| updatedAt: w.updated_at ?? undefined, | ||
| confidence: w.confidence ?? undefined, | ||
| sourceStage: w.source_stage ?? undefined, | ||
| relationships: w.relationships?.map(toRelRef) ?? undefined, | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Normalization stops at |
||
| }; | ||
| } | ||
|
|
||
|
|
@@ -192,13 +207,48 @@ export class LongTermMemory { | |
| entityType: string, | ||
| options?: { description?: string }, | ||
| ): Promise<Entity> { | ||
| const wire = await this.transport.request<WireEntity>("add_entity", { | ||
| const wire = await this.transport.request<WireEntity | WireMergedResolution>("add_entity", { | ||
| name, | ||
| entity_type: entityType, | ||
| type: entityType, | ||
| description: options?.description, | ||
| }); | ||
| return toEntity(wire); | ||
| // NAMS resolves-before-create: a sufficiently similar name merges onto an | ||
| // existing entity and the response carries no name/type — follow up with | ||
| // a GET for the canonical merged-into record. | ||
| if (wire && typeof wire === "object" && "resolution" in wire && wire.resolution === "merged") { | ||
| return this.resolveMergedEntity(wire, name, entityType); | ||
| } | ||
| return toEntity(wire as WireEntity); | ||
| } | ||
|
|
||
| /** | ||
| * Turn a `resolution: "merged"` create response into the canonical Entity. | ||
| * | ||
| * Fetches the merged-into entity; if that read fails (or no id was | ||
| * returned), synthesizes a minimal record from the request so the caller | ||
| * still gets a well-formed Entity. | ||
| */ | ||
| private async resolveMergedEntity( | ||
| wire: WireMergedResolution, | ||
| name: string, | ||
| entityType: string, | ||
| ): Promise<Entity> { | ||
| const mergedId = wire.merged_into ?? wire.id; | ||
| if (mergedId) { | ||
| try { | ||
| return await this.getEntity(mergedId); | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Missing Python's guard ( Sharper: |
||
| } catch (error) { | ||
| if (!(error instanceof MemoryError)) throw error; | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Same over-broad catch as Python: Caveat on the obvious fix: |
||
| // Canonical record unreadable — fall through to the synthesized form. | ||
| } | ||
| } | ||
| return toEntity({ | ||
| id: mergedId ?? "", | ||
| name, | ||
| type: entityType, | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Also |
||
| confidence: wire.confidence, | ||
| }); | ||
| } | ||
|
|
||
| async addPreference( | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Runs on every entity parse, not just merges, and strips identity fields.
"id": null-> client-side UUID viaMemoryEntry.id'sdefault_factory=uuid4.search_entitiescan then return entities whose id does not exist server-side; a laterget_entity_relationships(entity.id)404s with no clue why."created_at": null-> silentlynow().Restrict to fields with meaningful defaults (
confidence,description,subtype,updated_at).