Add create_entity and delete_entity MCP tools

README.md

@@ -52,6 +52,8 @@ npx @modelcontextprotocol/inspector@latest http://127.0.0.1:8201/sse --cli --met
 **Available tools:**
 - `get_guidelines(task: str)`: Get relevant guidelines for a specific task.
 - `save_trajectory(trajectory_data: str, task_id: str | None)`: Save a conversation trajectory and generate new tips.
+- `create_entity(content: str, entity_type: str, metadata: str | None, enable_conflict_resolution: bool)`: Create a single entity in the namespace.
+- `delete_entity(entity_id: str)`: Delete a specific entity by its ID.
 
 ## Documentation
 

kaizen/backend/milvus.py

@@ -33,8 +33,12 @@ def deserialize_content(content: str):
 
 
 class MilvusEntityBackend(BaseEntityBackend):
-    milvus = MilvusClient(**milvus_client_settings.model_dump())
-    embedding_model = SentenceTransformer(milvus_other_settings.embedding_model)
+    # Removed class attributes
+
+    def __init__(self, config=None):
+        super().__init__(config)
+        self.milvus = MilvusClient(**milvus_client_settings.model_dump())
+        self.embedding_model = SentenceTransformer(milvus_other_settings.embedding_model)
 
     def ready(self):
         _ = self.milvus.list_collections()
@@ -143,19 +147,21 @@ def update_entities(
             updates = []
             for entity in entities:
                 content_str = serialize_content(entity.content)
+                # Convert None metadata to empty dict for Milvus compatibility
+                metadata = entity.metadata if entity.metadata is not None else {}
                 entity_id = str(self.milvus.insert(collection_name=namespace_id, data={
                     'type': entity_type,
                     'content': content_str,
                     'created_at': int(now.timestamp()),
                     'embedding': self.embedding_model.encode(content_str),
-                    'metadata': entity.metadata
+                    'metadata': metadata
                 })['ids'][0])
                 updates.append(EntityUpdate(
                     id=entity_id,
                     type=entity_type,
                     content=entity.content,
                     event='ADD',
-                    metadata=entity.metadata
+                    metadata=metadata
                 ))
         return updates
 
@@ -187,9 +193,30 @@ def search_entities(
         return [parse_milvus_entity(i) for i in results]
 
     def delete_entity_by_id(self, namespace_id: str, entity_id: str):
-        entity_id = int(entity_id)
+        try:
+            entity_id_int = int(entity_id)
+        except ValueError:
+            raise KaizenException(f"Invalid entity ID: {entity_id}. Entity IDs must be numeric.")
         self.validate_namespace(namespace_id)
-        self.milvus.delete(collection_name=namespace_id, ids=[entity_id])
+
+        # Check if entity exists before deleting
+        existing = self.milvus.query(
+            collection_name=namespace_id,
+            filter=f"id == {entity_id_int}",
+            output_fields=["id"]
+        )
+        if not existing:
+            raise KaizenException(f"Entity with ID {entity_id} not found in namespace {namespace_id}.")
+
+        self.milvus.delete(collection_name=namespace_id, ids=[entity_id_int])
+
+    def close(self):
+        """Close Milvus connection."""
+        try:
+            if hasattr(self, 'milvus'):
+                self.milvus.close()
+        except Exception as e:
+            logger.warning(f"Error closing Milvus client: {e}")
 
 entity_schema = CollectionSchema(fields=[
     # Keep it as an INT64 or else you won't be able to list all entities.

kaizen/frontend/mcp/mcp_server.py

@@ -13,20 +13,32 @@
 from kaizen.frontend.client.kaizen_client import KaizenClient
 from kaizen.llm.tips.tips import generate_tips
 from kaizen.schema.core import Entity, RecordedEntity
-from kaizen.schema.exceptions import NamespaceNotFoundException
+from kaizen.schema.exceptions import KaizenException, NamespaceNotFoundException
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("entities-mcp")
 
 mcp = FastMCP("entities")
-client = KaizenClient()
+_client = None
+
+
+def get_client() -> KaizenClient:
+    """Get or create the KaizenClient singleton.
+
+    This lazy initialization allows tests to configure settings
+    before the client is created.
+    """
+    global _client
+    if _client is None:
+        _client = KaizenClient()
+    return _client
 
 
 def ensure_namespace():
     try:
-        client.get_namespace_details(kaizen_config.namespace_id)
+        get_client().get_namespace_details(kaizen_config.namespace_id)
     except NamespaceNotFoundException:
-        client.create_namespace(kaizen_config.namespace_id)
+        get_client().create_namespace(kaizen_config.namespace_id)
 
 
 @mcp.tool()
@@ -41,7 +53,7 @@ def get_guidelines(task: str) -> str:
     logger.info(f"Getting guidelines for task: {task}")
     ensure_namespace()
     # Get relevant guidelines
-    results = client.search_entities(
+    results = get_client().search_entities(
         namespace_id=kaizen_config.namespace_id,
         query=task,
         filters={"type": "guideline"},
@@ -85,14 +97,14 @@ def save_trajectory(
             )
         )
 
-    client.update_entities(
+    get_client().update_entities(
         namespace_id=kaizen_config.namespace_id,
         entities=entities,
         enable_conflict_resolution=False,
     )
     tips = generate_tips(messages)
 
-    client.update_entities(
+    get_client().update_entities(
         namespace_id=kaizen_config.namespace_id,
         entities=[
             Entity(
@@ -109,8 +121,103 @@ def save_trajectory(
         enable_conflict_resolution=True,
     )
 
-    return client.search_entities(
+    return get_client().search_entities(
         namespace_id=kaizen_config.namespace_id,
         filters={"type": "trajectory", "task_id": task_id},
         limit=1000,
     )
+
+
+@mcp.tool()
+def create_entity(
+    content: str,
+    entity_type: str,
+    metadata: str | None = None,
+    enable_conflict_resolution: bool = False
+) -> str:
+    """
+    Create a single entity in the namespace.
+
+    Args:
+        content: The searchable text or structured data for the entity
+        entity_type: The type/category of the entity (e.g., 'guideline', 'note', 'fact')
+        metadata: Optional JSON string containing arbitrary metadata related to the entity
+        enable_conflict_resolution: If True, uses LLM to check for conflicts with existing entities
+
+    Returns:
+        JSON string with the entity update details (ADD/UPDATE/DELETE/NONE) and entity ID
+    """
+    logger.info(f"Creating entity of type: {entity_type}")
+    ensure_namespace()
+
+    # Parse metadata if provided
+    metadata_dict = None
+    if metadata:
+        try:
+            metadata_dict = json.loads(metadata)
+        except json.JSONDecodeError as e:
+            logger.exception(f"Invalid JSON in metadata parameter: {str(e)}")
+            return json.dumps({
+                "error": "Invalid metadata JSON",
+                "message": f"Failed to parse metadata: {str(e)}",
+                "invalid_metadata": metadata
+            })
+
+    # Create the entity using the Entity schema
+    entity = Entity(
+        type=entity_type,
+        content=content,
+        metadata=metadata_dict
+    )
+
+    # Use KaizenClient.update_entities() to create the entity
+    updates = get_client().update_entities(
+        namespace_id=kaizen_config.namespace_id,
+        entities=[entity],
+        enable_conflict_resolution=enable_conflict_resolution
+    )
+
+    # Return the first (and only) update result
+    if updates:
+        update = updates[0]
+        return json.dumps({
+            "event": update.event,
+            "id": update.id,
+            "type": update.type,
+            "content": update.content,
+            "metadata": update.metadata
+        })
+    else:
+        return json.dumps({"error": "Entity creation failed"})
+
+
+@mcp.tool()
+def delete_entity(entity_id: str) -> str:
+    """
+    Delete a specific entity by its ID.
+
+    Args:
+        entity_id: The unique identifier of the entity to delete
+
+    Returns:
+        JSON string confirming deletion or error message
+    """
+    logger.info(f"Deleting entity: {entity_id}")
+    ensure_namespace()
+
+    try:
+        # Use KaizenClient.delete_entity_by_id() to delete the entity
+        get_client().delete_entity_by_id(
+            namespace_id=kaizen_config.namespace_id,
+            entity_id=entity_id
+        )
+        return json.dumps({
+            "success": True,
+            "message": f"Entity {entity_id} deleted successfully"
+        })
+    except KaizenException as e:
+        logger.exception(f"Error deleting entity {entity_id}: {str(e)}")
+        return json.dumps({
+            "success": False,
+            "error": str(e)
+        })