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..0d8b96d2 100644
--- a/kaizen/llm/tips/tips.py
+++ b/kaizen/llm/tips/tips.py
@@ -1,54 +1,158 @@
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: 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:
+ function_description = (
+ f"{assistant_response['function']['name']}({args_str})"
+ )
+
+ agent_steps.append(
+ {
+ "type": "action",
+ "content": function_description,
+ "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..dabdf84d
--- /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(description="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