From 88323d929913e31eb4e4b50975719112c2874cc7 Mon Sep 17 00:00:00 2001 From: Punleuk Oum Date: Fri, 16 Jan 2026 09:22:26 -0500 Subject: [PATCH 1/3] better tips? --- README.md | 16 +- kaizen/config/llm.py | 13 +- kaizen/frontend/mcp/mcp_server.py | 53 +++-- kaizen/llm/tips/prompts/generate_tips.jinja2 | 41 +++- kaizen/llm/tips/tips.py | 193 ++++++++++++++----- kaizen/schema/tips.py | 13 ++ kaizen/sync/phoenix_sync.py | 126 +++++++----- pyproject.toml | 2 +- 8 files changed, 323 insertions(+), 134 deletions(-) create mode 100644 kaizen/schema/tips.py diff --git a/README.md b/README.md index 6310eb7b..1f93cf05 100644 --- a/README.md +++ b/README.md @@ -53,14 +53,14 @@ All configuration variables are prefixed with `KAIZEN_`. **General Settings:** -| Variable | Description | Default | -|----------|-------------|---------| -| `KAIZEN_PROVIDER` | Backend provider (`milvus` or `filesystem`) | `milvus` | -| `KAIZEN_NAMESPACE_ID` | Namespace ID for isolation | `kaizen` | -| `KAIZEN_TIPS_MODEL` | Model for generating tips | `openai/gpt-4o` | -| `KAIZEN_CONFLICT_RESOLUTION_MODEL` | Model for resolving conflicts | `openai/gpt-4o` | -| `KAIZEN_CUSTOM_LLM_PROVIDER` | LiteLLM provider (use `openai` for proxy with custom models) | `openai` | -| `KAIZEN_EMBEDDING_MODEL` | Embedding model | `sentence-transformers/all-MiniLM-L6-v2` | +| Variable | Description | Default | +|----------|-------------------------------------------------------------------------------|------------------------------------------| +| `KAIZEN_PROVIDER` | Backend provider (`milvus` or `filesystem`) | `milvus` | +| `KAIZEN_NAMESPACE_ID` | Namespace ID for isolation | `kaizen` | +| `KAIZEN_TIPS_MODEL` | Model for generating tips (e.g. `openai/gpt-4o` for proxy with custom models) | `gpt-4o` | +| `KAIZEN_CONFLICT_RESOLUTION_MODEL` | Model for resolving conflicts (e.g. `openai/gpt-4o` for proxy with custom models) | `gpt-4o` | +| `KAIZEN_CUSTOM_LLM_PROVIDER` | LiteLLM provider (use `openai` for proxy with custom models) | `None` | +| `KAIZEN_EMBEDDING_MODEL` | Embedding model | `sentence-transformers/all-MiniLM-L6-v2` | **Milvus Backend Settings** (when `KAIZEN_PROVIDER=milvus`): diff --git a/kaizen/config/llm.py b/kaizen/config/llm.py index 7b106a87..99e687cf 100644 --- a/kaizen/config/llm.py +++ b/kaizen/config/llm.py @@ -1,10 +1,13 @@ +from pydantic import Field from pydantic_settings import BaseSettings, SettingsConfigDict + class LLMSettings(BaseSettings): - model_config = SettingsConfigDict(env_prefix='KAIZEN_') - tips_model: str = "openai/gpt-4o" - conflict_resolution_model: str = "openai/gpt-4o" - custom_llm_provider: str = "openai" + model_config = SettingsConfigDict(env_prefix="KAIZEN_") + tips_model: str = "gpt-4o" + conflict_resolution_model: str = "gpt-4o" + custom_llm_provider: str | None = Field(default=None) + # to reload settings call llm_settings.__init__() -llm_settings = LLMSettings() \ No newline at end of file +llm_settings = LLMSettings() diff --git a/kaizen/frontend/mcp/mcp_server.py b/kaizen/frontend/mcp/mcp_server.py index 2ae49da5..6e2ff8e9 100644 --- a/kaizen/frontend/mcp/mcp_server.py +++ b/kaizen/frontend/mcp/mcp_server.py @@ -3,6 +3,7 @@ This server provides a tool to get task-relevant guidelines. """ + import json import logging import uuid @@ -20,12 +21,14 @@ mcp = FastMCP("katas") client = KaizenClient() + def ensure_namespace(): try: client.get_namespace_details(kaizen_config.namespace_id) except NamespaceNotFoundException: client.create_namespace(kaizen_config.namespace_id) + @mcp.tool() def get_guidelines(task: str) -> str: """ @@ -41,7 +44,7 @@ def get_guidelines(task: str) -> str: results = client.search_entities( namespace_id=kaizen_config.namespace_id, query=task, - filters={"type": "guideline"} + filters={"type": "guideline"}, ) # Format the response @@ -54,7 +57,9 @@ def get_guidelines(task: str) -> str: @mcp.tool() -def save_trajectory(trajectory_data: str, task_id: str | None = None) -> list[RecordedEntity]: +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 @@ -67,33 +72,45 @@ def save_trajectory(trajectory_data: str, task_id: str | None = None) -> list[Re entities = [] messages = json.loads(trajectory_data) for message in messages: - entities.append(Entity( - type='trajectory', - content=message['content'] if isinstance(message['content'], str) else str(message['content']), - metadata={ - "task_id": task_id, - "message": message # store the original message for reference - } - )) + entities.append( + Entity( + type="trajectory", + content=message["content"] + if isinstance(message["content"], str) + else str(message["content"]), + metadata={ + "task_id": task_id, + "message": message, # store the original message for reference + }, + ) + ) client.update_entities( namespace_id=kaizen_config.namespace_id, entities=entities, - enable_conflict_resolution=False + enable_conflict_resolution=False, ) tips = generate_tips(messages) client.update_entities( namespace_id=kaizen_config.namespace_id, - entities=[Entity( - type='guideline', - content=tip, - ) for tip in tips], - enable_conflict_resolution=True + entities=[ + Entity( + type="guideline", + content=tip.content, + metadata={ + "category": tip.category, + "rationale": tip.rationale, + "trigger": tip.trigger, + }, + ) + for tip in tips + ], + enable_conflict_resolution=True, ) return client.search_entities( namespace_id=kaizen_config.namespace_id, filters={"type": "trajectory", "task_id": task_id}, - limit=1000 - ) \ No newline at end of file + limit=1000, + ) diff --git a/kaizen/llm/tips/prompts/generate_tips.jinja2 b/kaizen/llm/tips/prompts/generate_tips.jinja2 index 896ef4c2..9ed35a6b 100644 --- a/kaizen/llm/tips/prompts/generate_tips.jinja2 +++ b/kaizen/llm/tips/prompts/generate_tips.jinja2 @@ -1,7 +1,38 @@ -Examine the following agent trajectory. Extract guidelines or tips that can be used in the future to optimize the trajectory. -This includes remembering sequences of tool calls, or entities that will make the reasoning more efficient and avoid errors. +You are analyzing an AI agent's execution trajectory to extract actionable tips. -The tips should be short sentences that guide future iterations of the agent to perform better or faster, avoid errors, or learn from experience. -Generate as many tips as you think is needed. Output as a JSON formatted list of strings. +# Task Information +**Task:** {{task_instruction}} +**Status:** UNKNOWN +**Steps Taken:** {{num_steps}} -{{markdown_trajectory}} \ No newline at end of file +# Agent Trajectory +{{trajectory_summary}} + +# Your Task +Extract 3-5 actionable tips from this trajectory that would help AI agents perform similar tasks better. + +**Guidelines:** +1. Focus on patterns that worked or mistakes that were made +2. Be specific to what you observed in this trajectory +3. Each tip should have: + - Clear description of what to do (or avoid) + - Why it matters + - When to apply it + +{% if not constrained_decoding_supported %} +**Output Format (JSON):** +```json +{ + "tips": [ + { + "content": "Clear, actionable tip", + "rationale": "Why this tip helps", + "category": "strategy|recovery|optimization", + "trigger": "When to apply this tip" + } + ] +} +``` + +Generate tips now. Return ONLY the JSON, no other text. +{% endif %} \ No newline at end of file diff --git a/kaizen/llm/tips/tips.py b/kaizen/llm/tips/tips.py index e17f8c61..979bbe75 100644 --- a/kaizen/llm/tips/tips.py +++ b/kaizen/llm/tips/tips.py @@ -1,54 +1,159 @@ import json +from json import JSONDecodeError + +import litellm from jinja2 import Template -from litellm import completion +from litellm import completion, get_supported_openai_params, supports_response_schema from kaizen.config.llm import llm_settings from kaizen.utils.utils import clean_llm_response +from kaizen.schema.exceptions import KaizenException +from kaizen.schema.tips import TipGenerationResponse, Tip from pathlib import Path -def generate_tips(messages: list[dict]) -> list[str]: - markdown_trajectory = messages_to_markdown(messages) - prompt_file = Path(__file__).parent / "prompts/generate_tips.jinja2" - prompt = Template(prompt_file.read_text()).render(markdown_trajectory=markdown_trajectory) - response = completion( - model=llm_settings.tips_model, - messages=[{"role": "user", "content": prompt}], - custom_llm_provider=llm_settings.custom_llm_provider - ).choices[0].message.content - clean_response = clean_llm_response(response) - return json.loads(clean_response) -def messages_to_markdown(messages: list[dict]) -> str: +def parse_openai_agents_trajectory(messages: list[dict]) -> dict: """ - Convert a list of OpenAI-format messages to a Markdown string. + Parse OpenAI Agents SDK trajectory from streamer.to_input_list(). + + Returns: + dict with: + - task_instruction: The task description + - agent_steps: List of agent reasoning/actions + - function_calls: List of tool/function calls made + - num_steps: Total number of agent actions """ - md_lines = [] - - for msg in messages: - role = msg.get("role", "unknown").title() - content = msg.get("content", "") - - md_lines.append(f"## {role}") - - if isinstance(content, str): - md_lines.append(content) - elif isinstance(content, list): - for block in content: - if isinstance(block, dict): - if block.get("type") == "text": - md_lines.append(block.get("text", "")) - elif block.get("type") == "function_call": - func = block.get("function", {}) - md_lines.append(f"**Tool Call**: `{func.get('name')}`") - md_lines.append("```json") - md_lines.append(func.get("arguments", "")) - md_lines.append("```") - elif block.get("type") == "function_response": - md_lines.append(f"**Tool Result** ({block.get('id')}):") - md_lines.append("```") - md_lines.append(block.get("content", "")) - md_lines.append("```") - - md_lines.append("") # Empty line between messages - - return "\n".join(md_lines) \ No newline at end of file + agent_steps = [] + function_calls = [] + task_instruction = None + + for message in messages: + # Extract task instruction from first user message + if message.get("role") == "user" and task_instruction is None: + if isinstance(message["content"], str): + task_instruction = message["content"] + else: + raise KaizenException("First user message was not a task instruction.") + + # Extract assistant reasoning/messages + if message.get("role") == "assistant": + content = message.get("content", "") + if isinstance(content, str) and content.strip(): + agent_steps.append( + {"type": "reasoning", "content": content, "raw": message} + ) + + # Extract function calls + elif isinstance(content, list): + for assistant_response in content: + if assistant_response["type"] == "function_call": + function_call = { + "type": "function_call", + "name": assistant_response["function"]["name"], + "arguments": assistant_response["function"]["arguments"], + "call_id": assistant_response["id"], + "raw": assistant_response, + } + function_calls.append(function_call) + + # Add to agent steps as an action + args_str = (assistant_response["function"]["arguments"],) + try: + args = ( + json.loads(args_str) + if isinstance(args_str, str) + else args_str + ) + args_display = ", ".join( + f"{k}={json.dumps(v)}" for k, v in args.items() + ) + except JSONDecodeError: + args_display = args_str + + agent_steps.append( + { + "type": "actions", + "content": f"{assistant_response['function']['name']}({args_display})", + "raw": assistant_response, + } + ) + else: + raise KaizenException( + f"Unhandled assistant content type in list `{assistant_response['type']}`" + ) + else: + raise KaizenException( + f"Unhandled assistant content type `{type(content)}`" + ) + + steps_text = [] + for i, step in enumerate(agent_steps[:50], 1): + step_type = step["type"] + content = step["content"] + # Truncate long content + if len(content) > 2000: + content = content[:2000] + "..." + + if step_type == "reasoning": + steps_text.append(f"**Step {i} - Reasoning:**\n{content}") + elif step_type == "action": + steps_text.append(f"**Step {i} - Action:**\n{content}") + elif step_type == "observation": + steps_text.append(f"**Step {i} - Observation:**\n{content}") + + return { + "task_instruction": task_instruction or "Unknown task", + "trajectory_summary": "\n\n".join(steps_text), + "function_calls": function_calls, + "num_steps": len( + [s for s in agent_steps if s["type"] in ["action", "reasoning"]] + ), + } + + +def generate_tips(messages: list[dict]) -> list[Tip]: + prompt_file = Path(__file__).parent / "prompts/generate_tips.jinja2" + supports_response_format = "response_format" in get_supported_openai_params( + model=llm_settings.tips_model, + custom_llm_provider=llm_settings.custom_llm_provider, + ) + response_schema_enabled = supports_response_schema( + model=llm_settings.tips_model, + custom_llm_provider=llm_settings.custom_llm_provider, + ) + constrained_decoding_supported = ( + supports_response_format and response_schema_enabled + ) + trajectory_data = parse_openai_agents_trajectory(messages) + prompt = Template(prompt_file.read_text()).render( + task_instruction=trajectory_data["task_instruction"], + num_steps=trajectory_data["num_steps"], + trajectory_summary=trajectory_data["trajectory_summary"], + constrained_decoding_supported=constrained_decoding_supported, + ) + + if constrained_decoding_supported: + litellm.enable_json_schema_validation = True + clean_response = ( + completion( + model=llm_settings.tips_model, + messages=[{"role": "user", "content": prompt}], + response_format=TipGenerationResponse, + custom_llm_provider=llm_settings.custom_llm_provider, + ) + .choices[0] + .message.content + ) + else: + litellm.enable_json_schema_validation = False + response = ( + completion( + model=llm_settings.tips_model, + messages=[{"role": "user", "content": prompt}], + custom_llm_provider=llm_settings.custom_llm_provider, + ) + .choices[0] + .message.content + ) + clean_response = clean_llm_response(response) + return TipGenerationResponse.model_validate(json.loads(clean_response)).tips diff --git a/kaizen/schema/tips.py b/kaizen/schema/tips.py new file mode 100644 index 00000000..03e7d716 --- /dev/null +++ b/kaizen/schema/tips.py @@ -0,0 +1,13 @@ +from pydantic import BaseModel, Field +from typing import Literal + + +class Tip(BaseModel): + content: str = Field(description="Clear, actionable tip") + rationale: str = Field(description="Why this tip helps") + category: Literal["strategy", "recovery", "optimization"] + trigger: str = Field("When to apply this tip") + + +class TipGenerationResponse(BaseModel): + tips: list[Tip] diff --git a/kaizen/sync/phoenix_sync.py b/kaizen/sync/phoenix_sync.py index 14f7c187..e4585b90 100644 --- a/kaizen/sync/phoenix_sync.py +++ b/kaizen/sync/phoenix_sync.py @@ -28,6 +28,7 @@ @dataclass class SyncResult: """Result of a sync operation.""" + processed: int skipped: int tips_generated: int @@ -87,7 +88,7 @@ def _get_processed_span_ids(self) -> set[str]: entities = self.client.search_entities( namespace_id=self.namespace_id, filters={"type": "trajectory"}, - limit=10000 + limit=10000, ) return { e.metadata.get("span_id") @@ -105,6 +106,7 @@ def _parse_content(self, content: Any) -> Any: except json.JSONDecodeError: try: import ast + return ast.literal_eval(content) except (ValueError, SyntaxError): return content @@ -126,12 +128,14 @@ def _extract_messages_from_span(self, span: dict) -> list[dict]: role = attrs.get(f"gen_ai.prompt.{i}.role") content = attrs.get(f"gen_ai.prompt.{i}.content") if role and content is not None: - messages.append({ - "index": i, - "type": "prompt", - "role": role, - "content": self._parse_content(content) - }) + messages.append( + { + "index": i, + "type": "prompt", + "role": role, + "content": self._parse_content(content), + } + ) # Extract completion messages completion_indices = set() @@ -144,12 +148,14 @@ def _extract_messages_from_span(self, span: dict) -> list[dict]: role = attrs.get(f"gen_ai.completion.{i}.role") content = attrs.get(f"gen_ai.completion.{i}.content") if role and content is not None: - messages.append({ - "index": i, - "type": "completion", - "role": role, - "content": self._parse_content(content) - }) + messages.append( + { + "index": i, + "type": "completion", + "role": role, + "content": self._parse_content(content), + } + ) return messages @@ -184,21 +190,25 @@ def _convert_to_openai_format(self, content: Any, role: str) -> dict: thinking_parts.append(thinking) elif block_type == "tool_use": - tool_calls.append({ - "id": block.get("id", ""), - "type": "function", - "function": { - "name": block.get("name", ""), - "arguments": json.dumps(block.get("input", {})) + tool_calls.append( + { + "id": block.get("id", ""), + "type": "function", + "function": { + "name": block.get("name", ""), + "arguments": json.dumps(block.get("input", {})), + }, } - }) + ) elif block_type == "tool_result": - tool_results.append({ - "tool_call_id": block.get("tool_use_id", ""), - "content": block.get("content", ""), - "is_error": block.get("is_error", False) - }) + tool_results.append( + { + "tool_call_id": block.get("tool_use_id", ""), + "content": block.get("content", ""), + "is_error": block.get("is_error", False), + } + ) if role == "assistant": msg = {"role": "assistant"} @@ -233,11 +243,13 @@ def _extract_trajectory(self, span: dict) -> dict: if converted.get("role") == "tool" and "tool_results" in converted: for result in converted["tool_results"]: - openai_messages.append({ - "role": "tool", - "tool_call_id": result["tool_call_id"], - "content": result["content"] - }) + openai_messages.append( + { + "role": "tool", + "tool_call_id": result["tool_call_id"], + "content": result["content"], + } + ) else: openai_messages.append(converted) @@ -250,13 +262,14 @@ def _extract_trajectory(self, span: dict) -> dict: "usage": { "prompt_tokens": attrs.get("gen_ai.usage.prompt_tokens"), "completion_tokens": attrs.get("gen_ai.usage.completion_tokens"), - "total_tokens": attrs.get("llm.usage.total_tokens") - } + "total_tokens": attrs.get("llm.usage.total_tokens"), + }, } def _clean_trajectory(self, trajectory: dict) -> dict: """Clean up a trajectory by removing system reminders.""" import re + cleaned_messages = [] for msg in trajectory.get("messages", []): @@ -267,10 +280,10 @@ def _clean_trajectory(self, trajectory: dict) -> dict: content = msg["content"] if isinstance(content, str): content = re.sub( - r'.*?', - '', + r".*?", + "", content, - flags=re.DOTALL + flags=re.DOTALL, ).strip() if not content: continue @@ -290,23 +303,25 @@ def _process_trajectory(self, trajectory: dict) -> int: for msg in trajectory.get("messages", []): content = msg.get("content") if isinstance(content, str) and content: - entities.append(Entity( - type='trajectory', - content=content, - metadata={ - "trace_id": trajectory["trace_id"], - "span_id": trajectory["span_id"], - "model": trajectory["model"], - "role": msg.get("role"), - "timestamp": trajectory["timestamp"], - } - )) + entities.append( + Entity( + type="trajectory", + content=content, + metadata={ + "trace_id": trajectory["trace_id"], + "span_id": trajectory["span_id"], + "model": trajectory["model"], + "role": msg.get("role"), + "timestamp": trajectory["timestamp"], + }, + ) + ) if entities: self.client.update_entities( namespace_id=self.namespace_id, entities=entities, - enable_conflict_resolution=False + enable_conflict_resolution=False, ) # Generate tips from the trajectory @@ -315,19 +330,22 @@ def _process_trajectory(self, trajectory: dict) -> int: if tips: tip_entities = [ Entity( - type='guideline', - content=tip, + type="guideline", + content=tip.content, metadata={ + "category": tip.category, + "rationale": tip.rationale, + "trigger": tip.trigger, "source_trace_id": trajectory["trace_id"], "source_span_id": trajectory["span_id"], - } + }, ) for tip in tips ] self.client.update_entities( namespace_id=self.namespace_id, entities=tip_entities, - enable_conflict_resolution=True + enable_conflict_resolution=True, ) return len(tips) @@ -347,7 +365,9 @@ def sync( Returns: SyncResult with counts of processed, skipped, and tips generated """ - logger.info(f"Starting sync from {self.phoenix_url} to namespace '{self.namespace_id}'") + logger.info( + f"Starting sync from {self.phoenix_url} to namespace '{self.namespace_id}'" + ) self._ensure_namespace() @@ -405,7 +425,7 @@ def sync( processed=processed, skipped=skipped, tips_generated=tips_generated, - errors=errors + errors=errors, ) logger.info( diff --git a/pyproject.toml b/pyproject.toml index d08abdc3..900f897f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,6 +42,6 @@ addopts = "--ignore=explorations -m 'not phoenix'" markers = [ "e2e", "unit", - "phoenix: tests requiring Phoenix sync functionality (deselected by default)", + "phoenix" ] anyio_mode = "auto" \ No newline at end of file From 7966a5736e4b95333e7f8ab72b0998cf7ee23ef4 Mon Sep 17 00:00:00 2001 From: Punleuk Oum Date: Fri, 16 Jan 2026 09:31:16 -0500 Subject: [PATCH 2/3] Update tips.py --- kaizen/schema/tips.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/kaizen/schema/tips.py b/kaizen/schema/tips.py index 03e7d716..dabdf84d 100644 --- a/kaizen/schema/tips.py +++ b/kaizen/schema/tips.py @@ -6,7 +6,7 @@ class Tip(BaseModel): content: str = Field(description="Clear, actionable tip") rationale: str = Field(description="Why this tip helps") category: Literal["strategy", "recovery", "optimization"] - trigger: str = Field("When to apply this tip") + trigger: str = Field(description="When to apply this tip") class TipGenerationResponse(BaseModel): From ad7e158bb092d3a72175fa8c1132f883ccb32091 Mon Sep 17 00:00:00 2001 From: Punleuk Oum Date: Fri, 16 Jan 2026 09:35:57 -0500 Subject: [PATCH 3/3] you caught me --- kaizen/llm/tips/tips.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/kaizen/llm/tips/tips.py b/kaizen/llm/tips/tips.py index 979bbe75..0d8b96d2 100644 --- a/kaizen/llm/tips/tips.py +++ b/kaizen/llm/tips/tips.py @@ -57,23 +57,22 @@ def parse_openai_agents_trajectory(messages: list[dict]) -> dict: function_calls.append(function_call) # Add to agent steps as an action - args_str = (assistant_response["function"]["arguments"],) + args_str = assistant_response["function"]["arguments"] try: - args = ( - json.loads(args_str) - if isinstance(args_str, str) - else args_str - ) + args: dict = json.loads(args_str) args_display = ", ".join( f"{k}={json.dumps(v)}" for k, v in args.items() ) + function_description = f"{assistant_response['function']['name']}({args_display})" except JSONDecodeError: - args_display = args_str + function_description = ( + f"{assistant_response['function']['name']}({args_str})" + ) agent_steps.append( { - "type": "actions", - "content": f"{assistant_response['function']['name']}({args_display})", + "type": "action", + "content": function_description, "raw": assistant_response, } )