diff --git a/.secrets.baseline b/.secrets.baseline index 4ae4e63a..971c3820 100644 --- a/.secrets.baseline +++ b/.secrets.baseline @@ -3,7 +3,7 @@ "files": "^.secrets.baseline$", "lines": null }, - "generated_at": "2026-01-14T17:28:21Z", + "generated_at": "2026-01-15T18:30:24Z", "plugins_used": [ { "name": "AWSKeyDetector" @@ -95,6 +95,34 @@ "verified_result": null } ], + "README_extract_trajectories.md": [ + { + "hashed_secret": "48ae8db8c64fbf47ec3f8eec6f31f3a5b88594d2", + "is_secret": false, + "is_verified": false, + "line_number": 93, + "type": "Hex High Entropy String", + "verified_result": null + }, + { + "hashed_secret": "c4f1b0e724abd250d38f29d90b68c3f832ccb6a6", + "is_secret": false, + "is_verified": false, + "line_number": 94, + "type": "Hex High Entropy String", + "verified_result": null + } + ], + "README_phoenix_sync.md": [ + { + "hashed_secret": "ec3810e10fb78db55ce38b9c18d1c3eb1db739e0", + "is_secret": false, + "is_verified": false, + "line_number": 200, + "type": "Secret Keyword", + "verified_result": null + } + ], "explorations/claudecode/README.md": [ { "hashed_secret": "c7a8c334eef5d1749fface7d42c66f9ae5e8cf36", @@ -148,6 +176,16 @@ "type": "Base64 High Entropy String", "verified_result": null } + ], + "tests/unit/test_extract_trajectories.py": [ + { + "hashed_secret": "90bd1b48e958257948487b90bee080ba5ed00caa", + "is_secret": false, + "is_verified": false, + "line_number": 486, + "type": "Hex High Entropy String", + "verified_result": null + } ] }, "version": "0.13.1+ibm.64.dss", diff --git a/README.md b/README.md index 352e35d3..6310eb7b 100644 --- a/README.md +++ b/README.md @@ -210,8 +210,30 @@ Kaizen includes a command-line interface for managing namespaces and entities di ## Development -To run tests: +### Running Tests + +Run the default test suite: + +```bash +uv run pytest +``` + +#### Phoenix Sync Tests + +Tests for the Phoenix trajectory sync functionality are **skipped by default** since they require familiarity with the Phoenix integration. To include them: ```bash +# Run all tests including Phoenix tests +uv run pytest --run-phoenix + +# Run only Phoenix tests +uv run pytest -m phoenix + +# Run default tests (excludes Phoenix) uv run pytest ``` + +The Phoenix tests cover: +- `kaizen/sync/phoenix_sync.py` - Trajectory sync from Arize Phoenix +- `extract_trajectories.py` - Standalone trajectory extraction script +- CLI `sync phoenix` command diff --git a/README_extract_trajectories.md b/README_extract_trajectories.md new file mode 100644 index 00000000..9502bbab --- /dev/null +++ b/README_extract_trajectories.md @@ -0,0 +1,226 @@ +# Extract Trajectories from Arize Phoenix + +A Python tool to extract agent trajectories from Arize Phoenix traces and convert them to OpenAI chat completion message format. + +## Features + +- Fetches spans from Phoenix's REST API with pagination support +- Converts Anthropic/Claude message format to OpenAI-compatible format +- Preserves agent reasoning/thinking in a dedicated field +- Handles tool calls and tool responses +- Cleans system reminders and noise from messages +- Supports both JSON and human-readable text output + +## Requirements + +- Python 3.10+ +- Arize Phoenix running locally (default: `http://localhost:6006`) +- No external dependencies (uses only stdlib) + +## Usage + +### Command Line + +```bash +# Get recent trajectories as JSON (pretty-printed) +python3 extract_trajectories.py --limit 100 --pretty + +# Get a specific trace by ID +python3 extract_trajectories.py --trace-id a67264f4375a60a1ac26ef3061c7352d --pretty + +# Human-readable text format +python3 extract_trajectories.py --text --trace-id a67264f4375a60a1ac26ef3061c7352d + +# Include failed/error spans +python3 extract_trajectories.py --include-errors --pretty + +# Keep system reminders (don't clean) +python3 extract_trajectories.py --no-clean --pretty + +# Save to file +python3 extract_trajectories.py --limit 100 -o trajectories.json --pretty + +# Custom Phoenix URL +python3 extract_trajectories.py --url http://localhost:8080 --pretty +``` + +### CLI Options + +| Option | Description | +|--------|-------------| +| `--url` | Phoenix server URL (default: `http://localhost:6006`) | +| `--limit` | Maximum number of spans to fetch (default: 100) | +| `--include-errors` | Include failed/error spans | +| `--no-clean` | Don't remove system reminders from messages | +| `--output, -o` | Output file path (default: stdout) | +| `--pretty` | Pretty-print JSON output | +| `--trace-id` | Filter to a specific trace ID | +| `--text` | Output as human-readable text instead of JSON | + +### As a Library + +```python +from extract_trajectories import ( + get_trajectories, + get_trajectory_by_trace_id, + format_trajectory_as_text +) + +# Get all recent trajectories +trajectories = get_trajectories( + base_url="http://localhost:6006", + limit=100, + include_errors=False, + clean=True +) + +# Get a specific trajectory by trace ID +trajectory = get_trajectory_by_trace_id("a67264f4375a60a1ac26ef3061c7352d") + +# Format for human-readable display +if trajectory: + print(format_trajectory_as_text(trajectory, include_thinking=True)) +``` + +## Output Format + +### JSON Output + +Each trajectory is converted to an OpenAI-compatible format: + +```json +{ + "trace_id": "a67264f4375a60a1ac26ef3061c7352d", + "span_id": "c2b544103b5db193", + "model": "aws/claude-sonnet-4-5", + "timestamp": "2026-01-13T19:47:05.287716+00:00", + "messages": [ + { + "role": "user", + "content": "What states do I have teammates in?" + }, + { + "role": "assistant", + "thinking": "The user is asking me to read a file...", + "content": "I'll help you find out what states your teammates are in.", + "tool_calls": [ + { + "id": "tooluse_nIHmKHC8rZ3UetpCPKjbbH", + "type": "function", + "function": { + "name": "Read", + "arguments": "{\"file_path\": \"/path/to/states.txt\"}" + } + } + ] + }, + { + "role": "tool", + "tool_call_id": "tooluse_nIHmKHC8rZ3UetpCPKjbbH", + "content": "texas\nnew york\nmassachusetts" + }, + { + "role": "assistant", + "content": "Based on the file, you have teammates in Texas, New York, and Massachusetts." + } + ], + "usage": { + "prompt_tokens": 18282, + "completion_tokens": 42, + "total_tokens": 18324 + } +} +``` + +### Message Types + +| Role | Description | +|------|-------------| +| `user` | User input messages | +| `assistant` | Agent responses (may include `thinking` and `tool_calls`) | +| `tool` | Tool execution results (includes `tool_call_id`) | + +### Non-Standard Fields + +The `thinking` field on assistant messages preserves Claude's chain-of-thought reasoning. This is not part of the OpenAI spec but is useful for: +- Debugging agent behavior +- Understanding decision-making +- Training/fine-tuning datasets + +### Text Output + +The `--text` flag produces human-readable output: + +``` +=== Trajectory: a67264f4375a... === +Model: aws/claude-sonnet-4-5 +Timestamp: 2026-01-13T19:47:05.287716+00:00 + +[USER] +What states do I have teammates in? + +[ASSISTANT] + +The user is asking me to read a file... + + +I'll help you find out what states your teammates are in. + -> Tool call: Read + Args: {"file_path": "/path/to/states.txt"} + +[TOOL RESULT] (id: tooluse_nIHmKHC8rZ3U...) +texas +new york +massachusetts + +[ASSISTANT] +Based on the file, you have teammates in Texas, New York, and Massachusetts. +``` + +## API Reference + +### `get_trajectories(base_url, limit, include_errors, clean) -> list[dict]` + +Fetch and extract agent trajectories from Phoenix. + +**Parameters:** +- `base_url` (str): Phoenix server URL. Default: `"http://localhost:6006"` +- `limit` (int): Maximum spans to fetch. Default: `100` +- `include_errors` (bool): Include failed spans. Default: `False` +- `clean` (bool): Remove system reminders. Default: `True` + +**Returns:** List of trajectory dictionaries. + +### `get_trajectory_by_trace_id(trace_id, base_url) -> dict | None` + +Get a single trajectory by its trace ID. + +**Parameters:** +- `trace_id` (str): The trace ID to look up +- `base_url` (str): Phoenix server URL. Default: `"http://localhost:6006"` + +**Returns:** Trajectory dict or `None` if not found. + +### `format_trajectory_as_text(trajectory, include_thinking) -> str` + +Format a trajectory as human-readable text. + +**Parameters:** +- `trajectory` (dict): The trajectory dictionary +- `include_thinking` (bool): Include agent reasoning. Default: `True` + +**Returns:** Formatted string. + +## Phoenix API + +This tool uses the Phoenix REST API endpoint: + +``` +GET /v1/projects/default/spans?limit=N&cursor=CURSOR +``` + +The spans contain attributes like: +- `gen_ai.prompt.N.role` / `gen_ai.prompt.N.content` - Input messages +- `gen_ai.completion.N.role` / `gen_ai.completion.N.content` - Output messages +- `gen_ai.request.model` - Model used +- `gen_ai.usage.*` - Token usage statistics diff --git a/README_phoenix_sync.md b/README_phoenix_sync.md new file mode 100644 index 00000000..c2a15d20 --- /dev/null +++ b/README_phoenix_sync.md @@ -0,0 +1,201 @@ +# Phoenix Sync + +Sync agent trajectories from Arize Phoenix to Kaizen and automatically generate tips/guidelines. + +## Overview + +The Phoenix sync module: +1. Fetches agent trajectories from Phoenix's REST API +2. Deduplicates already-processed trajectories (by `span_id`) +3. Converts messages to OpenAI format +4. Generates tips/guidelines using LLM +5. Stores both trajectories and tips in Kaizen + +## Installation + +No additional dependencies required - uses only stdlib for Phoenix API calls. + +## Configuration + +### Environment Variables + +| Variable | Default | Description | +|----------|---------|-------------| +| `PHOENIX_URL` | `http://localhost:6006` | Phoenix server URL | +| `PHOENIX_PROJECT` | `default` | Phoenix project name | +| `KAIZEN_NAMESPACE_ID` | `kaizen` | Target namespace for stored entities | +| `KAIZEN_PROVIDER` | `milvus` | Backend provider (`milvus` or `filesystem`) | + +## Usage + +### CLI + +```bash +# Basic sync with defaults +uv run python -m kaizen.frontend.cli.cli sync phoenix + +# Custom Phoenix URL and namespace +uv run python -m kaizen.frontend.cli.cli sync phoenix \ + --url http://phoenix.example.com:6006 \ + --namespace my_namespace + +# Fetch more spans and include errors +uv run python -m kaizen.frontend.cli.cli sync phoenix \ + --limit 500 \ + --include-errors + +# Full options +uv run python -m kaizen.frontend.cli.cli sync phoenix \ + --url http://localhost:6006 \ + --namespace production \ + --project my_project \ + --limit 200 \ + --include-errors +``` + +### CLI Options + +| Option | Short | Description | +|--------|-------|-------------| +| `--url` | `-u` | Phoenix server URL | +| `--namespace` | `-n` | Target Kaizen namespace | +| `--project` | `-p` | Phoenix project name | +| `--limit` | | Max spans to fetch (default: 100) | +| `--include-errors` | | Include failed/error spans | + +### Python API + +```python +from kaizen.sync.phoenix_sync import PhoenixSync + +# Initialize syncer +syncer = PhoenixSync( + phoenix_url="http://localhost:6006", + namespace_id="my_namespace", + project="default" +) + +# Run sync +result = syncer.sync(limit=100, include_errors=False) + +print(f"Processed: {result.processed}") +print(f"Skipped: {result.skipped}") +print(f"Tips generated: {result.tips_generated}") +print(f"Errors: {result.errors}") +``` + +## How It Works + +### 1. Fetch Spans + +The syncer calls Phoenix's REST API: +``` +GET /v1/projects/{project}/spans?limit=N&cursor=CURSOR +``` + +Only `litellm_request` spans with prompt messages are processed. + +### 2. Deduplication + +Each processed trajectory stores `span_id` in its metadata. On subsequent syncs, already-processed span IDs are skipped. + +### 3. Message Conversion + +Anthropic/Claude message format is converted to OpenAI format: +- Tool use blocks → `tool_calls` array +- Tool results → separate `tool` role messages +- Thinking blocks → preserved in `thinking` field + +### 4. Tip Generation + +The `generate_tips()` function analyzes the trajectory and produces actionable guidelines using an LLM. + +### 5. Storage + +Two entity types are stored: +- `trajectory` - Individual messages with metadata (trace_id, span_id, model, role) +- `guideline` - Generated tips with conflict resolution enabled + +## Data Flow + +``` +┌─────────────┐ fetch ┌───────────────────┐ +│ Phoenix │ ───────────────→│ PhoenixSync │ +│ (spans) │ │ │ +└─────────────┘ └─────────┬─────────┘ + │ + ┌─────────────────────┼─────────────────────┐ + │ │ │ + ▼ ▼ ▼ + ┌───────────────┐ ┌───────────────┐ ┌───────────────┐ + │ Deduplicate │ │ Convert to │ │ Generate │ + │ (by span_id) │ │ OpenAI format │ │ Tips │ + └───────────────┘ └───────────────┘ └───────┬───────┘ + │ + ▼ + ┌───────────────┐ + │ Kaizen │ + │ Backend │ + └───────────────┘ +``` + +## Running on a Schedule + +### Cron + +```bash +# Sync every hour +0 * * * * cd /path/to/kaizen && uv run python -m kaizen.frontend.cli.cli sync phoenix --limit 100 +``` + +### Systemd Timer + +```ini +# /etc/systemd/system/kaizen-sync.service +[Unit] +Description=Kaizen Phoenix Sync + +[Service] +Type=oneshot +WorkingDirectory=/path/to/kaizen +ExecStart=/path/to/uv run python -m kaizen.frontend.cli.cli sync phoenix +Environment=PHOENIX_URL=http://localhost:6006 +Environment=KAIZEN_NAMESPACE_ID=production +``` + +```ini +# /etc/systemd/system/kaizen-sync.timer +[Unit] +Description=Run Kaizen Phoenix Sync hourly + +[Timer] +OnCalendar=hourly +Persistent=true + +[Install] +WantedBy=timers.target +``` + +## Troubleshooting + +### Connection refused + +Ensure Phoenix is running and accessible at the configured URL: +```bash +curl http://localhost:6006/v1/projects/default/spans?limit=1 +``` + +### No spans processed + +- Check that spans have `name="litellm_request"` +- Verify spans contain `gen_ai.prompt.*` attributes +- Use `--include-errors` to include failed spans + +### Tips not generating + +Ensure LLM API key is configured: +```bash +export OPENAI_API_KEY=sk-... +# or for other providers +export ANTHROPIC_API_KEY=sk-... +``` diff --git a/extract_trajectories.py b/extract_trajectories.py new file mode 100755 index 00000000..e9b6e245 --- /dev/null +++ b/extract_trajectories.py @@ -0,0 +1,495 @@ +#!/usr/bin/env python3 +""" +Extract agent trajectories from Arize Phoenix and convert to OpenAI chat completion format. + +This script fetches spans from Phoenix traces and transforms them into a format +compatible with OpenAI's chat completion messages, including: +- User utterances +- Agent reasoning (thinking) +- Tool calls +- Tool responses +- Agent responses +""" + +import json +import argparse +from collections import defaultdict +from typing import Any +import urllib.request + + +def fetch_spans(base_url: str, limit: int = 1000) -> list[dict]: + """Fetch all spans from Phoenix, handling pagination.""" + spans = [] + cursor = None + + while True: + url = f"{base_url}/v1/projects/default/spans?limit={min(limit - len(spans), 100)}" + if cursor: + url += f"&cursor={cursor}" + + with urllib.request.urlopen(url) as response: + data = json.loads(response.read().decode()) + + spans.extend(data.get("data", [])) + cursor = data.get("next_cursor") + + if not cursor or len(spans) >= limit: + break + + return spans + + +def parse_content(content: Any) -> Any: + """Parse content which may be a string representation of a list/dict.""" + if isinstance(content, str): + try: + return json.loads(content) + except json.JSONDecodeError: + # Try to parse as Python literal + try: + import ast + return ast.literal_eval(content) + except (ValueError, SyntaxError): + return content + return content + + +def extract_messages_from_span(span: dict) -> list[dict]: + """Extract messages from a single span's attributes.""" + attrs = span.get("attributes", {}) + messages = [] + + # Extract prompt messages + prompt_indices = set() + for key in attrs: + if key.startswith("gen_ai.prompt.") and key.endswith(".role"): + idx = int(key.split(".")[2]) + prompt_indices.add(idx) + + for i in sorted(prompt_indices): + 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": parse_content(content) + }) + + # Extract completion messages + completion_indices = set() + for key in attrs: + if key.startswith("gen_ai.completion.") and key.endswith(".role"): + idx = int(key.split(".")[2]) + completion_indices.add(idx) + + for i in sorted(completion_indices): + 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": parse_content(content) + }) + + return messages + + +def convert_anthropic_to_openai(content: Any, role: str) -> dict: + """Convert Anthropic message format to OpenAI format.""" + + if isinstance(content, str): + return {"role": role, "content": content} + + if not isinstance(content, list): + return {"role": role, "content": str(content)} + + # Process list of content blocks + text_parts = [] + tool_calls = [] + tool_results = [] + thinking_parts = [] + + for block in content: + if not isinstance(block, dict): + text_parts.append(str(block)) + continue + + block_type = block.get("type") + + if block_type == "text": + text = block.get("text", "") + if text and text != "(no content)": + text_parts.append(text) + + elif block_type == "thinking": + thinking = block.get("thinking", "") + if thinking: + 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", {})) + } + }) + + 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) + }) + + # Build OpenAI format message + if role == "assistant": + msg = {"role": "assistant"} + + # Include thinking as a separate field (non-standard but useful) + if thinking_parts: + msg["thinking"] = "\n\n".join(thinking_parts) + + if text_parts: + msg["content"] = "\n\n".join(text_parts) + elif not tool_calls: + msg["content"] = None + + if tool_calls: + msg["tool_calls"] = tool_calls + + return msg + + elif role == "user" and tool_results: + # Tool results come back as "user" role in Anthropic format + # In OpenAI format, each tool result is a separate message + return {"role": "tool", "tool_results": tool_results} + + else: + # Regular user message + content_text = "\n\n".join(text_parts) if text_parts else "" + return {"role": role, "content": content_text} + + +def extract_trajectory(span: dict) -> dict: + """Extract a complete trajectory from a span.""" + attrs = span.get("attributes", {}) + messages = extract_messages_from_span(span) + + openai_messages = [] + + for msg in messages: + role = msg["role"] + content = msg["content"] + + converted = convert_anthropic_to_openai(content, role) + + # Handle tool results (expand into individual messages) + 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"] + }) + else: + openai_messages.append(converted) + + # Add the completion if not already included + if messages and messages[-1]["type"] != "completion": + completion_indices = set() + for key in attrs: + if key.startswith("gen_ai.completion.") and key.endswith(".role"): + idx = int(key.split(".")[2]) + completion_indices.add(idx) + + for i in sorted(completion_indices): + role = attrs.get(f"gen_ai.completion.{i}.role") + content = attrs.get(f"gen_ai.completion.{i}.content") + if role and content: + converted = convert_anthropic_to_openai(parse_content(content), role) + openai_messages.append(converted) + + return { + "trace_id": span["context"]["trace_id"], + "span_id": span["context"]["span_id"], + "model": attrs.get("gen_ai.request.model", "unknown"), + "timestamp": span.get("start_time"), + "messages": openai_messages, + "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") + } + } + + +def filter_system_reminders(text: str) -> str: + """Remove system reminders from text content.""" + import re + return re.sub(r'.*?', '', text, flags=re.DOTALL).strip() + + +def clean_trajectory(trajectory: dict, remove_system_reminders: bool = True) -> dict: + """Clean up a trajectory by removing noise and system messages.""" + cleaned_messages = [] + + for msg in trajectory.get("messages", []): + # Skip empty messages + if not msg.get("content") and not msg.get("tool_calls"): + continue + + # Clean content + if remove_system_reminders and msg.get("content"): + content = msg["content"] + if isinstance(content, str): + content = filter_system_reminders(content) + if not content: + continue + msg = {**msg, "content": content} + + cleaned_messages.append(msg) + + return {**trajectory, "messages": cleaned_messages} + + +def get_trajectories( + base_url: str = "http://localhost:6006", + limit: int = 100, + include_errors: bool = False, + clean: bool = True +) -> list[dict]: + """ + Fetch and extract agent trajectories from Phoenix. + + Args: + base_url: Phoenix server URL + limit: Maximum number of spans to fetch + include_errors: Whether to include failed spans + clean: Whether to clean system reminders from messages + + Returns: + List of trajectories in OpenAI chat completion format + """ + spans = fetch_spans(base_url, limit) + + trajectories = [] + for span in spans: + # Filter to LLM request spans + if span.get("name") != "litellm_request": + continue + + # Filter errors if requested + if not include_errors and span.get("status_code") == "ERROR": + continue + + # Only include spans with actual messages + attrs = span.get("attributes", {}) + if not any(k.startswith("gen_ai.prompt.") for k in attrs): + continue + + trajectory = extract_trajectory(span) + + if clean: + trajectory = clean_trajectory(trajectory) + + # Only include if there are meaningful messages + if trajectory["messages"]: + trajectories.append(trajectory) + + return trajectories + + +def get_trajectory_by_trace_id(trace_id: str, base_url: str = "http://localhost:6006") -> dict | None: + """ + Convenience function to get a single trajectory by trace ID. + + Args: + trace_id: The trace ID to look up + base_url: Phoenix server URL + + Returns: + Trajectory dict or None if not found + """ + trajectories = get_trajectories( + base_url=base_url, + limit=1000, + include_errors=True, + clean=True + ) + for t in trajectories: + if t["trace_id"] == trace_id: + return t + return None + + +def format_trajectory_as_text(trajectory: dict, include_thinking: bool = True) -> str: + """ + Format a trajectory as human-readable text. + + Args: + trajectory: The trajectory dict + include_thinking: Whether to include agent thinking/reasoning + + Returns: + Formatted string representation + """ + lines = [] + lines.append(f"=== Trajectory: {trajectory['trace_id'][:12]}... ===") + lines.append(f"Model: {trajectory['model']}") + lines.append(f"Timestamp: {trajectory['timestamp']}") + lines.append("") + + # Build a mapping of tool_call_id to tool name and arguments for reference + tool_call_map = {} + for msg in trajectory.get("messages", []): + if msg.get("tool_calls"): + for tc in msg["tool_calls"]: + func = tc.get("function", {}) + tool_call_map[tc.get("id", "")] = { + "name": func.get("name", "unknown"), + "arguments": func.get("arguments", "{}") + } + + for msg in trajectory.get("messages", []): + role = msg.get("role", "unknown").upper() + + if role == "USER": + lines.append(f"[USER]") + lines.append(msg.get("content", "")) + lines.append("") + + elif role == "ASSISTANT": + lines.append(f"[ASSISTANT]") + if include_thinking and msg.get("thinking"): + lines.append(f"") + lines.append(msg["thinking"][:500] + "..." if len(msg.get("thinking", "")) > 500 else msg.get("thinking", "")) + lines.append("") + lines.append("") + if msg.get("content"): + lines.append(msg["content"]) + if msg.get("tool_calls"): + for tc in msg["tool_calls"]: + func = tc.get("function", {}) + tool_name = func.get("name", "unknown") + tool_id = tc.get("id", "unknown") + lines.append(f" -> Tool call: {tool_name} (id: {tool_id[:20]}...)") + args = func.get("arguments", "{}") + # Pretty print JSON arguments if possible + try: + args_obj = json.loads(args) + args = json.dumps(args_obj, indent=4) + except (json.JSONDecodeError, TypeError): + pass + lines.append(f" Arguments:") + for arg_line in args.split("\n"): + lines.append(f" {arg_line}") + lines.append("") + + elif role == "TOOL": + tool_call_id = msg.get("tool_call_id", "unknown") + tool_info = tool_call_map.get(tool_call_id, {}) + tool_name = tool_info.get("name", "unknown") + lines.append(f"[TOOL RESULT] {tool_name} (id: {tool_call_id[:20]}...)") + content = msg.get("content", "") + # Try to pretty print JSON content + try: + content_obj = json.loads(content) + content = json.dumps(content_obj, indent=2) + except (json.JSONDecodeError, TypeError): + pass + lines.append(f" Response:") + for content_line in content.split("\n")[:50]: # Limit to 50 lines + lines.append(f" {content_line}") + if content.count("\n") > 50: + lines.append(f" ... (truncated, {content.count(chr(10)) - 50} more lines)") + lines.append("") + + return "\n".join(lines) + + +def main(): + parser = argparse.ArgumentParser( + description="Extract agent trajectories from Arize Phoenix" + ) + parser.add_argument( + "--url", + default="http://localhost:6006", + help="Phoenix server URL" + ) + parser.add_argument( + "--limit", + type=int, + default=100, + help="Maximum number of spans to fetch" + ) + parser.add_argument( + "--include-errors", + action="store_true", + help="Include failed spans" + ) + parser.add_argument( + "--no-clean", + action="store_true", + help="Don't clean system reminders from messages" + ) + parser.add_argument( + "--output", + "-o", + help="Output file (default: stdout)" + ) + parser.add_argument( + "--pretty", + action="store_true", + help="Pretty print JSON output" + ) + parser.add_argument( + "--trace-id", + help="Filter to specific trace ID" + ) + parser.add_argument( + "--text", + action="store_true", + help="Output as human-readable text instead of JSON" + ) + + args = parser.parse_args() + + trajectories = get_trajectories( + base_url=args.url, + limit=args.limit, + include_errors=args.include_errors, + clean=not args.no_clean + ) + + if args.trace_id: + trajectories = [t for t in trajectories if t["trace_id"] == args.trace_id] + + # Sort by timestamp (most recent first) + trajectories.sort(key=lambda t: t.get("timestamp", ""), reverse=True) + + if args.text: + output = "\n\n".join(format_trajectory_as_text(t) for t in trajectories) + else: + output = json.dumps( + trajectories, + indent=2 if args.pretty else None, + default=str + ) + + if args.output: + with open(args.output, "w") as f: + f.write(output) + print(f"Wrote {len(trajectories)} trajectories to {args.output}") + else: + print(output) + + +if __name__ == "__main__": + main() diff --git a/kaizen/config/phoenix.py b/kaizen/config/phoenix.py new file mode 100644 index 00000000..a0564d5b --- /dev/null +++ b/kaizen/config/phoenix.py @@ -0,0 +1,11 @@ +from pydantic import Field +from pydantic_settings import BaseSettings, SettingsConfigDict + + +class PhoenixSettings(BaseSettings): + model_config = SettingsConfigDict(env_prefix='PHOENIX_') + url: str = Field(default='http://localhost:6006', description='Phoenix server URL') + project: str = Field(default='default', description='Phoenix project name') + + +phoenix_settings = PhoenixSettings() diff --git a/kaizen/frontend/cli/cli.py b/kaizen/frontend/cli/cli.py index a38207b7..4925ccdb 100644 --- a/kaizen/frontend/cli/cli.py +++ b/kaizen/frontend/cli/cli.py @@ -18,9 +18,11 @@ app = typer.Typer(help="Kaizen CLI - Manage entities and namespaces") namespaces_app = typer.Typer(help="Namespace management commands") entities_app = typer.Typer(help="Entity management commands") +sync_app = typer.Typer(help="Sync commands") app.add_typer(namespaces_app, name="namespaces") app.add_typer(entities_app, name="entities") +app.add_typer(sync_app, name="sync") console = Console() @@ -319,5 +321,58 @@ def show_entity( raise typer.Exit(1) +# ============================================================================= +# Sync Commands +# ============================================================================= + + +@sync_app.command("phoenix") +def sync_phoenix( + phoenix_url: Annotated[Optional[str], typer.Option("--url", "-u", help="Phoenix server URL")] = None, + namespace: Annotated[Optional[str], typer.Option("--namespace", "-n", help="Target namespace")] = None, + project: Annotated[Optional[str], typer.Option("--project", "-p", help="Phoenix project name")] = None, + limit: Annotated[int, typer.Option(help="Maximum number of spans to fetch")] = 100, + include_errors: Annotated[bool, typer.Option("--include-errors", help="Include failed/error spans")] = False, +): + """Sync trajectories from Arize Phoenix and generate tips.""" + from kaizen.sync.phoenix_sync import PhoenixSync + + syncer = PhoenixSync( + phoenix_url=phoenix_url, + namespace_id=namespace, + project=project, + ) + + console.print(f"[bold]Syncing from Phoenix[/bold]") + console.print(f" URL: {syncer.phoenix_url}") + console.print(f" Project: {syncer.project}") + console.print(f" Namespace: {syncer.namespace_id}") + console.print(f" Limit: {limit}") + console.print() + + try: + result = syncer.sync(limit=limit, include_errors=include_errors) + + table = Table(title="Sync Results") + table.add_column("Metric", style="cyan") + table.add_column("Count", justify="right") + + table.add_row("Trajectories processed", str(result.processed)) + table.add_row("Trajectories skipped (already synced)", str(result.skipped)) + table.add_row("Tips generated", str(result.tips_generated)) + table.add_row("Errors", str(len(result.errors))) + + console.print(table) + + if result.errors: + console.print("\n[red]Errors:[/red]") + for error in result.errors: + console.print(f" - {error}") + + except Exception as e: + console.print(f"[red]Sync failed: {e}[/red]") + raise typer.Exit(1) + + if __name__ == "__main__": app() diff --git a/kaizen/sync/__init__.py b/kaizen/sync/__init__.py new file mode 100644 index 00000000..b2506ff8 --- /dev/null +++ b/kaizen/sync/__init__.py @@ -0,0 +1,3 @@ +from kaizen.sync.phoenix_sync import PhoenixSync + +__all__ = ['PhoenixSync'] diff --git a/kaizen/sync/phoenix_sync.py b/kaizen/sync/phoenix_sync.py new file mode 100644 index 00000000..14f7c187 --- /dev/null +++ b/kaizen/sync/phoenix_sync.py @@ -0,0 +1,416 @@ +""" +Phoenix Sync - Fetch trajectories from Arize Phoenix and generate tips. + +This module provides functionality to: +1. Fetch agent trajectories from Phoenix's REST API +2. Deduplicate already-processed trajectories +3. Generate tips/guidelines from new trajectories +4. Store both trajectories and tips in the Kaizen backend +""" + +import json +import logging +import urllib.request +from dataclasses import dataclass +from typing import Any + +from kaizen.config.phoenix import phoenix_settings +from kaizen.config.kaizen import kaizen_config +from kaizen.frontend.client.kaizen_client import KaizenClient +from kaizen.llm.tips.tips import generate_tips +from kaizen.schema.core import Entity +from kaizen.schema.exceptions import NamespaceNotFoundException + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger("kaizen.sync.phoenix") + + +@dataclass +class SyncResult: + """Result of a sync operation.""" + processed: int + skipped: int + tips_generated: int + errors: list[str] + + +class PhoenixSync: + """Sync trajectories from Arize Phoenix to Kaizen.""" + + def __init__( + self, + phoenix_url: str | None = None, + namespace_id: str | None = None, + project: str | None = None, + ): + self.phoenix_url = phoenix_url or phoenix_settings.url + self.project = project or phoenix_settings.project + self.namespace_id = namespace_id or kaizen_config.namespace_id + self.client = KaizenClient() + + def _ensure_namespace(self): + """Ensure the target namespace exists.""" + try: + self.client.get_namespace_details(self.namespace_id) + except NamespaceNotFoundException: + self.client.create_namespace(self.namespace_id) + logger.info(f"Created namespace: {self.namespace_id}") + + def _fetch_spans(self, limit: int = 1000) -> list[dict]: + """Fetch spans from Phoenix, handling pagination.""" + spans = [] + cursor = None + + while True: + url = f"{self.phoenix_url}/v1/projects/{self.project}/spans?limit={min(limit - len(spans), 100)}" + if cursor: + url += f"&cursor={cursor}" + + try: + with urllib.request.urlopen(url, timeout=30) as response: + data = json.loads(response.read().decode()) + except Exception as e: + logger.error(f"Failed to fetch spans from Phoenix: {e}") + raise + + spans.extend(data.get("data", [])) + cursor = data.get("next_cursor") + + if not cursor or len(spans) >= limit: + break + + return spans + + def _get_processed_span_ids(self) -> set[str]: + """Get span_ids that have already been processed.""" + try: + entities = self.client.search_entities( + namespace_id=self.namespace_id, + filters={"type": "trajectory"}, + limit=10000 + ) + return { + e.metadata.get("span_id") + for e in entities + if e.metadata and e.metadata.get("span_id") + } + except NamespaceNotFoundException: + return set() + + def _parse_content(self, content: Any) -> Any: + """Parse content which may be a string representation of a list/dict.""" + if isinstance(content, str): + try: + return json.loads(content) + except json.JSONDecodeError: + try: + import ast + return ast.literal_eval(content) + except (ValueError, SyntaxError): + return content + return content + + def _extract_messages_from_span(self, span: dict) -> list[dict]: + """Extract messages from a single span's attributes.""" + attrs = span.get("attributes", {}) + messages = [] + + # Extract prompt messages + prompt_indices = set() + for key in attrs: + if key.startswith("gen_ai.prompt.") and key.endswith(".role"): + idx = int(key.split(".")[2]) + prompt_indices.add(idx) + + for i in sorted(prompt_indices): + 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) + }) + + # Extract completion messages + completion_indices = set() + for key in attrs: + if key.startswith("gen_ai.completion.") and key.endswith(".role"): + idx = int(key.split(".")[2]) + completion_indices.add(idx) + + for i in sorted(completion_indices): + 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) + }) + + return messages + + def _convert_to_openai_format(self, content: Any, role: str) -> dict: + """Convert Anthropic message format to OpenAI format.""" + if isinstance(content, str): + return {"role": role, "content": content} + + if not isinstance(content, list): + return {"role": role, "content": str(content)} + + text_parts = [] + tool_calls = [] + tool_results = [] + thinking_parts = [] + + for block in content: + if not isinstance(block, dict): + text_parts.append(str(block)) + continue + + block_type = block.get("type") + + if block_type == "text": + text = block.get("text", "") + if text and text != "(no content)": + text_parts.append(text) + + elif block_type == "thinking": + thinking = block.get("thinking", "") + if thinking: + 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", {})) + } + }) + + 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) + }) + + if role == "assistant": + msg = {"role": "assistant"} + if thinking_parts: + msg["thinking"] = "\n\n".join(thinking_parts) + if text_parts: + msg["content"] = "\n\n".join(text_parts) + elif not tool_calls: + msg["content"] = None + if tool_calls: + msg["tool_calls"] = tool_calls + return msg + + elif role == "user" and tool_results: + return {"role": "tool", "tool_results": tool_results} + + else: + content_text = "\n\n".join(text_parts) if text_parts else "" + return {"role": role, "content": content_text} + + def _extract_trajectory(self, span: dict) -> dict: + """Extract a complete trajectory from a span.""" + attrs = span.get("attributes", {}) + messages = self._extract_messages_from_span(span) + + openai_messages = [] + + for msg in messages: + role = msg["role"] + content = msg["content"] + converted = self._convert_to_openai_format(content, role) + + 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"] + }) + else: + openai_messages.append(converted) + + return { + "trace_id": span["context"]["trace_id"], + "span_id": span["context"]["span_id"], + "model": attrs.get("gen_ai.request.model", "unknown"), + "timestamp": span.get("start_time"), + "messages": openai_messages, + "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") + } + } + + 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", []): + if not msg.get("content") and not msg.get("tool_calls"): + continue + + if msg.get("content"): + content = msg["content"] + if isinstance(content, str): + content = re.sub( + r'.*?', + '', + content, + flags=re.DOTALL + ).strip() + if not content: + continue + msg = {**msg, "content": content} + + cleaned_messages.append(msg) + + return {**trajectory, "messages": cleaned_messages} + + def _process_trajectory(self, trajectory: dict) -> int: + """Process a single trajectory: store it and generate tips. + + Returns the number of tips generated. + """ + # Store trajectory messages + entities = [] + 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"], + } + )) + + if entities: + self.client.update_entities( + namespace_id=self.namespace_id, + entities=entities, + enable_conflict_resolution=False + ) + + # Generate tips from the trajectory + tips = generate_tips(trajectory["messages"]) + + if tips: + tip_entities = [ + Entity( + type='guideline', + content=tip, + metadata={ + "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 + ) + + return len(tips) + + def sync( + self, + limit: int = 100, + include_errors: bool = False, + ) -> SyncResult: + """ + Fetch new trajectories from Phoenix and generate tips. + + Args: + limit: Maximum number of spans to fetch from Phoenix + include_errors: Whether to include failed/error spans + + Returns: + SyncResult with counts of processed, skipped, and tips generated + """ + logger.info(f"Starting sync from {self.phoenix_url} to namespace '{self.namespace_id}'") + + self._ensure_namespace() + + # Fetch spans from Phoenix + spans = self._fetch_spans(limit) + logger.info(f"Fetched {len(spans)} spans from Phoenix") + + # Get already processed span IDs + processed_ids = self._get_processed_span_ids() + logger.info(f"Found {len(processed_ids)} already processed spans") + + processed = 0 + skipped = 0 + tips_generated = 0 + errors = [] + + for span in spans: + # Filter to LLM request spans + if span.get("name") != "litellm_request": + continue + + # Filter errors if requested + if not include_errors and span.get("status_code") == "ERROR": + continue + + # Check if already processed + span_id = span.get("context", {}).get("span_id") + if span_id in processed_ids: + skipped += 1 + continue + + # Only include spans with actual messages + attrs = span.get("attributes", {}) + if not any(k.startswith("gen_ai.prompt.") for k in attrs): + continue + + try: + trajectory = self._extract_trajectory(span) + trajectory = self._clean_trajectory(trajectory) + + if trajectory["messages"]: + tips_count = self._process_trajectory(trajectory) + processed += 1 + tips_generated += tips_count + logger.info( + f"Processed span {span_id[:12]}... - " + f"generated {tips_count} tips" + ) + except Exception as e: + error_msg = f"Error processing span {span_id}: {e}" + logger.error(error_msg) + errors.append(error_msg) + + result = SyncResult( + processed=processed, + skipped=skipped, + tips_generated=tips_generated, + errors=errors + ) + + logger.info( + f"Sync complete: {processed} processed, {skipped} skipped, " + f"{tips_generated} tips generated, {len(errors)} errors" + ) + + return result diff --git a/pyproject.toml b/pyproject.toml index da814086..d08abdc3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,9 +38,10 @@ include = ["kaizen"] package = true [tool.pytest.ini_options] -addopts = "--ignore=explorations" +addopts = "--ignore=explorations -m 'not phoenix'" markers = [ "e2e", "unit", + "phoenix: tests requiring Phoenix sync functionality (deselected by default)", ] anyio_mode = "auto" \ No newline at end of file diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 00000000..eced9fdd --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,27 @@ +"""Pytest configuration and fixtures.""" + +import pytest + + +def pytest_addoption(parser): + """Add custom command line options.""" + parser.addoption( + "--run-phoenix", + action="store_true", + default=False, + help="Run Phoenix sync tests (skipped by default)", + ) + + +def pytest_configure(config): + """Override marker filter when --run-phoenix is passed.""" + if config.getoption("--run-phoenix"): + # Remove the default marker filter to include phoenix tests + # Get current markexpr and modify it + markexpr = config.getoption("markexpr", default="") + if markexpr == "not phoenix": + config.option.markexpr = "" + elif "not phoenix" in markexpr: + # Remove "not phoenix" from the expression + new_expr = markexpr.replace("not phoenix and ", "").replace(" and not phoenix", "").replace("not phoenix", "") + config.option.markexpr = new_expr.strip() diff --git a/tests/unit/test_cli.py b/tests/unit/test_cli.py index b1759461..c2ac250d 100644 --- a/tests/unit/test_cli.py +++ b/tests/unit/test_cli.py @@ -583,3 +583,267 @@ def test_entities_help(self): assert "delete" in result.stdout assert "search" in result.stdout assert "show" in result.stdout + + def test_sync_help(self): + """Test sync subcommand help.""" + result = runner.invoke(app, ["sync", "--help"]) + + assert result.exit_code == 0 + assert "phoenix" in result.stdout + + +# ============================================================================= +# Sync Commands Tests +# ============================================================================= + + +@pytest.mark.unit +@pytest.mark.phoenix +class TestSyncPhoenix: + """Tests for 'kaizen sync phoenix' command.""" + + def test_sync_phoenix_default_params(self): + """Test sync phoenix with default parameters.""" + with patch("kaizen.sync.phoenix_sync.PhoenixSync") as MockSync: + mock_syncer = MagicMock() + mock_syncer.phoenix_url = "http://localhost:6006" + mock_syncer.project = "default" + mock_syncer.namespace_id = "test_ns" + mock_syncer.sync.return_value = MagicMock( + processed=5, + skipped=2, + tips_generated=10, + errors=[] + ) + MockSync.return_value = mock_syncer + + result = runner.invoke(app, ["sync", "phoenix"]) + + assert result.exit_code == 0 + mock_syncer.sync.assert_called_once_with(limit=100, include_errors=False) + + def test_sync_phoenix_with_custom_url(self): + """Test sync phoenix with custom Phoenix URL.""" + with patch("kaizen.sync.phoenix_sync.PhoenixSync") as MockSync: + mock_syncer = MagicMock() + mock_syncer.phoenix_url = "http://custom:8080" + mock_syncer.project = "default" + mock_syncer.namespace_id = "test_ns" + mock_syncer.sync.return_value = MagicMock( + processed=0, + skipped=0, + tips_generated=0, + errors=[] + ) + MockSync.return_value = mock_syncer + + result = runner.invoke(app, ["sync", "phoenix", "--url", "http://custom:8080"]) + + assert result.exit_code == 0 + MockSync.assert_called_once_with( + phoenix_url="http://custom:8080", + namespace_id=None, + project=None + ) + + def test_sync_phoenix_with_custom_namespace(self): + """Test sync phoenix with custom namespace.""" + with patch("kaizen.sync.phoenix_sync.PhoenixSync") as MockSync: + mock_syncer = MagicMock() + mock_syncer.phoenix_url = "http://localhost:6006" + mock_syncer.project = "default" + mock_syncer.namespace_id = "my_namespace" + mock_syncer.sync.return_value = MagicMock( + processed=0, + skipped=0, + tips_generated=0, + errors=[] + ) + MockSync.return_value = mock_syncer + + result = runner.invoke(app, ["sync", "phoenix", "--namespace", "my_namespace"]) + + assert result.exit_code == 0 + MockSync.assert_called_once_with( + phoenix_url=None, + namespace_id="my_namespace", + project=None + ) + + def test_sync_phoenix_with_custom_project(self): + """Test sync phoenix with custom project.""" + with patch("kaizen.sync.phoenix_sync.PhoenixSync") as MockSync: + mock_syncer = MagicMock() + mock_syncer.phoenix_url = "http://localhost:6006" + mock_syncer.project = "my_project" + mock_syncer.namespace_id = "test_ns" + mock_syncer.sync.return_value = MagicMock( + processed=0, + skipped=0, + tips_generated=0, + errors=[] + ) + MockSync.return_value = mock_syncer + + result = runner.invoke(app, ["sync", "phoenix", "--project", "my_project"]) + + assert result.exit_code == 0 + MockSync.assert_called_once_with( + phoenix_url=None, + namespace_id=None, + project="my_project" + ) + + def test_sync_phoenix_with_custom_limit(self): + """Test sync phoenix with custom limit.""" + with patch("kaizen.sync.phoenix_sync.PhoenixSync") as MockSync: + mock_syncer = MagicMock() + mock_syncer.phoenix_url = "http://localhost:6006" + mock_syncer.project = "default" + mock_syncer.namespace_id = "test_ns" + mock_syncer.sync.return_value = MagicMock( + processed=0, + skipped=0, + tips_generated=0, + errors=[] + ) + MockSync.return_value = mock_syncer + + result = runner.invoke(app, ["sync", "phoenix", "--limit", "50"]) + + assert result.exit_code == 0 + mock_syncer.sync.assert_called_once_with(limit=50, include_errors=False) + + def test_sync_phoenix_with_include_errors(self): + """Test sync phoenix with include-errors flag.""" + with patch("kaizen.sync.phoenix_sync.PhoenixSync") as MockSync: + mock_syncer = MagicMock() + mock_syncer.phoenix_url = "http://localhost:6006" + mock_syncer.project = "default" + mock_syncer.namespace_id = "test_ns" + mock_syncer.sync.return_value = MagicMock( + processed=0, + skipped=0, + tips_generated=0, + errors=[] + ) + MockSync.return_value = mock_syncer + + result = runner.invoke(app, ["sync", "phoenix", "--include-errors"]) + + assert result.exit_code == 0 + mock_syncer.sync.assert_called_once_with(limit=100, include_errors=True) + + def test_sync_phoenix_displays_results(self): + """Test sync phoenix displays results in output.""" + with patch("kaizen.sync.phoenix_sync.PhoenixSync") as MockSync: + mock_syncer = MagicMock() + mock_syncer.phoenix_url = "http://localhost:6006" + mock_syncer.project = "default" + mock_syncer.namespace_id = "test_ns" + mock_syncer.sync.return_value = MagicMock( + processed=10, + skipped=5, + tips_generated=20, + errors=[] + ) + MockSync.return_value = mock_syncer + + result = runner.invoke(app, ["sync", "phoenix"]) + + assert result.exit_code == 0 + assert "Sync Results" in result.stdout + assert "10" in result.stdout # processed + assert "5" in result.stdout # skipped + assert "20" in result.stdout # tips_generated + + def test_sync_phoenix_displays_errors(self): + """Test sync phoenix displays errors if any.""" + with patch("kaizen.sync.phoenix_sync.PhoenixSync") as MockSync: + mock_syncer = MagicMock() + mock_syncer.phoenix_url = "http://localhost:6006" + mock_syncer.project = "default" + mock_syncer.namespace_id = "test_ns" + mock_syncer.sync.return_value = MagicMock( + processed=1, + skipped=0, + tips_generated=0, + errors=["Error processing span abc: Connection failed"] + ) + MockSync.return_value = mock_syncer + + result = runner.invoke(app, ["sync", "phoenix"]) + + assert result.exit_code == 0 + assert "Errors:" in result.stdout + assert "Connection failed" in result.stdout + + def test_sync_phoenix_handles_exception(self): + """Test sync phoenix handles exceptions gracefully.""" + with patch("kaizen.sync.phoenix_sync.PhoenixSync") as MockSync: + mock_syncer = MagicMock() + mock_syncer.phoenix_url = "http://localhost:6006" + mock_syncer.project = "default" + mock_syncer.namespace_id = "test_ns" + mock_syncer.sync.side_effect = Exception("Phoenix server unreachable") + MockSync.return_value = mock_syncer + + result = runner.invoke(app, ["sync", "phoenix"]) + + assert result.exit_code == 1 + assert "Sync failed" in result.stdout + assert "Phoenix server unreachable" in result.stdout + + def test_sync_phoenix_displays_parameters(self): + """Test sync phoenix displays sync parameters.""" + with patch("kaizen.sync.phoenix_sync.PhoenixSync") as MockSync: + mock_syncer = MagicMock() + mock_syncer.phoenix_url = "http://test:6006" + mock_syncer.project = "test_project" + mock_syncer.namespace_id = "test_namespace" + mock_syncer.sync.return_value = MagicMock( + processed=0, + skipped=0, + tips_generated=0, + errors=[] + ) + MockSync.return_value = mock_syncer + + result = runner.invoke(app, ["sync", "phoenix"]) + + assert result.exit_code == 0 + assert "http://test:6006" in result.stdout + assert "test_project" in result.stdout + assert "test_namespace" in result.stdout + + def test_sync_phoenix_all_options(self): + """Test sync phoenix with all options combined.""" + with patch("kaizen.sync.phoenix_sync.PhoenixSync") as MockSync: + mock_syncer = MagicMock() + mock_syncer.phoenix_url = "http://custom:9000" + mock_syncer.project = "prod" + mock_syncer.namespace_id = "production" + mock_syncer.sync.return_value = MagicMock( + processed=100, + skipped=50, + tips_generated=200, + errors=[] + ) + MockSync.return_value = mock_syncer + + result = runner.invoke(app, [ + "sync", "phoenix", + "--url", "http://custom:9000", + "--namespace", "production", + "--project", "prod", + "--limit", "500", + "--include-errors" + ]) + + assert result.exit_code == 0 + MockSync.assert_called_once_with( + phoenix_url="http://custom:9000", + namespace_id="production", + project="prod" + ) + mock_syncer.sync.assert_called_once_with(limit=500, include_errors=True) diff --git a/tests/unit/test_extract_trajectories.py b/tests/unit/test_extract_trajectories.py new file mode 100644 index 00000000..b41a9b7c --- /dev/null +++ b/tests/unit/test_extract_trajectories.py @@ -0,0 +1,817 @@ +"""Tests for extract_trajectories.py standalone script functions.""" + +import json +import os +import sys +from unittest.mock import MagicMock, patch, Mock + +import pytest + +# Add project root to path for importing the standalone script +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(__file__)))) + +from extract_trajectories import ( + parse_content, + extract_messages_from_span, + convert_anthropic_to_openai, + extract_trajectory, + filter_system_reminders, + clean_trajectory, + format_trajectory_as_text, + get_trajectories, +) + +# Mark all tests in this module as phoenix tests (skipped by default) +pytestmark = pytest.mark.phoenix + + +# ============================================================================= +# parse_content() Tests +# ============================================================================= + + +@pytest.mark.unit +class TestParseContent: + """Tests for parse_content function.""" + + def test_parse_content_json_dict(self): + """Test parsing JSON dict string.""" + content = '{"key": "value"}' + result = parse_content(content) + assert result == {"key": "value"} + + def test_parse_content_json_list(self): + """Test parsing JSON list string.""" + content = '[{"type": "text", "text": "hello"}]' + result = parse_content(content) + assert result == [{"type": "text", "text": "hello"}] + + def test_parse_content_python_literal_dict(self): + """Test parsing Python literal dict (single quotes).""" + content = "{'key': 'value'}" + result = parse_content(content) + assert result == {"key": "value"} + + def test_parse_content_python_literal_list(self): + """Test parsing Python literal list.""" + content = "[{'type': 'text'}]" + result = parse_content(content) + assert result == [{"type": "text"}] + + def test_parse_content_plain_string(self): + """Test plain string is returned unchanged.""" + content = "Just a plain message" + result = parse_content(content) + assert result == "Just a plain message" + + def test_parse_content_passthrough_dict(self): + """Test dict is passed through.""" + content = {"already": "dict"} + result = parse_content(content) + assert result == {"already": "dict"} + + def test_parse_content_passthrough_list(self): + """Test list is passed through.""" + content = [1, 2, 3] + result = parse_content(content) + assert result == [1, 2, 3] + + def test_parse_content_invalid_syntax(self): + """Test invalid syntax returns original string.""" + content = "not {valid json" + result = parse_content(content) + assert result == content + + +# ============================================================================= +# extract_messages_from_span() Tests +# ============================================================================= + + +@pytest.mark.unit +class TestExtractMessagesFromSpan: + """Tests for extract_messages_from_span function.""" + + def test_extract_single_prompt(self): + """Test extracting a single prompt.""" + span = { + "attributes": { + "gen_ai.prompt.0.role": "user", + "gen_ai.prompt.0.content": "Hello" + } + } + messages = extract_messages_from_span(span) + assert len(messages) == 1 + assert messages[0] == { + "index": 0, + "type": "prompt", + "role": "user", + "content": "Hello" + } + + def test_extract_multiple_prompts_sorted(self): + """Test that multiple prompts are sorted by index.""" + span = { + "attributes": { + "gen_ai.prompt.2.role": "user", + "gen_ai.prompt.2.content": "Third", + "gen_ai.prompt.0.role": "system", + "gen_ai.prompt.0.content": "First", + "gen_ai.prompt.1.role": "user", + "gen_ai.prompt.1.content": "Second" + } + } + messages = extract_messages_from_span(span) + assert len(messages) == 3 + assert messages[0]["content"] == "First" + assert messages[1]["content"] == "Second" + assert messages[2]["content"] == "Third" + + def test_extract_completion_messages(self): + """Test extracting completion messages.""" + span = { + "attributes": { + "gen_ai.completion.0.role": "assistant", + "gen_ai.completion.0.content": "I can help with that." + } + } + messages = extract_messages_from_span(span) + assert len(messages) == 1 + assert messages[0]["type"] == "completion" + assert messages[0]["role"] == "assistant" + + def test_extract_mixed_prompts_and_completions(self): + """Test extracting both prompts and completions.""" + span = { + "attributes": { + "gen_ai.prompt.0.role": "user", + "gen_ai.prompt.0.content": "Question", + "gen_ai.completion.0.role": "assistant", + "gen_ai.completion.0.content": "Answer" + } + } + messages = extract_messages_from_span(span) + prompts = [m for m in messages if m["type"] == "prompt"] + completions = [m for m in messages if m["type"] == "completion"] + assert len(prompts) == 1 + assert len(completions) == 1 + + def test_extract_empty_attributes(self): + """Test extracting from span with no relevant attributes.""" + span = {"attributes": {"other.attr": "value"}} + messages = extract_messages_from_span(span) + assert messages == [] + + def test_extract_parses_json_content(self): + """Test that JSON content is parsed.""" + span = { + "attributes": { + "gen_ai.prompt.0.role": "assistant", + "gen_ai.prompt.0.content": '[{"type": "text", "text": "Hi"}]' + } + } + messages = extract_messages_from_span(span) + assert messages[0]["content"] == [{"type": "text", "text": "Hi"}] + + def test_extract_skips_missing_role(self): + """Test that entries without role are skipped.""" + span = { + "attributes": { + "gen_ai.prompt.0.content": "No role here" + } + } + messages = extract_messages_from_span(span) + assert messages == [] + + +# ============================================================================= +# convert_anthropic_to_openai() Tests +# ============================================================================= + + +@pytest.mark.unit +class TestConvertAnthropicToOpenai: + """Tests for convert_anthropic_to_openai function.""" + + def test_convert_string_content(self): + """Test converting simple string content.""" + result = convert_anthropic_to_openai("Hello", "user") + assert result == {"role": "user", "content": "Hello"} + + def test_convert_text_block(self): + """Test converting text block.""" + content = [{"type": "text", "text": "Hello world"}] + result = convert_anthropic_to_openai(content, "assistant") + assert result["role"] == "assistant" + assert result["content"] == "Hello world" + + def test_convert_thinking_block(self): + """Test converting thinking block.""" + content = [ + {"type": "thinking", "thinking": "Let me think about this..."}, + {"type": "text", "text": "Here's my answer."} + ] + result = convert_anthropic_to_openai(content, "assistant") + assert result["thinking"] == "Let me think about this..." + assert result["content"] == "Here's my answer." + + def test_convert_multiple_thinking_blocks(self): + """Test converting multiple thinking blocks.""" + content = [ + {"type": "thinking", "thinking": "First thought"}, + {"type": "thinking", "thinking": "Second thought"}, + {"type": "text", "text": "Final answer"} + ] + result = convert_anthropic_to_openai(content, "assistant") + assert "First thought" in result["thinking"] + assert "Second thought" in result["thinking"] + + def test_convert_tool_use_block(self): + """Test converting tool_use block.""" + content = [ + { + "type": "tool_use", + "id": "tool_abc", + "name": "search", + "input": {"query": "test"} + } + ] + result = convert_anthropic_to_openai(content, "assistant") + assert result["role"] == "assistant" + assert len(result["tool_calls"]) == 1 + tool_call = result["tool_calls"][0] + assert tool_call["id"] == "tool_abc" + assert tool_call["type"] == "function" + assert tool_call["function"]["name"] == "search" + assert json.loads(tool_call["function"]["arguments"]) == {"query": "test"} + + def test_convert_multiple_tool_use_blocks(self): + """Test converting multiple tool_use blocks.""" + content = [ + {"type": "tool_use", "id": "t1", "name": "read", "input": {}}, + {"type": "tool_use", "id": "t2", "name": "write", "input": {}} + ] + result = convert_anthropic_to_openai(content, "assistant") + assert len(result["tool_calls"]) == 2 + + def test_convert_tool_result_block(self): + """Test converting tool_result block.""" + content = [ + { + "type": "tool_result", + "tool_use_id": "tool_123", + "content": "Result data", + "is_error": False + } + ] + result = convert_anthropic_to_openai(content, "user") + assert result["role"] == "tool" + assert "tool_results" in result + assert result["tool_results"][0]["tool_call_id"] == "tool_123" + assert result["tool_results"][0]["content"] == "Result data" + + def test_convert_tool_result_with_error(self): + """Test converting tool_result with error flag.""" + content = [ + { + "type": "tool_result", + "tool_use_id": "tool_err", + "content": "Error occurred", + "is_error": True + } + ] + result = convert_anthropic_to_openai(content, "user") + assert result["tool_results"][0]["is_error"] is True + + def test_convert_filters_no_content_text(self): + """Test that (no content) placeholder is filtered.""" + content = [ + {"type": "text", "text": "(no content)"}, + {"type": "text", "text": "Real text"} + ] + result = convert_anthropic_to_openai(content, "assistant") + assert result["content"] == "Real text" + + def test_convert_filters_empty_text(self): + """Test that empty text is filtered.""" + content = [ + {"type": "text", "text": ""}, + {"type": "text", "text": "Content"} + ] + result = convert_anthropic_to_openai(content, "assistant") + assert result["content"] == "Content" + + def test_convert_assistant_tool_only_no_content(self): + """Test assistant with only tool calls has null content.""" + content = [ + {"type": "tool_use", "id": "t1", "name": "test", "input": {}} + ] + result = convert_anthropic_to_openai(content, "assistant") + assert result.get("content") is None + assert "tool_calls" in result + + def test_convert_non_dict_in_list(self): + """Test handling non-dict items in content list.""" + content = ["string item", {"type": "text", "text": "dict item"}] + result = convert_anthropic_to_openai(content, "user") + assert "string item" in result["content"] + assert "dict item" in result["content"] + + def test_convert_non_list_non_string(self): + """Test converting non-list, non-string content.""" + result = convert_anthropic_to_openai(42, "user") + assert result == {"role": "user", "content": "42"} + + def test_convert_user_regular_message(self): + """Test converting regular user message (not tool result).""" + content = [{"type": "text", "text": "User question"}] + result = convert_anthropic_to_openai(content, "user") + assert result["role"] == "user" + assert result["content"] == "User question" + + +# ============================================================================= +# filter_system_reminders() Tests +# ============================================================================= + + +@pytest.mark.unit +class TestFilterSystemReminders: + """Tests for filter_system_reminders function.""" + + def test_filter_single_reminder(self): + """Test filtering a single system reminder.""" + text = "Before reminder after" + result = filter_system_reminders(text) + assert result == "Before after" + + def test_filter_multiple_reminders(self): + """Test filtering multiple system reminders.""" + text = "firstmiddlesecond" + result = filter_system_reminders(text) + assert result == "middle" + + def test_filter_multiline_reminder(self): + """Test filtering multiline system reminder.""" + text = "Start\n\nLine 1\nLine 2\n\nEnd" + result = filter_system_reminders(text) + assert "" not in result + assert "Start" in result + assert "End" in result + + def test_filter_no_reminders(self): + """Test text without reminders is unchanged.""" + text = "No reminders here" + result = filter_system_reminders(text) + assert result == "No reminders here" + + def test_filter_empty_string(self): + """Test filtering empty string.""" + result = filter_system_reminders("") + assert result == "" + + def test_filter_only_reminder(self): + """Test text that is only a reminder.""" + text = "Only reminder content" + result = filter_system_reminders(text) + assert result == "" + + +# ============================================================================= +# clean_trajectory() Tests +# ============================================================================= + + +@pytest.mark.unit +class TestCleanTrajectory: + """Tests for clean_trajectory function.""" + + def test_clean_removes_system_reminders(self): + """Test that system reminders are removed from content.""" + trajectory = { + "trace_id": "test", + "messages": [ + {"role": "user", "content": "Hi ignore there"} + ] + } + cleaned = clean_trajectory(trajectory) + assert "" not in cleaned["messages"][0]["content"] + assert "Hi" in cleaned["messages"][0]["content"] + + def test_clean_removes_empty_messages(self): + """Test that empty messages are removed.""" + trajectory = { + "trace_id": "test", + "messages": [ + {"role": "user", "content": "Valid"}, + {"role": "assistant", "content": ""}, + {"role": "assistant", "content": None}, + {"role": "user", "content": "Also valid"} + ] + } + cleaned = clean_trajectory(trajectory) + assert len(cleaned["messages"]) == 2 + + def test_clean_preserves_tool_calls(self): + """Test that messages with tool_calls are preserved.""" + trajectory = { + "trace_id": "test", + "messages": [ + {"role": "assistant", "tool_calls": [{"id": "1"}]} + ] + } + cleaned = clean_trajectory(trajectory) + assert len(cleaned["messages"]) == 1 + + def test_clean_with_remove_reminders_false(self): + """Test that reminders are preserved when flag is False.""" + trajectory = { + "trace_id": "test", + "messages": [ + {"role": "user", "content": "keep me"} + ] + } + cleaned = clean_trajectory(trajectory, remove_system_reminders=False) + assert "" in cleaned["messages"][0]["content"] + + def test_clean_removes_only_reminder_messages(self): + """Test that messages containing only reminders are removed.""" + trajectory = { + "trace_id": "test", + "messages": [ + {"role": "user", "content": "Keep"}, + {"role": "assistant", "content": "only reminder"} + ] + } + cleaned = clean_trajectory(trajectory) + assert len(cleaned["messages"]) == 1 + assert cleaned["messages"][0]["content"] == "Keep" + + def test_clean_preserves_metadata(self): + """Test that trajectory metadata is preserved.""" + trajectory = { + "trace_id": "t1", + "span_id": "s1", + "model": "claude-3", + "messages": [{"role": "user", "content": "Hi"}] + } + cleaned = clean_trajectory(trajectory) + assert cleaned["trace_id"] == "t1" + assert cleaned["span_id"] == "s1" + assert cleaned["model"] == "claude-3" + + +# ============================================================================= +# format_trajectory_as_text() Tests +# ============================================================================= + + +@pytest.mark.unit +class TestFormatTrajectoryAsText: + """Tests for format_trajectory_as_text function.""" + + def test_format_user_message(self): + """Test formatting user message.""" + trajectory = { + "trace_id": "abc123def456", + "model": "claude-3", + "timestamp": "2024-01-15T10:00:00Z", + "messages": [ + {"role": "user", "content": "What is 2+2?"} + ] + } + result = format_trajectory_as_text(trajectory) + assert "[USER]" in result + assert "What is 2+2?" in result + assert "abc123def456"[:12] in result + + def test_format_assistant_message(self): + """Test formatting assistant message.""" + trajectory = { + "trace_id": "abc", + "model": "claude-3", + "timestamp": "2024-01-15", + "messages": [ + {"role": "assistant", "content": "The answer is 4."} + ] + } + result = format_trajectory_as_text(trajectory) + assert "[ASSISTANT]" in result + assert "The answer is 4." in result + + def test_format_with_thinking(self): + """Test formatting assistant message with thinking.""" + trajectory = { + "trace_id": "abc", + "model": "claude-3", + "timestamp": "2024-01-15", + "messages": [ + { + "role": "assistant", + "thinking": "Let me calculate...", + "content": "The answer is 4." + } + ] + } + result = format_trajectory_as_text(trajectory, include_thinking=True) + assert "" in result + assert "Let me calculate..." in result + assert "" in result + + def test_format_without_thinking(self): + """Test formatting with thinking disabled.""" + trajectory = { + "trace_id": "abc", + "model": "claude-3", + "timestamp": "2024-01-15", + "messages": [ + { + "role": "assistant", + "thinking": "Hidden thought", + "content": "Visible content" + } + ] + } + result = format_trajectory_as_text(trajectory, include_thinking=False) + assert "Hidden thought" not in result + assert "Visible content" in result + + def test_format_tool_calls(self): + """Test formatting tool calls.""" + trajectory = { + "trace_id": "abc", + "model": "claude-3", + "timestamp": "2024-01-15", + "messages": [ + { + "role": "assistant", + "tool_calls": [ + { + "id": "tool_12345678901234567890", + "function": { + "name": "read_file", + "arguments": '{"path": "/test.txt"}' + } + } + ] + } + ] + } + result = format_trajectory_as_text(trajectory) + assert "Tool call: read_file" in result + assert "Arguments:" in result + + def test_format_tool_results(self): + """Test formatting tool results.""" + trajectory = { + "trace_id": "abc", + "model": "claude-3", + "timestamp": "2024-01-15", + "messages": [ + { + "role": "assistant", + "tool_calls": [ + {"id": "tool_abc", "function": {"name": "read", "arguments": "{}"}} + ] + }, + { + "role": "tool", + "tool_call_id": "tool_abc", + "content": "File contents here" + } + ] + } + result = format_trajectory_as_text(trajectory) + assert "[TOOL RESULT]" in result + assert "File contents here" in result + + def test_format_long_thinking_truncated(self): + """Test that long thinking is truncated.""" + long_thinking = "A" * 600 + trajectory = { + "trace_id": "abc", + "model": "claude-3", + "timestamp": "2024-01-15", + "messages": [ + {"role": "assistant", "thinking": long_thinking, "content": "Answer"} + ] + } + result = format_trajectory_as_text(trajectory) + assert "..." in result + assert len(result) < len(long_thinking) + 500 + + def test_format_header_info(self): + """Test that header info is included.""" + trajectory = { + "trace_id": "trace_123456789", + "model": "claude-3-opus", + "timestamp": "2024-01-15T10:30:00Z", + "messages": [] + } + result = format_trajectory_as_text(trajectory) + assert "Trajectory:" in result + assert "Model: claude-3-opus" in result + assert "Timestamp: 2024-01-15T10:30:00Z" in result + + +# ============================================================================= +# extract_trajectory() Tests +# ============================================================================= + + +@pytest.mark.unit +class TestExtractTrajectory: + """Tests for extract_trajectory function.""" + + def test_extract_basic_trajectory(self): + """Test extracting a basic trajectory.""" + span = { + "context": {"trace_id": "trace_1", "span_id": "span_1"}, + "start_time": "2024-01-15T10:00:00Z", + "attributes": { + "gen_ai.request.model": "claude-3", + "gen_ai.prompt.0.role": "user", + "gen_ai.prompt.0.content": "Hello", + "gen_ai.completion.0.role": "assistant", + "gen_ai.completion.0.content": "Hi there!" + } + } + trajectory = extract_trajectory(span) + + assert trajectory["trace_id"] == "trace_1" + assert trajectory["span_id"] == "span_1" + assert trajectory["model"] == "claude-3" + assert trajectory["timestamp"] == "2024-01-15T10:00:00Z" + assert len(trajectory["messages"]) >= 2 + + def test_extract_trajectory_with_usage(self): + """Test extracting trajectory with usage info.""" + span = { + "context": {"trace_id": "t1", "span_id": "s1"}, + "start_time": "2024-01-15", + "attributes": { + "gen_ai.request.model": "claude-3", + "gen_ai.prompt.0.role": "user", + "gen_ai.prompt.0.content": "Test", + "gen_ai.usage.prompt_tokens": 100, + "gen_ai.usage.completion_tokens": 50, + "llm.usage.total_tokens": 150 + } + } + trajectory = extract_trajectory(span) + + assert trajectory["usage"]["prompt_tokens"] == 100 + assert trajectory["usage"]["completion_tokens"] == 50 + assert trajectory["usage"]["total_tokens"] == 150 + + def test_extract_trajectory_unknown_model(self): + """Test extracting trajectory without model info.""" + span = { + "context": {"trace_id": "t1", "span_id": "s1"}, + "start_time": "2024-01-15", + "attributes": { + "gen_ai.prompt.0.role": "user", + "gen_ai.prompt.0.content": "Test" + } + } + trajectory = extract_trajectory(span) + + assert trajectory["model"] == "unknown" + + def test_extract_trajectory_expands_tool_results(self): + """Test that tool results are expanded into individual messages.""" + tool_result_content = json.dumps([ + {"type": "tool_result", "tool_use_id": "t1", "content": "Result 1"}, + {"type": "tool_result", "tool_use_id": "t2", "content": "Result 2"} + ]) + span = { + "context": {"trace_id": "t1", "span_id": "s1"}, + "start_time": "2024-01-15", + "attributes": { + "gen_ai.request.model": "claude-3", + "gen_ai.prompt.0.role": "user", + "gen_ai.prompt.0.content": tool_result_content + } + } + trajectory = extract_trajectory(span) + + tool_messages = [m for m in trajectory["messages"] if m.get("role") == "tool"] + assert len(tool_messages) == 2 + + +# ============================================================================= +# get_trajectories() Tests (with mocked network) +# ============================================================================= + + +@pytest.mark.unit +class TestGetTrajectories: + """Tests for get_trajectories function.""" + + @patch("extract_trajectories.fetch_spans") + def test_get_trajectories_filters_non_llm_spans(self, mock_fetch): + """Test that non-LLM spans are filtered.""" + mock_fetch.return_value = [ + {"name": "other_span", "attributes": {}} + ] + + result = get_trajectories() + + assert result == [] + + @patch("extract_trajectories.fetch_spans") + def test_get_trajectories_filters_error_spans(self, mock_fetch): + """Test that error spans are filtered by default.""" + mock_fetch.return_value = [ + { + "name": "litellm_request", + "status_code": "ERROR", + "context": {"trace_id": "t1", "span_id": "s1"}, + "attributes": { + "gen_ai.prompt.0.role": "user", + "gen_ai.prompt.0.content": "test" + } + } + ] + + result = get_trajectories(include_errors=False) + + assert result == [] + + @patch("extract_trajectories.fetch_spans") + def test_get_trajectories_includes_errors_when_requested(self, mock_fetch): + """Test that error spans are included when flag is set.""" + mock_fetch.return_value = [ + { + "name": "litellm_request", + "status_code": "ERROR", + "context": {"trace_id": "t1", "span_id": "s1"}, + "start_time": "2024-01-15", + "attributes": { + "gen_ai.request.model": "test", + "gen_ai.prompt.0.role": "user", + "gen_ai.prompt.0.content": "test message" + } + } + ] + + result = get_trajectories(include_errors=True) + + assert len(result) == 1 + + @patch("extract_trajectories.fetch_spans") + def test_get_trajectories_filters_empty_messages(self, mock_fetch): + """Test that spans without messages are filtered.""" + mock_fetch.return_value = [ + { + "name": "litellm_request", + "context": {"trace_id": "t1", "span_id": "s1"}, + "attributes": {} # No gen_ai.prompt.* attributes + } + ] + + result = get_trajectories() + + assert result == [] + + @patch("extract_trajectories.fetch_spans") + def test_get_trajectories_cleans_by_default(self, mock_fetch): + """Test that trajectories are cleaned by default.""" + mock_fetch.return_value = [ + { + "name": "litellm_request", + "context": {"trace_id": "t1", "span_id": "s1"}, + "start_time": "2024-01-15", + "attributes": { + "gen_ai.request.model": "test", + "gen_ai.prompt.0.role": "user", + "gen_ai.prompt.0.content": "Hi remove" + } + } + ] + + result = get_trajectories(clean=True) + + assert len(result) == 1 + assert "" not in result[0]["messages"][0]["content"] + + @patch("extract_trajectories.fetch_spans") + def test_get_trajectories_no_clean(self, mock_fetch): + """Test trajectories without cleaning.""" + mock_fetch.return_value = [ + { + "name": "litellm_request", + "context": {"trace_id": "t1", "span_id": "s1"}, + "start_time": "2024-01-15", + "attributes": { + "gen_ai.request.model": "test", + "gen_ai.prompt.0.role": "user", + "gen_ai.prompt.0.content": "keep" + } + } + ] + + result = get_trajectories(clean=False) + + assert len(result) == 1 + assert "" in result[0]["messages"][0]["content"] diff --git a/tests/unit/test_phoenix_sync.py b/tests/unit/test_phoenix_sync.py new file mode 100644 index 00000000..96c0a007 --- /dev/null +++ b/tests/unit/test_phoenix_sync.py @@ -0,0 +1,797 @@ +"""Tests for Phoenix Sync functionality.""" + +import json +from unittest.mock import MagicMock, patch, Mock + +import pytest + +from kaizen.sync.phoenix_sync import PhoenixSync, SyncResult + +# Mark all tests in this module as phoenix tests (skipped by default) +pytestmark = pytest.mark.phoenix + + +@pytest.fixture +def phoenix_sync(): + """Create a PhoenixSync instance with mocked client.""" + with patch("kaizen.sync.phoenix_sync.KaizenClient") as mock_client_class: + mock_client = MagicMock() + mock_client_class.return_value = mock_client + sync = PhoenixSync( + phoenix_url="http://test-phoenix:6006", + namespace_id="test_namespace", + project="test_project" + ) + sync.client = mock_client + yield sync + + +# ============================================================================= +# _parse_content() Tests +# ============================================================================= + + +@pytest.mark.unit +class TestParseContent: + """Tests for _parse_content method.""" + + def test_parse_content_json_string(self, phoenix_sync): + """Test parsing a JSON string.""" + content = '{"key": "value", "number": 42}' + result = phoenix_sync._parse_content(content) + assert result == {"key": "value", "number": 42} + + def test_parse_content_json_list(self, phoenix_sync): + """Test parsing a JSON list string.""" + content = '[{"type": "text", "text": "hello"}]' + result = phoenix_sync._parse_content(content) + assert result == [{"type": "text", "text": "hello"}] + + def test_parse_content_python_literal(self, phoenix_sync): + """Test parsing a Python literal string.""" + content = "{'key': 'value'}" # Single quotes - not valid JSON + result = phoenix_sync._parse_content(content) + assert result == {"key": "value"} + + def test_parse_content_plain_string(self, phoenix_sync): + """Test that plain strings are returned as-is.""" + content = "This is just plain text" + result = phoenix_sync._parse_content(content) + assert result == "This is just plain text" + + def test_parse_content_passthrough_dict(self, phoenix_sync): + """Test that dicts are passed through unchanged.""" + content = {"already": "parsed"} + result = phoenix_sync._parse_content(content) + assert result == {"already": "parsed"} + + def test_parse_content_passthrough_list(self, phoenix_sync): + """Test that lists are passed through unchanged.""" + content = [{"type": "text"}] + result = phoenix_sync._parse_content(content) + assert result == [{"type": "text"}] + + def test_parse_content_invalid_json_returns_string(self, phoenix_sync): + """Test that invalid JSON/Python returns the original string.""" + content = "not valid {json or python" + result = phoenix_sync._parse_content(content) + assert result == content + + +# ============================================================================= +# _extract_messages_from_span() Tests +# ============================================================================= + + +@pytest.mark.unit +class TestExtractMessagesFromSpan: + """Tests for _extract_messages_from_span method.""" + + def test_extract_single_prompt(self, phoenix_sync): + """Test extracting a single prompt message.""" + span = { + "attributes": { + "gen_ai.prompt.0.role": "user", + "gen_ai.prompt.0.content": "Hello, world!" + } + } + messages = phoenix_sync._extract_messages_from_span(span) + assert len(messages) == 1 + assert messages[0]["role"] == "user" + assert messages[0]["content"] == "Hello, world!" + assert messages[0]["type"] == "prompt" + assert messages[0]["index"] == 0 + + def test_extract_multiple_prompts(self, phoenix_sync): + """Test extracting multiple prompt messages.""" + span = { + "attributes": { + "gen_ai.prompt.0.role": "system", + "gen_ai.prompt.0.content": "You are a helpful assistant.", + "gen_ai.prompt.1.role": "user", + "gen_ai.prompt.1.content": "What is 2+2?" + } + } + messages = phoenix_sync._extract_messages_from_span(span) + assert len(messages) == 2 + assert messages[0]["role"] == "system" + assert messages[1]["role"] == "user" + + def test_extract_with_completion(self, phoenix_sync): + """Test extracting prompts and completions.""" + span = { + "attributes": { + "gen_ai.prompt.0.role": "user", + "gen_ai.prompt.0.content": "Hi", + "gen_ai.completion.0.role": "assistant", + "gen_ai.completion.0.content": "Hello! How can I help?" + } + } + messages = phoenix_sync._extract_messages_from_span(span) + assert len(messages) == 2 + prompts = [m for m in messages if m["type"] == "prompt"] + completions = [m for m in messages if m["type"] == "completion"] + assert len(prompts) == 1 + assert len(completions) == 1 + + def test_extract_empty_span(self, phoenix_sync): + """Test extracting from span with no messages.""" + span = {"attributes": {}} + messages = phoenix_sync._extract_messages_from_span(span) + assert messages == [] + + def test_extract_parses_json_content(self, phoenix_sync): + """Test that JSON content in attributes is parsed.""" + span = { + "attributes": { + "gen_ai.prompt.0.role": "assistant", + "gen_ai.prompt.0.content": '[{"type": "text", "text": "Hello"}]' + } + } + messages = phoenix_sync._extract_messages_from_span(span) + assert len(messages) == 1 + assert messages[0]["content"] == [{"type": "text", "text": "Hello"}] + + def test_extract_handles_non_sequential_indices(self, phoenix_sync): + """Test handling non-sequential message indices.""" + span = { + "attributes": { + "gen_ai.prompt.0.role": "user", + "gen_ai.prompt.0.content": "First", + "gen_ai.prompt.5.role": "user", + "gen_ai.prompt.5.content": "Second" + } + } + messages = phoenix_sync._extract_messages_from_span(span) + assert len(messages) == 2 + # Should be sorted by index + assert messages[0]["index"] == 0 + assert messages[1]["index"] == 5 + + +# ============================================================================= +# _convert_to_openai_format() Tests +# ============================================================================= + + +@pytest.mark.unit +class TestConvertToOpenAIFormat: + """Tests for _convert_to_openai_format method.""" + + def test_convert_simple_string(self, phoenix_sync): + """Test converting a simple string message.""" + result = phoenix_sync._convert_to_openai_format("Hello", "user") + assert result == {"role": "user", "content": "Hello"} + + def test_convert_text_block(self, phoenix_sync): + """Test converting Anthropic text block.""" + content = [{"type": "text", "text": "Hello, world!"}] + result = phoenix_sync._convert_to_openai_format(content, "assistant") + assert result["role"] == "assistant" + assert result["content"] == "Hello, world!" + + def test_convert_multiple_text_blocks(self, phoenix_sync): + """Test converting multiple text blocks.""" + content = [ + {"type": "text", "text": "First part"}, + {"type": "text", "text": "Second part"} + ] + result = phoenix_sync._convert_to_openai_format(content, "assistant") + assert result["content"] == "First part\n\nSecond part" + + def test_convert_thinking_block(self, phoenix_sync): + """Test converting Anthropic thinking block.""" + content = [ + {"type": "thinking", "thinking": "Let me analyze this..."}, + {"type": "text", "text": "The answer is 42."} + ] + result = phoenix_sync._convert_to_openai_format(content, "assistant") + assert result["role"] == "assistant" + assert result["thinking"] == "Let me analyze this..." + assert result["content"] == "The answer is 42." + + def test_convert_tool_use_block(self, phoenix_sync): + """Test converting Anthropic tool_use block.""" + content = [ + { + "type": "tool_use", + "id": "tool_123", + "name": "read_file", + "input": {"path": "/tmp/test.txt"} + } + ] + result = phoenix_sync._convert_to_openai_format(content, "assistant") + assert result["role"] == "assistant" + assert "tool_calls" in result + assert len(result["tool_calls"]) == 1 + assert result["tool_calls"][0]["id"] == "tool_123" + assert result["tool_calls"][0]["type"] == "function" + assert result["tool_calls"][0]["function"]["name"] == "read_file" + assert json.loads(result["tool_calls"][0]["function"]["arguments"]) == {"path": "/tmp/test.txt"} + + def test_convert_tool_result_block(self, phoenix_sync): + """Test converting Anthropic tool_result block.""" + content = [ + { + "type": "tool_result", + "tool_use_id": "tool_123", + "content": "File contents here", + "is_error": False + } + ] + result = phoenix_sync._convert_to_openai_format(content, "user") + assert result["role"] == "tool" + assert "tool_results" in result + assert result["tool_results"][0]["tool_call_id"] == "tool_123" + assert result["tool_results"][0]["content"] == "File contents here" + + def test_convert_mixed_content_blocks(self, phoenix_sync): + """Test converting mixed content blocks.""" + content = [ + {"type": "thinking", "thinking": "I need to read the file first"}, + {"type": "text", "text": "Let me check that file."}, + { + "type": "tool_use", + "id": "tool_456", + "name": "read_file", + "input": {"path": "/etc/hosts"} + } + ] + result = phoenix_sync._convert_to_openai_format(content, "assistant") + assert result["role"] == "assistant" + assert result["thinking"] == "I need to read the file first" + assert result["content"] == "Let me check that file." + assert len(result["tool_calls"]) == 1 + + def test_convert_filters_no_content_text(self, phoenix_sync): + """Test that '(no content)' text is filtered out.""" + content = [ + {"type": "text", "text": "(no content)"}, + {"type": "text", "text": "Real content"} + ] + result = phoenix_sync._convert_to_openai_format(content, "assistant") + assert result["content"] == "Real content" + + def test_convert_assistant_only_tool_calls(self, phoenix_sync): + """Test assistant message with only tool calls (no text).""" + content = [ + { + "type": "tool_use", + "id": "tool_789", + "name": "bash", + "input": {"command": "ls"} + } + ] + result = phoenix_sync._convert_to_openai_format(content, "assistant") + assert result["role"] == "assistant" + assert result.get("content") is None + assert len(result["tool_calls"]) == 1 + + def test_convert_non_dict_in_list(self, phoenix_sync): + """Test handling non-dict items in content list.""" + content = ["plain string", {"type": "text", "text": "dict item"}] + result = phoenix_sync._convert_to_openai_format(content, "user") + assert "plain string" in result["content"] + + def test_convert_non_list_non_string(self, phoenix_sync): + """Test handling content that is neither list nor string.""" + result = phoenix_sync._convert_to_openai_format(12345, "user") + assert result == {"role": "user", "content": "12345"} + + +# ============================================================================= +# _extract_trajectory() Tests +# ============================================================================= + + +@pytest.mark.unit +class TestExtractTrajectory: + """Tests for _extract_trajectory method.""" + + def test_extract_full_trajectory(self, phoenix_sync): + """Test extracting a complete trajectory.""" + span = { + "context": { + "trace_id": "trace_abc123", + "span_id": "span_xyz789" + }, + "start_time": "2024-01-15T10:30:00Z", + "attributes": { + "gen_ai.request.model": "claude-3-opus", + "gen_ai.prompt.0.role": "user", + "gen_ai.prompt.0.content": "What is 2+2?", + "gen_ai.completion.0.role": "assistant", + "gen_ai.completion.0.content": "2+2 equals 4.", + "gen_ai.usage.prompt_tokens": 10, + "gen_ai.usage.completion_tokens": 8, + "llm.usage.total_tokens": 18 + } + } + trajectory = phoenix_sync._extract_trajectory(span) + + assert trajectory["trace_id"] == "trace_abc123" + assert trajectory["span_id"] == "span_xyz789" + assert trajectory["model"] == "claude-3-opus" + assert trajectory["timestamp"] == "2024-01-15T10:30:00Z" + assert len(trajectory["messages"]) == 2 + assert trajectory["usage"]["prompt_tokens"] == 10 + assert trajectory["usage"]["completion_tokens"] == 8 + assert trajectory["usage"]["total_tokens"] == 18 + + def test_extract_trajectory_with_tool_calls(self, phoenix_sync): + """Test extracting trajectory with tool calls.""" + tool_use_content = json.dumps([ + {"type": "text", "text": "I'll read that file."}, + {"type": "tool_use", "id": "tool_1", "name": "read_file", "input": {"path": "/test"}} + ]) + tool_result_content = json.dumps([ + {"type": "tool_result", "tool_use_id": "tool_1", "content": "file contents"} + ]) + + span = { + "context": {"trace_id": "trace_1", "span_id": "span_1"}, + "start_time": "2024-01-15T10:30:00Z", + "attributes": { + "gen_ai.request.model": "claude-3", + "gen_ai.prompt.0.role": "user", + "gen_ai.prompt.0.content": "Read /test", + "gen_ai.prompt.1.role": "assistant", + "gen_ai.prompt.1.content": tool_use_content, + "gen_ai.prompt.2.role": "user", + "gen_ai.prompt.2.content": tool_result_content, + "gen_ai.completion.0.role": "assistant", + "gen_ai.completion.0.content": "The file contains: file contents" + } + } + trajectory = phoenix_sync._extract_trajectory(span) + + # Should have: user, assistant with tool_call, tool result, assistant response + messages = trajectory["messages"] + assert any(m.get("tool_calls") for m in messages) + assert any(m.get("role") == "tool" for m in messages) + + +# ============================================================================= +# _clean_trajectory() Tests +# ============================================================================= + + +@pytest.mark.unit +class TestCleanTrajectory: + """Tests for _clean_trajectory method.""" + + def test_clean_removes_system_reminders(self, phoenix_sync): + """Test that system reminders are removed.""" + trajectory = { + "trace_id": "test", + "messages": [ + { + "role": "user", + "content": "Hello This is a reminder there" + } + ] + } + cleaned = phoenix_sync._clean_trajectory(trajectory) + assert "" not in cleaned["messages"][0]["content"] + assert "Hello" in cleaned["messages"][0]["content"] + assert "there" in cleaned["messages"][0]["content"] + + def test_clean_removes_multiline_system_reminders(self, phoenix_sync): + """Test that multiline system reminders are removed.""" + trajectory = { + "trace_id": "test", + "messages": [ + { + "role": "assistant", + "content": "Start\n\nLine 1\nLine 2\n\nEnd" + } + ] + } + cleaned = phoenix_sync._clean_trajectory(trajectory) + assert "" not in cleaned["messages"][0]["content"] + assert "Start" in cleaned["messages"][0]["content"] + assert "End" in cleaned["messages"][0]["content"] + + def test_clean_removes_empty_messages(self, phoenix_sync): + """Test that empty messages are removed.""" + trajectory = { + "trace_id": "test", + "messages": [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": ""}, + {"role": "assistant", "content": None}, + {"role": "user", "content": "World"} + ] + } + cleaned = phoenix_sync._clean_trajectory(trajectory) + assert len(cleaned["messages"]) == 2 + assert cleaned["messages"][0]["content"] == "Hello" + assert cleaned["messages"][1]["content"] == "World" + + def test_clean_preserves_tool_calls(self, phoenix_sync): + """Test that messages with tool_calls but no content are preserved.""" + trajectory = { + "trace_id": "test", + "messages": [ + { + "role": "assistant", + "tool_calls": [{"id": "1", "function": {"name": "test"}}] + } + ] + } + cleaned = phoenix_sync._clean_trajectory(trajectory) + assert len(cleaned["messages"]) == 1 + assert "tool_calls" in cleaned["messages"][0] + + def test_clean_removes_only_reminder_content(self, phoenix_sync): + """Test that messages with only system reminders are removed.""" + trajectory = { + "trace_id": "test", + "messages": [ + {"role": "user", "content": "Valid"}, + {"role": "assistant", "content": "Only reminder"}, + {"role": "user", "content": "Also valid"} + ] + } + cleaned = phoenix_sync._clean_trajectory(trajectory) + assert len(cleaned["messages"]) == 2 + + def test_clean_preserves_non_string_content(self, phoenix_sync): + """Test that non-string content is preserved.""" + trajectory = { + "trace_id": "test", + "messages": [ + {"role": "user", "content": ["list", "content"]} + ] + } + cleaned = phoenix_sync._clean_trajectory(trajectory) + assert len(cleaned["messages"]) == 1 + assert cleaned["messages"][0]["content"] == ["list", "content"] + + +# ============================================================================= +# sync() Tests +# ============================================================================= + + +@pytest.mark.unit +class TestSync: + """Tests for sync method.""" + + @patch("kaizen.sync.phoenix_sync.urllib.request.urlopen") + @patch("kaizen.sync.phoenix_sync.generate_tips") + def test_sync_creates_namespace_if_not_exists(self, mock_generate_tips, mock_urlopen, phoenix_sync): + """Test that sync creates namespace if it doesn't exist.""" + from kaizen.schema.exceptions import NamespaceNotFoundException + + phoenix_sync.client.get_namespace_details.side_effect = NamespaceNotFoundException() + mock_response = MagicMock() + mock_response.read.return_value = b'{"data": [], "next_cursor": null}' + mock_response.__enter__ = Mock(return_value=mock_response) + mock_response.__exit__ = Mock(return_value=False) + mock_urlopen.return_value = mock_response + + phoenix_sync.sync(limit=10) + + phoenix_sync.client.create_namespace.assert_called_once_with("test_namespace") + + @patch("kaizen.sync.phoenix_sync.urllib.request.urlopen") + @patch("kaizen.sync.phoenix_sync.generate_tips") + def test_sync_skips_already_processed(self, mock_generate_tips, mock_urlopen, phoenix_sync): + """Test that already processed spans are skipped.""" + mock_response = MagicMock() + mock_response.read.return_value = json.dumps({ + "data": [ + { + "name": "litellm_request", + "context": {"trace_id": "t1", "span_id": "already_processed"}, + "attributes": { + "gen_ai.prompt.0.role": "user", + "gen_ai.prompt.0.content": "test" + } + } + ], + "next_cursor": None + }).encode() + mock_response.__enter__ = Mock(return_value=mock_response) + mock_response.__exit__ = Mock(return_value=False) + mock_urlopen.return_value = mock_response + + # Mock that this span was already processed + mock_entity = MagicMock() + mock_entity.metadata = {"span_id": "already_processed"} + phoenix_sync.client.search_entities.return_value = [mock_entity] + + result = phoenix_sync.sync(limit=10) + + assert result.skipped == 1 + assert result.processed == 0 + + @patch("kaizen.sync.phoenix_sync.urllib.request.urlopen") + @patch("kaizen.sync.phoenix_sync.generate_tips") + def test_sync_filters_error_spans(self, mock_generate_tips, mock_urlopen, phoenix_sync): + """Test that error spans are filtered by default.""" + mock_response = MagicMock() + mock_response.read.return_value = json.dumps({ + "data": [ + { + "name": "litellm_request", + "status_code": "ERROR", + "context": {"trace_id": "t1", "span_id": "s1"}, + "attributes": { + "gen_ai.prompt.0.role": "user", + "gen_ai.prompt.0.content": "test" + } + } + ], + "next_cursor": None + }).encode() + mock_response.__enter__ = Mock(return_value=mock_response) + mock_response.__exit__ = Mock(return_value=False) + mock_urlopen.return_value = mock_response + + phoenix_sync.client.search_entities.return_value = [] + + result = phoenix_sync.sync(limit=10, include_errors=False) + + assert result.processed == 0 + + @patch("kaizen.sync.phoenix_sync.urllib.request.urlopen") + @patch("kaizen.sync.phoenix_sync.generate_tips") + def test_sync_includes_error_spans_when_requested(self, mock_generate_tips, mock_urlopen, phoenix_sync): + """Test that error spans are included when include_errors=True.""" + mock_response = MagicMock() + mock_response.read.return_value = json.dumps({ + "data": [ + { + "name": "litellm_request", + "status_code": "ERROR", + "context": {"trace_id": "t1", "span_id": "s1"}, + "start_time": "2024-01-15T10:00:00Z", + "attributes": { + "gen_ai.request.model": "test-model", + "gen_ai.prompt.0.role": "user", + "gen_ai.prompt.0.content": "test message" + } + } + ], + "next_cursor": None + }).encode() + mock_response.__enter__ = Mock(return_value=mock_response) + mock_response.__exit__ = Mock(return_value=False) + mock_urlopen.return_value = mock_response + + phoenix_sync.client.search_entities.return_value = [] + mock_generate_tips.return_value = [] + + result = phoenix_sync.sync(limit=10, include_errors=True) + + assert result.processed == 1 + + @patch("kaizen.sync.phoenix_sync.urllib.request.urlopen") + @patch("kaizen.sync.phoenix_sync.generate_tips") + def test_sync_filters_non_llm_spans(self, mock_generate_tips, mock_urlopen, phoenix_sync): + """Test that non-LLM spans are filtered out.""" + mock_response = MagicMock() + mock_response.read.return_value = json.dumps({ + "data": [ + { + "name": "some_other_span", + "context": {"trace_id": "t1", "span_id": "s1"}, + "attributes": {} + } + ], + "next_cursor": None + }).encode() + mock_response.__enter__ = Mock(return_value=mock_response) + mock_response.__exit__ = Mock(return_value=False) + mock_urlopen.return_value = mock_response + + phoenix_sync.client.search_entities.return_value = [] + + result = phoenix_sync.sync(limit=10) + + assert result.processed == 0 + + @patch("kaizen.sync.phoenix_sync.urllib.request.urlopen") + @patch("kaizen.sync.phoenix_sync.generate_tips") + def test_sync_processes_valid_spans(self, mock_generate_tips, mock_urlopen, phoenix_sync): + """Test that valid spans are processed.""" + mock_response = MagicMock() + mock_response.read.return_value = json.dumps({ + "data": [ + { + "name": "litellm_request", + "context": {"trace_id": "t1", "span_id": "s1"}, + "start_time": "2024-01-15T10:00:00Z", + "attributes": { + "gen_ai.request.model": "claude-3", + "gen_ai.prompt.0.role": "user", + "gen_ai.prompt.0.content": "Hello" + } + } + ], + "next_cursor": None + }).encode() + mock_response.__enter__ = Mock(return_value=mock_response) + mock_response.__exit__ = Mock(return_value=False) + mock_urlopen.return_value = mock_response + + phoenix_sync.client.search_entities.return_value = [] + mock_generate_tips.return_value = ["Tip 1", "Tip 2"] + + result = phoenix_sync.sync(limit=10) + + assert result.processed == 1 + assert result.tips_generated == 2 + phoenix_sync.client.update_entities.assert_called() + + @patch("kaizen.sync.phoenix_sync.urllib.request.urlopen") + @patch("kaizen.sync.phoenix_sync.generate_tips") + def test_sync_returns_correct_counts(self, mock_generate_tips, mock_urlopen, phoenix_sync): + """Test that sync returns correct counts in SyncResult.""" + mock_response = MagicMock() + mock_response.read.return_value = json.dumps({ + "data": [ + { + "name": "litellm_request", + "context": {"trace_id": "t1", "span_id": "new_span"}, + "start_time": "2024-01-15T10:00:00Z", + "attributes": { + "gen_ai.request.model": "claude-3", + "gen_ai.prompt.0.role": "user", + "gen_ai.prompt.0.content": "New message" + } + }, + { + "name": "litellm_request", + "context": {"trace_id": "t2", "span_id": "old_span"}, + "start_time": "2024-01-15T09:00:00Z", + "attributes": { + "gen_ai.request.model": "claude-3", + "gen_ai.prompt.0.role": "user", + "gen_ai.prompt.0.content": "Old message" + } + } + ], + "next_cursor": None + }).encode() + mock_response.__enter__ = Mock(return_value=mock_response) + mock_response.__exit__ = Mock(return_value=False) + mock_urlopen.return_value = mock_response + + # old_span was already processed + mock_entity = MagicMock() + mock_entity.metadata = {"span_id": "old_span"} + phoenix_sync.client.search_entities.return_value = [mock_entity] + mock_generate_tips.return_value = ["Generated tip"] + + result = phoenix_sync.sync(limit=10) + + assert isinstance(result, SyncResult) + assert result.processed == 1 + assert result.skipped == 1 + assert result.tips_generated == 1 + assert result.errors == [] + + @patch("kaizen.sync.phoenix_sync.urllib.request.urlopen") + @patch("kaizen.sync.phoenix_sync.generate_tips") + def test_sync_handles_processing_errors(self, mock_generate_tips, mock_urlopen, phoenix_sync): + """Test that processing errors are captured.""" + mock_response = MagicMock() + mock_response.read.return_value = json.dumps({ + "data": [ + { + "name": "litellm_request", + "context": {"trace_id": "t1", "span_id": "s1"}, + "start_time": "2024-01-15T10:00:00Z", + "attributes": { + "gen_ai.request.model": "claude-3", + "gen_ai.prompt.0.role": "user", + "gen_ai.prompt.0.content": "test" + } + } + ], + "next_cursor": None + }).encode() + mock_response.__enter__ = Mock(return_value=mock_response) + mock_response.__exit__ = Mock(return_value=False) + mock_urlopen.return_value = mock_response + + phoenix_sync.client.search_entities.return_value = [] + mock_generate_tips.side_effect = Exception("Tip generation failed") + + result = phoenix_sync.sync(limit=10) + + assert result.processed == 0 + assert len(result.errors) == 1 + assert "Tip generation failed" in result.errors[0] + + +# ============================================================================= +# _ensure_namespace() Tests +# ============================================================================= + + +@pytest.mark.unit +class TestEnsureNamespace: + """Tests for _ensure_namespace method.""" + + def test_ensure_namespace_exists(self, phoenix_sync): + """Test that existing namespace is not recreated.""" + phoenix_sync.client.get_namespace_details.return_value = MagicMock() + + phoenix_sync._ensure_namespace() + + phoenix_sync.client.create_namespace.assert_not_called() + + def test_ensure_namespace_creates_if_missing(self, phoenix_sync): + """Test that missing namespace is created.""" + from kaizen.schema.exceptions import NamespaceNotFoundException + phoenix_sync.client.get_namespace_details.side_effect = NamespaceNotFoundException() + + phoenix_sync._ensure_namespace() + + phoenix_sync.client.create_namespace.assert_called_once_with("test_namespace") + + +# ============================================================================= +# _get_processed_span_ids() Tests +# ============================================================================= + + +@pytest.mark.unit +class TestGetProcessedSpanIds: + """Tests for _get_processed_span_ids method.""" + + def test_get_processed_span_ids_empty(self, phoenix_sync): + """Test getting processed IDs when none exist.""" + phoenix_sync.client.search_entities.return_value = [] + + result = phoenix_sync._get_processed_span_ids() + + assert result == set() + + def test_get_processed_span_ids_with_entities(self, phoenix_sync): + """Test getting processed IDs from existing entities.""" + entity1 = MagicMock() + entity1.metadata = {"span_id": "span_1"} + entity2 = MagicMock() + entity2.metadata = {"span_id": "span_2"} + entity3 = MagicMock() + entity3.metadata = None # No metadata + + phoenix_sync.client.search_entities.return_value = [entity1, entity2, entity3] + + result = phoenix_sync._get_processed_span_ids() + + assert result == {"span_1", "span_2"} + + def test_get_processed_span_ids_namespace_not_found(self, phoenix_sync): + """Test that missing namespace returns empty set.""" + from kaizen.schema.exceptions import NamespaceNotFoundException + phoenix_sync.client.search_entities.side_effect = NamespaceNotFoundException() + + result = phoenix_sync._get_processed_span_ids() + + assert result == set()