From 8ffbc67c7c33e655a1bdc95df62b8fd242a7e201 Mon Sep 17 00:00:00 2001 From: Vatche Isahagian Date: Fri, 23 Jan 2026 20:02:22 -0500 Subject: [PATCH 1/6] Refactor: Rename 'kata' to 'entity' across codebase --- kaizen/backend/base.py | 4 ++-- kaizen/backend/filesystem.py | 6 +++--- kaizen/backend/milvus.py | 6 +++--- kaizen/config/milvus.py | 2 +- kaizen/db/sqlite_manager.py | 2 +- kaizen/frontend/client/kaizen_client.py | 12 ++++++------ kaizen/frontend/mcp/mcp_server.py | 8 ++++---- kaizen/schema/conflict_resolution.py | 2 +- tests/unit/test_client.py | 26 ++++++++++++------------- 9 files changed, 34 insertions(+), 34 deletions(-) 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/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') From 5d77adf47a8ef521873fbdc401a71ddf1b8c9635 Mon Sep 17 00:00:00 2001 From: Vatche Isahagian Date: Mon, 26 Jan 2026 10:24:01 -0500 Subject: [PATCH 2/6] changes to address reviews --- kaizen/sync/phoenix_sync.py | 22 +++++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/kaizen/sync/phoenix_sync.py b/kaizen/sync/phoenix_sync.py index 3192aa72..7148a4ab 100644 --- a/kaizen/sync/phoenix_sync.py +++ b/kaizen/sync/phoenix_sync.py @@ -145,6 +145,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: + pass + + 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,8 +181,8 @@ 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 + output_msgs = [{"role": "assistant", "content": output_val}] + #pass except: pass @@ -184,6 +192,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: + pass + + 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") @@ -503,7 +519,7 @@ def sync( ) except Exception as e: error_msg = f"Error processing span {span_id}: {e}" - logger.error(error_msg) + logger.error(error_msg) errors.append(error_msg) result = SyncResult( From 851ad2d7439d6e3bf2b2f16882ea17bacd532f66 Mon Sep 17 00:00:00 2001 From: Vatche Isahagian Date: Mon, 26 Jan 2026 10:49:05 -0500 Subject: [PATCH 3/6] fixes for coderabbit --- kaizen/sync/phoenix_sync.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/kaizen/sync/phoenix_sync.py b/kaizen/sync/phoenix_sync.py index 7148a4ab..3e5da76c 100644 --- a/kaizen/sync/phoenix_sync.py +++ b/kaizen/sync/phoenix_sync.py @@ -364,7 +364,7 @@ def _extract_trajectory(self, span: dict) -> dict: "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"), + "total_tokens": attrs.get("gen_ai.usage.total_tokens") or attrs.get("llm.token_count.total"), }, } From 3be2a288010b8b039c90290ef91fcd6aa0c4be43 Mon Sep 17 00:00:00 2001 From: Vatche Isahagian Date: Mon, 26 Jan 2026 11:33:04 -0500 Subject: [PATCH 4/6] fix token usage extraction and improve error logging --- kaizen/sync/phoenix_sync.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/kaizen/sync/phoenix_sync.py b/kaizen/sync/phoenix_sync.py index 3e5da76c..261ce8d0 100644 --- a/kaizen/sync/phoenix_sync.py +++ b/kaizen/sync/phoenix_sync.py @@ -132,8 +132,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: {input_val}") if input_msgs: # Handle OpenInference format @@ -148,8 +148,8 @@ def _extract_messages_from_span(self, span: dict) -> list[dict]: if isinstance(msg, str): try: msg = self._parse_content(msg) - except: - pass + except Exception as e: + logger.debug(f"Failed to parse input message string: {e}. Payload: {msg}") if not isinstance(msg, dict): continue @@ -183,8 +183,8 @@ def _extract_messages_from_span(self, span: dict) -> list[dict]: # Fallback for simple string output output_msgs = [{"role": "assistant", "content": output_val}] #pass - except: - pass + except Exception as e: + logger.debug(f"Failed to parse output.value: {e}. Payload: {output_val}") if output_msgs: if isinstance(output_msgs, str): @@ -195,8 +195,8 @@ def _extract_messages_from_span(self, span: dict) -> list[dict]: if isinstance(msg, str): try: msg = self._parse_content(msg) - except: - pass + except Exception as e: + logger.debug(f"Failed to parse output message string: {e}. Payload: {msg}") if not isinstance(msg, dict): continue @@ -362,9 +362,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("gen_ai.usage.total_tokens") or attrs.get("llm.token_count.total"), + "prompt_tokens": attrs.get("gen_ai.usage.prompt_tokens") or attrs.get("llm.token_count.prompt") or attrs.get("llm.usage.prompt_tokens"), + "completion_tokens": attrs.get("gen_ai.usage.completion_tokens") or attrs.get("llm.token_count.completion") or attrs.get("llm.usage.completion_tokens"), + "total_tokens": attrs.get("gen_ai.usage.total_tokens") or attrs.get("llm.token_count.total") or attrs.get("llm.usage.total_tokens"), }, } @@ -519,7 +519,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( From 55916c69b0d77f889b8f355db6bc9b7f7645eb82 Mon Sep 17 00:00:00 2001 From: Vatche Isahagian Date: Mon, 26 Jan 2026 11:44:27 -0500 Subject: [PATCH 5/6] fix token usage extraction and secure error logging --- kaizen/sync/phoenix_sync.py | 25 +++++++++++++++++++++---- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/kaizen/sync/phoenix_sync.py b/kaizen/sync/phoenix_sync.py index 261ce8d0..ed5b00e1 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): @@ -133,7 +150,7 @@ def _extract_messages_from_span(self, span: dict) -> list[dict]: elif isinstance(parsed_input, list): # rare but possible input_msgs = parsed_input except Exception as e: - logger.debug(f"Failed to parse input.value: {e}. Payload: {input_val}") + logger.debug(f"Failed to parse input.value: {e}. Payload: {self._format_payload_summary(input_val)}") if input_msgs: # Handle OpenInference format @@ -149,7 +166,7 @@ def _extract_messages_from_span(self, span: dict) -> list[dict]: try: msg = self._parse_content(msg) except Exception as e: - logger.debug(f"Failed to parse input message string: {e}. Payload: {msg}") + logger.debug(f"Failed to parse input message string: {e}. Payload: {self._format_payload_summary(msg)}") if not isinstance(msg, dict): continue @@ -184,7 +201,7 @@ def _extract_messages_from_span(self, span: dict) -> list[dict]: output_msgs = [{"role": "assistant", "content": output_val}] #pass except Exception as e: - logger.debug(f"Failed to parse output.value: {e}. Payload: {output_val}") + logger.debug(f"Failed to parse output.value: {e}. Payload: {self._format_payload_summary(output_val)}") if output_msgs: if isinstance(output_msgs, str): @@ -196,7 +213,7 @@ def _extract_messages_from_span(self, span: dict) -> list[dict]: try: msg = self._parse_content(msg) except Exception as e: - logger.debug(f"Failed to parse output message string: {e}. Payload: {msg}") + logger.debug(f"Failed to parse output message string: {e}. Payload: {self._format_payload_summary(msg)}") if not isinstance(msg, dict): continue From 5e7b19474f52c9ed0ea8cd0b4cfd649c7e412ac6 Mon Sep 17 00:00:00 2001 From: Vatche Isahagian Date: Mon, 26 Jan 2026 12:19:14 -0500 Subject: [PATCH 6/6] secure logging and robust token usage extraction --- kaizen/sync/phoenix_sync.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/kaizen/sync/phoenix_sync.py b/kaizen/sync/phoenix_sync.py index ed5b00e1..d83c59d5 100644 --- a/kaizen/sync/phoenix_sync.py +++ b/kaizen/sync/phoenix_sync.py @@ -379,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") or attrs.get("llm.usage.prompt_tokens"), - "completion_tokens": attrs.get("gen_ai.usage.completion_tokens") or attrs.get("llm.token_count.completion") or attrs.get("llm.usage.completion_tokens"), - "total_tokens": attrs.get("gen_ai.usage.total_tokens") or attrs.get("llm.token_count.total") or attrs.get("llm.usage.total_tokens"), + "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), }, }