Tips Improvement

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`):
 

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()
+llm_settings = LLMSettings()

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
-    )
+        limit=1000,
+    )

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}}
+# 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 %}

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)
+    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