diff --git a/extract_trajectories.py b/extract_trajectories.py new file mode 100644 index 00000000..592536b7 --- /dev/null +++ b/extract_trajectories.py @@ -0,0 +1,584 @@ +#!/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 spans from an Arize Phoenix server until the requested limit is reached or no more spans are available. + + Parameters: + base_url (str): Base URL of the Phoenix server (e.g., "http://localhost:6006"). + limit (int): Maximum number of spans to retrieve. + + Returns: + list[dict]: A list of span dictionaries retrieved from the server; length will be at most `limit` and may be smaller if the server has fewer spans. + """ + 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: + """ + Normalize content by parsing string representations of JSON or Python literals. + + If `content` is a string, the function first attempts to parse it as JSON. If JSON parsing fails, it attempts to parse the string as a Python literal (e.g., dict or list) using `ast.literal_eval`. If both parses fail, the original string is returned. Non-string inputs are returned unchanged. + + Parameters: + content (Any): The value to normalize; commonly a string containing a serialized JSON or Python literal, or an already-parsed object. + + Returns: + Any: The parsed object when parsing succeeds, or the original `content` if parsing is not applicable or fails. + """ + 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]: + """ + Extracts prompt and completion messages from a span's attributes. + + Scans the span's `attributes` for keys of the form `gen_ai.prompt.{i}.role` / `gen_ai.prompt.{i}.content` + and `gen_ai.completion.{i}.role` / `gen_ai.completion.{i}.content`, collects matching entries, and + returns them as a list of message dictionaries sorted by their index within each type. + + Parameters: + span (dict): A span dictionary (as returned by Phoenix) containing an "attributes" mapping. + + Returns: + list[dict]: A list of messages where each message contains: + - index (int): The numeric index parsed from the attribute keys. + - type (str): Either `"prompt"` or `"completion"`. + - role (str): The role value from the span attributes. + - content (Any): The parsed content (normalized from strings when applicable). + """ + 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 a message expressed in Anthropic block format into an OpenAI-style message dictionary. + + Handles three input shapes: + - string: returned as a single message with the given role and the string as content. + - non-list: coerced to string and returned as a single message. + - list of block dictionaries: processes block types (`text`, `thinking`, `tool_use`, `tool_result`) and maps them into OpenAI-style fields. For `role == "assistant"` the result may include `thinking` (non-standard), `content`, and `tool_calls`. For `role == "user"` with tool results, returns a `role: "tool"` message containing `tool_results`. Otherwise returns a message with `role` and joined `content`. + + Parameters: + content (Any): Anthropic-formatted message (string, other scalar, or list of block dicts). + role (str): Role to assign in the resulting message (e.g., "assistant", "user"). + + Returns: + dict: An OpenAI-style message dictionary. Examples: + - {"role": "", "content": ""} for simple inputs, + - {"role": "assistant", "thinking": "<...>", "content": "<...>", "tool_calls": [...]} for assistant blocks, + - {"role": "tool", "tool_results": [...]} when user-sourced tool results are present. + """ + + 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: + """ + Builds a complete trajectory dictionary from a Phoenix span, converting Anthropic-style prompts/completions and tool results into OpenAI-style messages. + + Parameters: + span (dict): A Phoenix span dictionary containing at least `attributes` and `context` keys. Attributes may include `gen_ai.prompt.*`, `gen_ai.completion.*`, model and usage fields. + + Returns: + dict: A trajectory with the following keys: + - trace_id (str): Trace identifier from span["context"]["trace_id"]. + - span_id (str): Span identifier from span["context"]["span_id"]. + - model (str): Model name from `gen_ai.request.model` or "unknown" if absent. + - timestamp: Span start time from `span["start_time"]`. + - messages (list[dict]): OpenAI-style messages converted from prompts, completions, and expanded tool results. Tool result messages include `role: "tool"`, `tool_call_id`, and `content`. + - usage (dict): Token usage fields with keys `prompt_tokens`, `completion_tokens`, and `total_tokens` populated from span attributes when available. + """ + 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: + """ + Strip segments wrapped in ... from the given text. + + Removes all occurrences of ... (including their contents, possibly spanning multiple lines) and trims leading/trailing whitespace. + + Returns: + The cleaned string with system reminder segments removed. + """ + import re + return re.sub(r'.*?', '', text, flags=re.DOTALL).strip() + + +def clean_trajectory(trajectory: dict, remove_system_reminders: bool = True) -> dict: + """ + Remove empty messages and optionally strip segments from message content in a trajectory. + + Parameters: + trajectory (dict): Trajectory dictionary containing a "messages" list and other metadata. + remove_system_reminders (bool): If True, remove text segments wrapped in ... from string message content. Defaults to True. + + Returns: + dict: A copy of the input trajectory with "messages" replaced by a filtered list where: + - Messages with no `content` and no `tool_calls` are removed. + - If `remove_system_reminders` is enabled, messages whose content becomes empty after stripping are removed. + """ + 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 spans from a Phoenix server and convert them into trajectories in OpenAI chat completion format. + + Parameters: + base_url (str): Phoenix server base URL. + limit (int): Maximum number of spans to fetch. + include_errors (bool): If True, include spans with error status; otherwise exclude them. + clean (bool): If True, clean message contents (e.g., remove system reminders) before returning. + + Returns: + list[dict]: List of trajectory dictionaries. Each trajectory contains keys such as `trace_id`, `span_id`, `model`, `timestamp`, `messages`, and `usage`. + """ + 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: + """ + Produce a human-readable multi-line text representation of a trajectory for inspection. + + Parameters: + trajectory (dict): Trajectory object containing keys like `trace_id`, `model`, `timestamp`, and a `messages` list where each message may have `role`, `content`, `thinking`, `tool_calls`, and `tool_call_id`. + include_thinking (bool): If True, include agent "thinking" or reasoning blocks when present; otherwise omit them. + + Returns: + str: A formatted multi-line string summarizing the trajectory, including a header (trace id, model, timestamp) and sequentially rendered messages for USER, ASSISTANT (with optional thinking and tool call details), and TOOL results. + """ + 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(): + """ + Entry point for the CLI that extracts agent trajectories from an Arize Phoenix server and writes them to stdout or a file. + + Parses command-line options, fetches trajectories according to the provided flags (URL, limit, include-errors, cleaning), optionally filters by trace ID, sorts results by timestamp (most recent first), and outputs either JSON (optionally pretty-printed) or a human-readable text representation. When an output file is specified, writes the results to that file and prints a brief confirmation with the number of trajectories written. + """ + 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() \ No newline at end of file diff --git a/kaizen/frontend/cli/cli.py b/kaizen/frontend/cli/cli.py index a38207b7..95a070ba 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,73 @@ 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, +): + """ + Import trajectories from Arize Phoenix into the local system and generate tips. + + Prints the sync configuration, runs the Phoenix sync process, displays a results table + with metrics (trajectories processed, trajectories skipped, tips generated, and errors), + and prints error details when present. On unexpected failure the command exits with + code 1. + + Parameters: + phoenix_url (Optional[str]): Phoenix server URL to connect to. If None, the syncer + will use its default or configured URL. + namespace (Optional[str]): Target Kaizen namespace ID to store generated tips. + project (Optional[str]): Phoenix project name to sync. + limit (int): Maximum number of spans/trajectories to fetch from Phoenix. + include_errors (bool): If True, include failed/error spans in the sync. + """ + 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() + app() \ No newline at end of file diff --git a/kaizen/sync/phoenix_sync.py b/kaizen/sync/phoenix_sync.py new file mode 100644 index 00000000..d968a32b --- /dev/null +++ b/kaizen/sync/phoenix_sync.py @@ -0,0 +1,514 @@ +""" +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, + ): + """ + Initialize the PhoenixSync instance. + + Parameters: + phoenix_url (str | None): Optional override for the Arize Phoenix API base URL; if omitted, the value from `phoenix_settings.url` is used. + namespace_id (str | None): Optional override for the target Kaizen namespace ID; if omitted, the value from `kaizen_config.namespace_id` is used. + project (str | None): Optional override for the Phoenix project name; if omitted, the value from `phoenix_settings.project` is used. + """ + 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 Kaizen namespace identified by self.namespace_id exists, creating it if it does not. + + Creates the namespace through the Kaizen client when missing and logs the creation. + """ + 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]: + """ + Retrieve up to `limit` span objects from the Phoenix API, following pagination until the requested number is reached or no further pages are available. + + Parameters: + limit (int): Maximum number of spans to fetch. + + Returns: + list[dict]: A list of span objects returned by Phoenix. + + Raises: + Exception: Propagates any exception raised while making HTTP requests or parsing responses. + """ + 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]: + """ + Return the set of span IDs for trajectories already stored in the target Kaizen namespace. + + Returns: + set[str]: Span ID strings found in trajectory entities' metadata; returns an empty set if no span IDs are found or if the namespace does not exist. + """ + 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 a value that may be a JSON or Python literal encoded as a string. + + If `content` is a string, this function first attempts to parse it as JSON. If JSON parsing fails, it then attempts to evaluate it as a Python literal (e.g., list, dict, tuple, number, boolean) using ast.literal_eval. If both attempts fail, the original string is returned unchanged. Non-string inputs are returned as-is. + + Parameters: + content (Any): The value to parse, or a string containing a JSON object/array or a Python literal. + + Returns: + Any: The parsed Python object when parsing succeeds, otherwise the original `content`. + """ + 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]: + """ + Collect prompt and completion messages from a span's `gen_ai` attributes. + + Searches the span's attributes for matching pairs `gen_ai.prompt.{i}.role`/`gen_ai.prompt.{i}.content` + and `gen_ai.completion.{i}.role`/`gen_ai.completion.{i}.content`, parses each content value via + _self._parse_content_, and returns the messages found. + + Returns: + list[dict]: List of message dictionaries. Each dictionary contains: + - 'index' (int): the numeric index extracted from the attribute keys. + - 'type' (str): either 'prompt' or 'completion'. + - 'role' (str): the role string from the attribute. + - 'content' (Any): the parsed content value. + """ + 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: + """ + Translate Anthropic-style message blocks into an OpenAI-compatible message dictionary. + + Accepts `content` that may be a string, any non-list value, or a list of block dictionaries. + Recognizes block types: + - "text": collected into the message content (skips "(no content)"), + - "thinking": collected into a `thinking` field, + - "tool_use": converted into `tool_calls` entries, + - "tool_result": converted into `tool_results` entries. + + @param content: The message payload to convert; either a plain string, another scalar (coerced to string), or a list of block dicts as described above. + @param role: The sender role (e.g., "assistant" or "user") which influences the output shape. + + @returns: + dict: An OpenAI-like message. Possible keys: + - "role": the provided role or "tool" for user tool results, + - "content": concatenated text blocks or None when assistant has only tool calls, + - "thinking": concatenated thinking blocks (if any), + - "tool_calls": list of function call descriptors (if any), + - "tool_results": list of tool result objects (only when returning a tool message). + """ + 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: + """ + Builds a trajectory dictionary from a Phoenix span. + + Parses the span's attributes and extracted messages, converts each message into OpenAI-compatible message structures (including expanded tool result messages), and returns a consolidated trajectory record. + + Returns: + dict: Trajectory with the following keys: + - trace_id (str): Span's trace identifier. + - span_id (str): Span's span identifier. + - model (str): Model name from `gen_ai.request.model` or `"unknown"`. + - timestamp: Span start time value. + - messages (list[dict]): OpenAI-style messages. Tool result messages are represented as + {"role": "tool", "tool_call_id": , "content": } while regular messages follow + standard OpenAI message shapes (e.g., {"role": "user"|"assistant", "content": , ...}). + - usage (dict): Token usage with keys: + - prompt_tokens + - completion_tokens + - total_tokens + """ + 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: + """ + Remove messages that have no content and no tool calls, and strip `` tags from string message content. + + Returns: + dict: A trajectory dictionary identical to the input but with the `messages` list filtered so that: + - Any message with neither `content` nor `tool_calls` is removed. + - If a message's `content` is a string, any `...` sections are removed and the content is trimmed. + """ + 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: + """ + Store trajectory messages as 'trajectory' entities and generate/store tips as 'guideline' entities. + + Returns: + int: The number of tips generated and stored. + """ + # 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: + """ + Orchestrates a full sync: fetches spans from Phoenix, filters and processes trajectories, and generates tips stored in Kaizen. + + Parameters: + limit (int): Maximum number of spans to fetch from Phoenix for this run. + include_errors (bool): If True, include spans with status "ERROR"; otherwise such spans are skipped. + + Returns: + SyncResult: Counts of processed spans, skipped spans, tips generated, and a list of error messages encountered while processing. + """ + 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 \ No newline at end of file