diff --git a/kaizen/backend/base.py b/kaizen/backend/base.py index 73ed6665..29c0e985 100644 --- a/kaizen/backend/base.py +++ b/kaizen/backend/base.py @@ -5,10 +5,10 @@ from kaizen.schema.conflict_resolution import EntityUpdate logging.basicConfig(level=logging.INFO) -logger = logging.getLogger("katas-db") +logger = logging.getLogger("entities-db") -class BaseKataBackend(ABC): +class BaseEntityBackend(ABC): def __init__(self, config: BaseSettings | None = None): pass diff --git a/kaizen/backend/filesystem.py b/kaizen/backend/filesystem.py index 01a6712c..4b380f99 100644 --- a/kaizen/backend/filesystem.py +++ b/kaizen/backend/filesystem.py @@ -5,7 +5,7 @@ from pathlib import Path from threading import Lock -from kaizen.backend.base import BaseKataBackend +from kaizen.backend.base import BaseEntityBackend from kaizen.config.filesystem import filesystem_settings from kaizen.llm.conflict_resolution.conflict_resolution import resolve_conflicts from kaizen.schema.conflict_resolution import EntityUpdate @@ -17,10 +17,10 @@ ) logging.basicConfig(level=logging.INFO) -logger = logging.getLogger("katas-db.filesystem") +logger = logging.getLogger("entities-db.filesystem") -class FilesystemKataBackend(BaseKataBackend): +class FilesystemEntityBackend(BaseEntityBackend): """A filesystem-based backend that stores data in JSON files. This backend uses simple text matching for search (no embeddings). diff --git a/kaizen/backend/milvus.py b/kaizen/backend/milvus.py index ac98c78d..4bfa83d9 100644 --- a/kaizen/backend/milvus.py +++ b/kaizen/backend/milvus.py @@ -3,7 +3,7 @@ import logging import uuid -from kaizen.backend.base import BaseKataBackend +from kaizen.backend.base import BaseEntityBackend from kaizen.config.milvus import milvus_client_settings, milvus_other_settings from kaizen.db.sqlite_manager import SQLiteManager from kaizen.llm.conflict_resolution.conflict_resolution import resolve_conflicts @@ -14,7 +14,7 @@ from sentence_transformers import SentenceTransformer logging.basicConfig(level=logging.INFO) -logger = logging.getLogger("katas-db.milvus") +logger = logging.getLogger("entities-db.milvus") def serialize_content(content) -> str: @@ -32,7 +32,7 @@ def deserialize_content(content: str): return content -class MilvusKataBackend(BaseKataBackend): +class MilvusEntityBackend(BaseEntityBackend): milvus = MilvusClient(**milvus_client_settings.model_dump()) embedding_model = SentenceTransformer(milvus_other_settings.embedding_model) diff --git a/kaizen/config/milvus.py b/kaizen/config/milvus.py index c0a712fd..1e03b087 100644 --- a/kaizen/config/milvus.py +++ b/kaizen/config/milvus.py @@ -4,7 +4,7 @@ class MilvusDBSettings(BaseSettings): model_config = SettingsConfigDict(env_prefix='KAIZEN_') - uri: str = Field(default='katas.milvus.db') + uri: str = Field(default='entities.milvus.db') user: str = Field(default='') password: str = Field(default='') db_name: str = Field(default='') diff --git a/kaizen/db/sqlite_manager.py b/kaizen/db/sqlite_manager.py index 3568ef47..10344163 100644 --- a/kaizen/db/sqlite_manager.py +++ b/kaizen/db/sqlite_manager.py @@ -23,7 +23,7 @@ def convert_timestamp(time: bytes) -> datetime.datetime: class SQLiteManager: """A database for any resources that can't be generalized across backends.""" - def __init__(self, db_path: str = 'katas.sqlite.db'): + def __init__(self, db_path: str = 'entities.sqlite.db'): self.db_path = db_path def _create_namespace_table(self): diff --git a/kaizen/frontend/client/kaizen_client.py b/kaizen/frontend/client/kaizen_client.py index 19e0a195..9861f30e 100644 --- a/kaizen/frontend/client/kaizen_client.py +++ b/kaizen/frontend/client/kaizen_client.py @@ -4,19 +4,19 @@ from kaizen.config.kaizen import KaizenConfig class KaizenClient: - """Wrapper client around kaizen kata backends.""" + """Wrapper client around kaizen entity backends.""" def __init__(self, config: KaizenConfig | None = None): """Initialize the Kaizen client.""" self.config = config or KaizenConfig() if self.config.backend == 'milvus': - from kaizen.backend.milvus import MilvusKataBackend - self.backend = MilvusKataBackend(self.config.settings) + from kaizen.backend.milvus import MilvusEntityBackend + self.backend = MilvusEntityBackend(self.config.settings) elif self.config.backend == 'filesystem': - from kaizen.backend.filesystem import FilesystemKataBackend - self.backend = FilesystemKataBackend(self.config.settings) + from kaizen.backend.filesystem import FilesystemEntityBackend + self.backend = FilesystemEntityBackend(self.config.settings) else: - raise NotImplementedError(f'Kata backend not implemented: {self.config.backend}') + raise NotImplementedError(f'Entity backend not implemented: {self.config.backend}') def ready(self) -> bool: """Check if the backend is healthy.""" diff --git a/kaizen/frontend/mcp/mcp_server.py b/kaizen/frontend/mcp/mcp_server.py index 6e2ff8e9..3cbd07e8 100644 --- a/kaizen/frontend/mcp/mcp_server.py +++ b/kaizen/frontend/mcp/mcp_server.py @@ -1,5 +1,5 @@ """ -Katas MCP Server +Kaizen MCP Server This server provides a tool to get task-relevant guidelines. """ @@ -16,9 +16,9 @@ from kaizen.schema.exceptions import NamespaceNotFoundException logging.basicConfig(level=logging.INFO) -logger = logging.getLogger("katas-mcp") +logger = logging.getLogger("entities-mcp") -mcp = FastMCP("katas") +mcp = FastMCP("entities") client = KaizenClient() @@ -61,7 +61,7 @@ def save_trajectory( trajectory_data: str, task_id: str | None = None ) -> list[RecordedEntity]: """ - Save the full agent trajectory to the Kata DB and generate tips + Save the full agent trajectory to the Entity DB and generate tips Args: trajectory_data: A JSON formatted OpenAI conversation. diff --git a/kaizen/schema/conflict_resolution.py b/kaizen/schema/conflict_resolution.py index 31c38d5a..5dc36dee 100644 --- a/kaizen/schema/conflict_resolution.py +++ b/kaizen/schema/conflict_resolution.py @@ -14,7 +14,7 @@ def from_recorded_entities(entities: list[RecordedEntity]) -> list['SimpleEntity class EntityUpdate(BaseModel): - """Produced by the LLM, to be processed by a kata backend.""" + """Produced by the LLM, to be processed by a entity backend.""" id: str = Field(description='The unique ID of an entity.') type: str = Field(description='The type of the entity.') content: str | list | dict = Field(description='The content of the entity.') diff --git a/kaizen/sync/phoenix_sync.py b/kaizen/sync/phoenix_sync.py index 3192aa72..d83c59d5 100644 --- a/kaizen/sync/phoenix_sync.py +++ b/kaizen/sync/phoenix_sync.py @@ -98,6 +98,23 @@ def _get_processed_span_ids(self) -> set[str]: except NamespaceNotFoundException: return set() + def _format_payload_summary(self, payload: Any) -> str: + """Format a payload summary for secure logging (avoid PII).""" + type_name = type(payload).__name__ + + if isinstance(payload, str): + length = len(payload) + preview = payload[:50] + "..." if length > 50 else payload + # Replace newlines in preview to keep logs on one line + preview = preview.replace("\n", "\\n") + return f"<{type_name} length={length} preview='{preview}'>" + + if isinstance(payload, (dict, list)): + length = len(payload) + return f"<{type_name} length={length}>" + + return f"<{type_name}>" + def _parse_content(self, content: Any) -> Any: """Parse content which may be a string representation of a list/dict.""" if isinstance(content, str): @@ -132,8 +149,8 @@ def _extract_messages_from_span(self, span: dict) -> list[dict]: input_msgs = parsed_input["messages"] elif isinstance(parsed_input, list): # rare but possible input_msgs = parsed_input - except: - pass + except Exception as e: + logger.debug(f"Failed to parse input.value: {e}. Payload: {self._format_payload_summary(input_val)}") if input_msgs: # Handle OpenInference format @@ -145,6 +162,14 @@ def _extract_messages_from_span(self, span: dict) -> list[dict]: for i, msg in enumerate(input_msgs): # OpenInference often uses message.role / message.content keys in flattened export # but via API it might be cleaner. Let's handle dict access safely. + if isinstance(msg, str): + try: + msg = self._parse_content(msg) + except Exception as e: + logger.debug(f"Failed to parse input message string: {e}. Payload: {self._format_payload_summary(msg)}") + + if not isinstance(msg, dict): + continue role = msg.get("message.role") or msg.get("role") content = msg.get("message.content") or msg.get("content") tool_calls = msg.get("message.tool_calls") or msg.get("tool_calls") @@ -173,10 +198,10 @@ def _extract_messages_from_span(self, span: dict) -> list[dict]: output_msgs = [c["message"] for c in parsed_output["choices"]] else: # Fallback for simple string output - # output_msgs = [{"role": "assistant", "content": output_val}] - pass - except: - pass + output_msgs = [{"role": "assistant", "content": output_val}] + #pass + except Exception as e: + logger.debug(f"Failed to parse output.value: {e}. Payload: {self._format_payload_summary(output_val)}") if output_msgs: if isinstance(output_msgs, str): @@ -184,6 +209,14 @@ def _extract_messages_from_span(self, span: dict) -> list[dict]: if isinstance(output_msgs, list): for i, msg in enumerate(output_msgs): + if isinstance(msg, str): + try: + msg = self._parse_content(msg) + except Exception as e: + logger.debug(f"Failed to parse output message string: {e}. Payload: {self._format_payload_summary(msg)}") + + if not isinstance(msg, dict): + continue role = msg.get("message.role") or msg.get("role") content = msg.get("message.content") or msg.get("content") tool_calls = msg.get("message.tool_calls") or msg.get("tool_calls") @@ -346,9 +379,9 @@ def _extract_trajectory(self, span: dict) -> dict: "timestamp": span.get("start_time"), "messages": openai_messages, "usage": { - "prompt_tokens": attrs.get("gen_ai.usage.prompt_tokens") or attrs.get("llm.token_count.prompt"), - "completion_tokens": attrs.get("gen_ai.usage.completion_tokens") or attrs.get("llm.token_count.completion"), - "total_tokens": attrs.get("llm.usage.total_tokens") or attrs.get("llm.token_count.total"), + "prompt_tokens": next((v for v in [attrs.get("gen_ai.usage.prompt_tokens"), attrs.get("llm.token_count.prompt"), attrs.get("llm.usage.prompt_tokens")] if v is not None), None), + "completion_tokens": next((v for v in [attrs.get("gen_ai.usage.completion_tokens"), attrs.get("llm.token_count.completion"), attrs.get("llm.usage.completion_tokens")] if v is not None), None), + "total_tokens": next((v for v in [attrs.get("gen_ai.usage.total_tokens"), attrs.get("llm.token_count.total"), attrs.get("llm.usage.total_tokens")] if v is not None), None), }, } @@ -503,7 +536,7 @@ def sync( ) except Exception as e: error_msg = f"Error processing span {span_id}: {e}" - logger.error(error_msg) + logger.exception(error_msg) errors.append(error_msg) result = SyncResult( diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 4f1a1dac..5d8901be 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -4,7 +4,7 @@ import datetime import pytest -from kaizen.backend.base import BaseKataBackend +from kaizen.backend.base import BaseEntityBackend from kaizen.schema.core import Entity, Namespace, RecordedEntity from kaizen.schema.conflict_resolution import EntityUpdate from kaizen.schema.exceptions import NamespaceNotFoundException, NamespaceAlreadyExistsException @@ -21,14 +21,14 @@ def test_health_check(kaizen_client: KaizenClient, monkeypatch): # Client should return False derived from API response def ready(self) -> bool: return False - monkeypatch.setattr(kaizen_client.backend, 'ready', ready.__get__(kaizen_client, BaseKataBackend)) + monkeypatch.setattr(kaizen_client.backend, 'ready', ready.__get__(kaizen_client, BaseEntityBackend)) health_status = kaizen_client.ready() assert health_status == False # Client should return True derived from API response def ready(self) -> bool: return True - monkeypatch.setattr(kaizen_client.backend, 'ready', ready.__get__(kaizen_client, BaseKataBackend)) + monkeypatch.setattr(kaizen_client.backend, 'ready', ready.__get__(kaizen_client, BaseEntityBackend)) health_status = kaizen_client.ready() assert health_status == True @@ -39,7 +39,7 @@ def test_create_namespace(kaizen_client: KaizenClient, monkeypatch): def create_namespace(self, namespace_id=None) -> Namespace: return Namespace(id='foobar', created_at=created_at) - monkeypatch.setattr(kaizen_client.backend, 'create_namespace', create_namespace.__get__(kaizen_client.backend, BaseKataBackend)) + monkeypatch.setattr(kaizen_client.backend, 'create_namespace', create_namespace.__get__(kaizen_client.backend, BaseEntityBackend)) result = kaizen_client.create_namespace(namespace_id='foobar') assert result.id == 'foobar' assert result.created_at == created_at @@ -50,7 +50,7 @@ def test_create_namespace_already_exists(kaizen_client: KaizenClient, monkeypatc def create_namespace(self, namespace_id=None) -> Namespace: raise NamespaceAlreadyExistsException() - monkeypatch.setattr(kaizen_client.backend, 'create_namespace', create_namespace.__get__(kaizen_client.backend, BaseKataBackend)) + monkeypatch.setattr(kaizen_client.backend, 'create_namespace', create_namespace.__get__(kaizen_client.backend, BaseEntityBackend)) with pytest.raises(NamespaceAlreadyExistsException) as e: kaizen_client.create_namespace(namespace_id='foobar') @@ -61,7 +61,7 @@ def test_get_namespace_details(kaizen_client: KaizenClient, monkeypatch): def get_namespace_details(self, namespace_id=None) -> Namespace: return Namespace(id='foobar', created_at=created_at) - monkeypatch.setattr(kaizen_client.backend, 'get_namespace_details', get_namespace_details.__get__(kaizen_client.backend, BaseKataBackend)) + monkeypatch.setattr(kaizen_client.backend, 'get_namespace_details', get_namespace_details.__get__(kaizen_client.backend, BaseEntityBackend)) result = kaizen_client.get_namespace_details(namespace_id='foobar') assert result.id == 'foobar' assert result.created_at == created_at @@ -72,7 +72,7 @@ def test_get_namespace_details_nonexistent(kaizen_client: KaizenClient, monkeypa def get_namespace_details(self, namespace_id=None) -> Namespace: raise NamespaceNotFoundException() - monkeypatch.setattr(kaizen_client.backend, 'get_namespace_details', get_namespace_details.__get__(kaizen_client.backend, BaseKataBackend)) + monkeypatch.setattr(kaizen_client.backend, 'get_namespace_details', get_namespace_details.__get__(kaizen_client.backend, BaseEntityBackend)) with pytest.raises(NamespaceNotFoundException) as e: kaizen_client.get_namespace_details(namespace_id='foobar') @@ -84,7 +84,7 @@ def test_search_namespaces(kaizen_client: KaizenClient, monkeypatch): def search_namespaces(self, limit=10) -> list[Namespace]: return [Namespace(id='foobar', created_at=created_at)] - monkeypatch.setattr(kaizen_client.backend, 'search_namespaces', search_namespaces.__get__(kaizen_client.backend, BaseKataBackend)) + monkeypatch.setattr(kaizen_client.backend, 'search_namespaces', search_namespaces.__get__(kaizen_client.backend, BaseEntityBackend)) result = kaizen_client.search_namespaces() assert result[0].id == 'foobar' assert result[0].created_at == created_at @@ -95,7 +95,7 @@ def test_delete_namespace(kaizen_client: KaizenClient, monkeypatch): def delete_namespace(self, namespace_id): pass - monkeypatch.setattr(kaizen_client.backend, 'delete_namespace', delete_namespace.__get__(kaizen_client.backend, BaseKataBackend)) + monkeypatch.setattr(kaizen_client.backend, 'delete_namespace', delete_namespace.__get__(kaizen_client.backend, BaseEntityBackend)) kaizen_client.delete_namespace(namespace_id='foobar') @pytest.mark.unit @@ -104,7 +104,7 @@ def test_update_entities(kaizen_client: KaizenClient, monkeypatch): def update_entities(self, namespace_id, entity, enable_conflict_resolution=True) -> list[EntityUpdate]: return [EntityUpdate(id='1', type='fact', content="User's name is Foobar", event='ADD')] - monkeypatch.setattr(kaizen_client.backend, 'update_entities', update_entities.__get__(kaizen_client.backend, BaseKataBackend)) + monkeypatch.setattr(kaizen_client.backend, 'update_entities', update_entities.__get__(kaizen_client.backend, BaseEntityBackend)) result = kaizen_client.update_entities(namespace_id='foobar', entities=[Entity(type='fact', content="User's name is Foobar.")]) assert result[0].id == '1' assert result[0].content == "User's name is Foobar" @@ -117,7 +117,7 @@ def test_search_entities(kaizen_client: KaizenClient, monkeypatch): def search_entities(self, namespace_id, query, filters, limit=10) -> list[RecordedEntity]: return [RecordedEntity(id='1', type='fact', created_at=created_at, content="User's name is Foobar.")] - monkeypatch.setattr(kaizen_client.backend, 'search_entities', search_entities.__get__(kaizen_client.backend, BaseKataBackend)) + monkeypatch.setattr(kaizen_client.backend, 'search_entities', search_entities.__get__(kaizen_client.backend, BaseEntityBackend)) result = kaizen_client.search_entities(namespace_id='foobar', query='name') assert result[0].id == '1' assert result[0].content == "User's name is Foobar." @@ -130,7 +130,7 @@ def test_get_all_entities(kaizen_client: KaizenClient, monkeypatch): def search_entities(self, namespace_id, query, filters, limit=10) -> list[RecordedEntity]: return [RecordedEntity(id='1', type='fact', created_at=created_at, content="User's name is Foobar.")] - monkeypatch.setattr(kaizen_client.backend, 'search_entities', search_entities.__get__(kaizen_client.backend, BaseKataBackend)) + monkeypatch.setattr(kaizen_client.backend, 'search_entities', search_entities.__get__(kaizen_client.backend, BaseEntityBackend)) result = kaizen_client.search_entities(namespace_id='foobar', query='name') assert result[0].id == '1' assert result[0].content == "User's name is Foobar." @@ -142,5 +142,5 @@ def test_delete_entity(kaizen_client: KaizenClient, monkeypatch): def delete_entity_by_id(self, namespace_id, entity_id): pass - monkeypatch.setattr(kaizen_client.backend, 'delete_entity_by_id', delete_entity_by_id.__get__(kaizen_client.backend, BaseKataBackend)) + monkeypatch.setattr(kaizen_client.backend, 'delete_entity_by_id', delete_entity_by_id.__get__(kaizen_client.backend, BaseEntityBackend)) kaizen_client.delete_entity_by_id(namespace_id='foobar', entity_id='1')