diff --git a/.gitignore b/.gitignore index 1b3a147..d8ff38e 100644 --- a/.gitignore +++ b/.gitignore @@ -37,4 +37,6 @@ config.json .vscode/ .github/ .env +.agents/ +skills-lock.json tests/ \ No newline at end of file diff --git a/README.md b/README.md index d4a3426..c44a503 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # Floki 🧭 — MLflow Experiment Agentic Chatbot -![image](assets/floki-banner.png) +![image](assets/floki-banner-2.png) > **⚠️ Work in Progress:** This project is actively being developed. Features, structure, and documentation may change frequently. diff --git a/assets/floki-banner-2.png b/assets/floki-banner-2.png new file mode 100644 index 0000000..3799d16 Binary files /dev/null and b/assets/floki-banner-2.png differ diff --git a/diary.md b/diary.md index 39afce4..a9b959f 100644 --- a/diary.md +++ b/diary.md @@ -30,4 +30,16 @@ I also just noticed that once I raise a PR, Copilot suggests changes after revie This is actually quite useful, especially when Copilot is writing some code which I don't care enough about to understand in a lot of detail. Added a pretty banner! -![image](assets/floki-banner.png) \ No newline at end of file +![image](assets/floki-banner.png) + +## 5/4/2026 + +Integrated Langfuse users and sessions. But one issue is that on Langfuse, the session traces don't show `input` and `output`, so need to fix that. + +Also identified a bug. The `list_runs()` tool outputs metrics and params for each run, which explodes the number of tokens in the following prompt in the chain. Need to remove that and reserve these things for a separate tool. + +## 5/14/2026 + +Kind of fixed the last issue about the tool. Split it into two. + +Right now working on using the structured responses from langchain. The outputs kind of suck, because the model just returns a message, and any tables that get rendered are taken from intermediate tool calls, which is bad. But the structured response doesn't seem to be working. Langfuse or Groq dashboards don't show the intent routing model even being called. \ No newline at end of file diff --git a/docs/superpowers/specs/2026-05-21-mlflow-agent-manual-test-checklist-design.md b/docs/superpowers/specs/2026-05-21-mlflow-agent-manual-test-checklist-design.md new file mode 100644 index 0000000..fa68f4c --- /dev/null +++ b/docs/superpowers/specs/2026-05-21-mlflow-agent-manual-test-checklist-design.md @@ -0,0 +1,232 @@ +# Manual MLflow Agent Test Checklist — Design + +Date: 2026-05-21 +Owner: Saumit Paul +Status: Draft + +## Overview +Design a **manual prompt checklist** (with expected outcomes) that validates a CLI agent’s real-world MLflow workflows. The checklist is dataset-agnostic and uses placeholders, so it can run against any shared MLflow tracking DB. + +## Goals +- Provide a repeatable manual test suite that mirrors ML engineer workflows. +- Ensure answers are grounded in MLflow data with concrete experiment/run references. +- Validate summaries, comparisons, regressions, artifacts, failures, and reproducibility. +- Make pass/fail decisions unambiguous through explicit acceptance criteria. + +## Non-goals +- Automating tests in pytest or building test harnesses. +- Writing synthetic data or seeding MLflow. +- Enforcing UI/UX formatting beyond the required fields. + +## Assumptions & Constraints +- All engineers use a shared MLflow DB with populated experiments and runs. +- Prompts are run via the CLI agent, which can call MLflow tools. +- Placeholders (e.g., ``) will be replaced by real values. +- The agent should not hallucinate experiments/runs. + +## Checklist Format (per item) +**ID** · **Goal** · **Prompt** · **Expected outcome** · **Pass/Fail checks** · **Notes/Assumptions** + +## Acceptance Criteria (applies to every item) +1. **Grounding:** Response cites concrete experiment/run IDs or names from MLflow; no invented entities. +2. **Completeness:** Required fields are present (e.g., run_id, metric values, date ranges). +3. **Correctness:** Rankings/aggregations match MLflow data (e.g., top 5 truly top by metric). +4. **Missing-data handling:** Absent fields are explicitly called out. +5. **Concision:** Avoid full metric/param dumps unless asked. +6. **Tooling:** Uses MLflow tools; does not answer from assumptions. + +--- + +## Category A — Recency & Team Activity + +**REC-1** +**Goal:** Summarize the last week’s work. +**Prompt:** “Summarize experiments and runs created or updated in `` (last 7 days), grouped by experiment.” +**Expected outcome:** List of experiments (name + ID), run counts, notable metrics/tags, brief summary per experiment. +**Pass/Fail checks:** Grounded IDs, date range stated, counts present. +**Notes:** Use `` placeholder. + +**REC-2** +**Goal:** Identify most active experiments. +**Prompt:** “Which experiments had the most new runs in ``?” +**Expected outcome:** Ranked list with experiment IDs, run counts, and date range. +**Pass/Fail checks:** Correct ordering, counts shown. + +**REC-3** +**Goal:** Highlight significant recent runs. +**Prompt:** “Show the top 3 best runs added in `` by `` across all experiments.” +**Expected outcome:** Table with run_id, experiment_name, metric value, key params. +**Pass/Fail checks:** Proper ranking; run IDs present. + +## Category B — Experiment Discovery & Metadata + +**DISC-1** +**Goal:** Find experiments by dataset/model tags. +**Prompt:** “List experiments that mention `` or `` in their tags or names.” +**Expected outcome:** Experiment IDs/names with tag evidence. +**Pass/Fail checks:** No hallucinations; evidence cited. + +**DISC-2** +**Goal:** Provide a quick experiment overview. +**Prompt:** “Give a one‑paragraph overview of `` including run count and top metric.” +**Expected outcome:** Experiment ID, run count, top metric summary. +**Pass/Fail checks:** Includes ID and metric name/value. + +**DISC-3** +**Goal:** Inventory experiments by owner/team tag. +**Prompt:** “Show experiments tagged with `=`.” +**Expected outcome:** Experiment list with IDs and tag proof. +**Pass/Fail checks:** Tag evidence present. + +## Category C — Run Comparison & Leaderboards + +**COMP-1** +**Goal:** Compare top runs within an experiment. +**Prompt:** “Compare the top 5 runs in `` by ``.” +**Expected outcome:** Table with run_id, metric, key params; sorted by metric. +**Pass/Fail checks:** Ordering correct; includes run IDs and metric values. + +**COMP-2** +**Goal:** Compare two specific runs. +**Prompt:** “Compare run `` vs `` on metrics and params.” +**Expected outcome:** Side‑by‑side comparison of key metrics/params. +**Pass/Fail checks:** Both run IDs referenced; no missing fields unacknowledged. + +**COMP-3** +**Goal:** Compare runs by a tag filter. +**Prompt:** “Among runs tagged `=` in ``, show the best 3 by ``.” +**Expected outcome:** Filtered ranked list with run IDs and metric values. +**Pass/Fail checks:** Filter applied; ranking correct. + +## Category D — Regression Tracking + +**REG-1** +**Goal:** Detect metric regression over time. +**Prompt:** “Has `` regressed in `` over `` compared to the prior period?” +**Expected outcome:** Regression/Stable/Improved + cited runs/metrics. +**Pass/Fail checks:** Includes referenced runs and date windows. + +**REG-2** +**Goal:** Identify last known good run. +**Prompt:** “What was the last run in `` before `` that achieved `` ≥ ``?” +**Expected outcome:** Single run_id with metric value and timestamp. +**Pass/Fail checks:** Run ID and timestamp present. + +**REG-3** +**Goal:** Explain regression drivers. +**Prompt:** “If there’s a regression, which params changed most between the best run last period and best run this period?” +**Expected outcome:** Param diffs tied to specific runs. +**Pass/Fail checks:** References both runs and their params. + +## Category E — Hyperparameter Sweeps + +**HYP-1** +**Goal:** Identify impactful hyperparameters. +**Prompt:** “Which hyperparameters most influence `` in ``?” +**Expected outcome:** Ranked list with evidence from runs. +**Pass/Fail checks:** Mentions run IDs or aggregated evidence. + +**HYP-2** +**Goal:** Find optimal param value. +**Prompt:** “For ``, which value correlates with the best ``?” +**Expected outcome:** Best value + supporting runs. +**Pass/Fail checks:** Value and run evidence present. + +**HYP-3** +**Goal:** Sweep coverage. +**Prompt:** “How many unique values were tried for `` in ``?” +**Expected outcome:** Count and list of values (if small). +**Pass/Fail checks:** Count present; values or explicit truncation. + +## Category F — Artifact Inspection + +**ART-1** +**Goal:** Locate best model artifact. +**Prompt:** “Show the model artifact path for the best run in `` by ``.” +**Expected outcome:** Run_id, artifact URI/path, metric value. +**Pass/Fail checks:** All fields present. + +**ART-2** +**Goal:** Retrieve evaluation artifacts. +**Prompt:** “List available evaluation artifacts (e.g., confusion matrix) for run ``.” +**Expected outcome:** Artifact list with paths. +**Pass/Fail checks:** Run_id referenced; artifacts enumerated or ‘none found’. + +**ART-3** +**Goal:** Compare artifacts across runs. +**Prompt:** “Compare the evaluation artifacts for the top 2 runs in ``.” +**Expected outcome:** Run IDs and artifact differences. +**Pass/Fail checks:** Evidence of both runs and artifacts. + +## Category G — Failure & Debugging + +**DBG-1** +**Goal:** Find failed runs. +**Prompt:** “List runs marked FAILED in `` with their error tags or notes.” +**Expected outcome:** Run IDs, failure status, error tags/notes. +**Pass/Fail checks:** Status and error evidence included. + +**DBG-2** +**Goal:** Detect missing metrics. +**Prompt:** “Are there runs in `` missing ``?” +**Expected outcome:** Run IDs missing the metric and count. +**Pass/Fail checks:** Explicit missing list or clear ‘none’. + +**DBG-3** +**Goal:** Negative test for nonexistent experiment. +**Prompt:** “Summarize ``.” +**Expected outcome:** Clear ‘not found’ response with a suggestion to list experiments. +**Pass/Fail checks:** No hallucinated results. + +## Category H — Reproducibility & Lineage + +**REP-1** +**Goal:** Produce a reproducibility recipe. +**Prompt:** “Give a reproducibility recipe for run `` (params, git SHA, data version, env).” +**Expected outcome:** All fields listed; missing fields explicitly stated. +**Pass/Fail checks:** Clear list + missing data callouts. + +**REP-2** +**Goal:** Identify lineage for a best run. +**Prompt:** “For the best run in ``, list dataset version, code version, and model artifact.” +**Expected outcome:** Run_id, dataset tag, git SHA/tag, artifact path. +**Pass/Fail checks:** All fields present or missing noted. + +**REP-3** +**Goal:** Detect inconsistent tagging. +**Prompt:** “Which runs in `` are missing a git SHA tag?” +**Expected outcome:** Run IDs missing the tag; count. +**Pass/Fail checks:** Count and run IDs or explicit ‘none’. + +## Category I — Performance & Governance + +**PERF-1** +**Goal:** Identify slow runs. +**Prompt:** “Which runs in `` are slowest (by duration), and how do their params differ?” +**Expected outcome:** Ranked durations + run IDs + parameter diffs. +**Pass/Fail checks:** Durations and run IDs included. + +**PERF-2** +**Goal:** Find unusually long runs. +**Prompt:** “List runs with duration > `` in ``.” +**Expected outcome:** Run IDs, durations, experiment names. +**Pass/Fail checks:** Threshold applied; IDs present. + +**GOV-1** +**Goal:** Identify stale experiments. +**Prompt:** “Which experiments have no runs in ``?” +**Expected outcome:** Experiment IDs/names with zero‑run confirmation. +**Pass/Fail checks:** IDs present; date range stated. + +**GOV-2** +**Goal:** Identify orphaned runs. +**Prompt:** “Are there runs without a meaningful experiment description/tag?” +**Expected outcome:** Run IDs or experiment IDs lacking descriptions/tags. +**Pass/Fail checks:** Evidence of missing metadata. + +--- + +## Execution Notes +- Replace placeholders with real values from your MLflow DB. +- Start with REC-1 to validate basic connectivity and grounding. +- If any item fails grounding or correctness, capture the prompt and response verbatim for triage. diff --git a/run_agent.sh b/run_agent.sh index ad25b9c..c601090 100644 --- a/run_agent.sh +++ b/run_agent.sh @@ -1,3 +1,3 @@ #!/bin/bash -python ./src/agent/langchain_agent.py \ No newline at end of file +python ./src/app.py \ No newline at end of file diff --git a/src/agent/__init__.py b/src/agent/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/agent/agent_middleware.py b/src/agent/agent_middleware.py new file mode 100644 index 0000000..17d139f --- /dev/null +++ b/src/agent/agent_middleware.py @@ -0,0 +1,70 @@ +from langchain.agents.middleware import wrap_tool_call +from langchain.agents.structured_output import ProviderStrategy +from langchain.messages import ToolMessage +import logging + + +# JSON Schema for BlockResponse (strict, provider-agnostic). +BLOCK_RESPONSE_SCHEMA = { + "title": "BlockResponse", + "type": "object", + "properties": { + "blocks": { + "type": "array", + "minItems": 1, + "items": { + "oneOf": [ + { + "type": "object", + "properties": { + "type": {"const": "text"}, + "markdown": {"type": "string"}, + }, + "required": ["type", "markdown"], + "additionalProperties": False, + }, + { + "type": "object", + "properties": { + "type": {"const": "table"}, + "markdown": {"type": "string"}, + }, + "required": ["type", "markdown"], + "additionalProperties": False, + }, + ] + }, + } + }, + "required": ["blocks"], + "additionalProperties": False, +} + +_logger = logging.getLogger(__name__) + + +@wrap_tool_call +def handle_tool_errors(request, handler): + """Handle tool execution errors with custom messages. + + This middleware wraps tool invocation and converts exceptions into a + structured ToolMessage so the agent receives a friendly error instead of + an exception bubbling up. + """ + try: + return handler(request) + except Exception as e: + return ToolMessage( + content=f"Tool error: Please check your input and try again. ({str(e)})", + tool_call_id=request.tool_call["id"] + ) + + +@wrap_tool_call +def set_block_response_schema(request, handler): + """Force the agent to use the BlockResponse JSON schema for structured output.""" + try: + request.response_format = ProviderStrategy(schema=BLOCK_RESPONSE_SCHEMA) + except Exception as e: + _logger.exception("Failed to set block response schema: %s", e) + return handler(request) diff --git a/src/agent/console_ui.py b/src/agent/console_ui.py index 7896979..9dfb9ef 100644 --- a/src/agent/console_ui.py +++ b/src/agent/console_ui.py @@ -1,6 +1,8 @@ from rich.console import Console from rich.table import Table from rich import box +from rich.markdown import Markdown +from rich.text import Text import json import re from typing import Any, List, Optional @@ -16,12 +18,16 @@ def print_text(text: str): console.print(text) -def print_welcome(app_name: str = "Floki", version: str = "v0.1"): - """Print a colorful ASCII welcome banner with brief usage hints. +def print_user(text: str): + """Print a user message with a clear 'You' tag.""" + try: + console.print(f"[bold white on dark_blue] You [/bold white on dark_blue] {text}") + except Exception: + console.print(f"You: {text}") - This is intended to be called once after the agent finishes initialization - and before the user starts interacting with the CLI. - """ + +# Purple → Blue gradient with glow effect +def print_welcome(app_name: str = "Floki", version: str = "v0.1"): from rich.panel import Panel from rich.align import Align from rich.text import Text @@ -38,15 +44,39 @@ def print_welcome(app_name: str = "Floki", version: str = "v0.1"): " ▀ ▀ ", ] - banner = "\n".join(banner_lines) + # Purple → Blue gradient + colors = [ + "#dd22ff", # bright magenta + "#cc22ff", + "#bb22ff", + "#9922ff", # purple + "#7722ff", + "#5533ff", # purple-blue + "#3355ff", # blue + "#2266ff", # bright blue + ] + t = Text() - t.append(banner + "\n", style="bold magenta") + + # Add glow layer (dim version as shadow) + banner_glow = "\n".join(banner_lines) + t.append(banner_glow + "\n\n\n", style="dim #5533ff") + + # Overlay gradient banner on top + t_banner = Text() + for line, color in zip(banner_lines, colors): + t_banner.append(line + "\n", style=f"bold {color}") + + t = Text() + for line, color in zip(banner_lines, colors): + t.append(line + "\n", style=f"bold {color}") + t.append("\n") t.append(f" {app_name} {version} 🧭\n", style="bold white on dark_green") t.append("\n") t.append("Welcome! Type your question and press Enter. Type 'exit' to quit.\n", style="cyan") t.append("Tool outputs (tables/JSON) will appear below the assistant message when available.\n", style="dim") - panel = Panel(Align.center(t), padding=(1, 2), border_style="magenta") + panel = Panel(Align.center(t), padding=(1, 2), border_style="bright_blue") console.print(panel) @@ -123,124 +153,125 @@ def _print_list_of_dicts(lst: List[dict]): console.print(table) -def print_result(result: Any): - """Render a langchain/agent result structure nicely. +def _render_block_response(block_resp: dict): + """Render a BlockResponse dict which contains ordered blocks. + + Each block is expected to be a dict with keys: + - type: "text" or "table" + - markdown: string containing markdown + """ + blocks = block_resp.get("blocks") or [] + for blk in blocks: + if not isinstance(blk, dict): + console.print(str(blk)) + continue + btype = blk.get("type") + md = blk.get("markdown", "") + if btype == "text": + try: + console.print(Markdown(md)) + except Exception: + console.print(md) + elif btype == "table": + # Prefer to render as a parsed markdown table for nicer column alignment + md_table = _parse_markdown_table(md) + if md_table: + _print_list_of_dicts(md_table) + else: + try: + console.print(Markdown(md)) + except Exception: + console.print(md) + else: + # Unknown block type; print raw representation + if isinstance(md, str) and md.strip(): + try: + console.print(Markdown(md)) + except Exception: + console.print(md) + else: + console.print(str(blk)) + - Strategy: - - Prefer structured tool outputs (JSON) found in tool messages. - - If not found, try to parse a markdown table from the agent text. - - Fallback to pretty JSON or plain text. +def print_result(result: Any): + """Render a langchain/agent result structure assuming BlockResponse schema. + + Behavior: + - Prefer the assistant's final message and expect it to be a BlockResponse + (JSON with a top-level 'blocks' array). Each block is rendered in order. + - If the final message does not contain a valid BlockResponse, fall back to + printing the content as Markdown or a parsed markdown table when possible. + - Prints a clear 'Floki' tag before assistant outputs. """ if not result: console.print("") return + # Prefer structured_response from ToolStrategy when available + if isinstance(result, dict) and isinstance(result.get("structured_response"), dict): + _render_block_response(result["structured_response"]) + return + messages = None if isinstance(result, dict): messages = result.get('messages') elif hasattr(result, 'messages'): messages = result.messages + # If there are no messages, maybe the result itself is already a BlockResponse if not messages: - console.print(result) + if isinstance(result, dict) and isinstance(result.get("blocks"), list): + _render_block_response(result) + return + # Fallback: print raw/pretty + try: + if isinstance(result, str): + console.print(Markdown(result)) + else: + console.print(result) + except Exception: + console.print(str(result)) return - # Show assistant final text first (if present) - printed_assistant = False + # Focus on the assistant's final message only (ignore intermediate tool outputs) + last = messages[-1] try: - last = messages[-1] - assistant_text = getattr(last, 'content') if hasattr(last, 'content') or isinstance(last, dict) else None + if isinstance(last, dict): + content = last.get('content') + else: + content = getattr(last, 'content', None) except Exception: - assistant_text = None - if isinstance(assistant_text, str) and assistant_text.strip(): - console.print(assistant_text) - printed_assistant = True - - # First pass: search messages from newest to oldest for structured outputs - rendered_structured = False - for msg in reversed(messages): - # tool calls metadata might be in additional_kwargs or tool_calls - tw = None - try: - tw = getattr(msg, 'additional_kwargs', None) or getattr(msg, 'tool_calls', None) or (msg.get('tool_calls') if isinstance(msg, dict) else None) - except Exception: - tw = None - content = None - try: - content = getattr(msg, 'content') - except Exception: - try: - content = msg.get('content') if isinstance(msg, dict) else None - except Exception: - content = None - if not content: - continue - extracted = _extract_json(content) - if extracted is not None: - # Print any tool-call metadata associated with this message only - if tw and not rendered_structured: - console.print(f"\n[bold cyan]Parsed tool output (from message):[/bold cyan] {tw}") - elif not rendered_structured: - console.print(f"\n[bold cyan]Parsed tool output:[/bold cyan]") - - # Render the structured JSON from this message as supplemental output - if isinstance(extracted, list) and extracted and isinstance(extracted[0], dict): - _print_list_of_dicts(extracted) - rendered_structured = True - continue - elif isinstance(extracted, dict): - try: - console.print_json(json.dumps(extracted)) - except Exception: - console.print(str(extracted)) - rendered_structured = True - continue - - # If we rendered any structured supplemental output, stop here (assistant text already shown) - if rendered_structured: + # If content is already a dict-like BlockResponse + if isinstance(content, dict) and isinstance(content.get('blocks'), list): + _render_block_response(content) return - # Second pass: try to find markdown table in message texts - for msg in messages: - try: - content = getattr(msg, 'content') - except Exception: - try: - content = msg.get('content') if isinstance(msg, dict) else None - except Exception: - content = None - if not content: - continue + # If content is a string, try to extract JSON (handles embedded/quoted JSON) + if isinstance(content, str): + extracted = _extract_json(content) + if isinstance(extracted, dict) and isinstance(extracted.get('blocks'), list): + _render_block_response(extracted) + return + # Try parsing a markdown table as a fallback, but render remaining text too md_table = _parse_markdown_table(content) if md_table: _print_list_of_dicts(md_table) + # print remaining text after table if present + import re as _re + m = _re.search(r"(\|[^\n]+\n\|[ \-:|]+\n(?:\|.*\n)*)", content, _re.DOTALL) + if m: + rest = content.replace(m.group(1), "", 1).strip() + if rest: + console.print(Markdown(rest)) return + # Otherwise render as Markdown + console.print(Markdown(content)) + return - # Fallback: print the last message content prettily - last = messages[-1] - try: - content = getattr(last, 'content') - except Exception: - try: - content = last.get('content') if isinstance(last, dict) else str(last) - except Exception: - content = str(last) - - # If we already printed the assistant text above, avoid printing it again + # Final fallback: print stringified content try: - if printed_assistant and isinstance(content, str) and isinstance(assistant_text, str) and content.strip() == assistant_text.strip(): - return + console.print(str(content)) except Exception: - pass - - # try JSON - extracted = _extract_json(content) if isinstance(content, str) else None - if extracted is not None: - try: - console.print_json(json.dumps(extracted)) - except Exception: - console.print(extracted) - else: console.print(content) diff --git a/src/agent/langchain_agent.py b/src/agent/langchain_agent.py index 4396267..7a9be5e 100644 --- a/src/agent/langchain_agent.py +++ b/src/agent/langchain_agent.py @@ -4,6 +4,7 @@ This agent uses LangGraph to orchestrate LLM reasoning and MLflow tool calls. """ + import os import json import sys @@ -12,8 +13,10 @@ sys.path.append(os.path.dirname(os.path.dirname(__file__))) from llm.inference_engine import GroqEngine, get_llm_from_config from mlflow_tools import data_access -from langfuse import get_client -from langfuse.langchain import CallbackHandler +from llm.tracing import setup_langfuse, propagate_attributes +from langgraph.checkpoint.memory import InMemorySaver +from langchain.agents.structured_output import ToolStrategy +from agent.agent_middleware import handle_tool_errors, BLOCK_RESPONSE_SCHEMA from dotenv import load_dotenv @@ -43,48 +46,36 @@ logging.info("Loading model from config: %s", llm_config.get('groq_model', 'Not Set')) llm = get_llm_from_config(llm_config) -# Langfuse setup — use environment variables only for simplicity -langfuse_handler = None -public_key = os.getenv('LANGFUSE_PUBLIC_KEY') -host = os.getenv('LANGFUSE_BASE_URL') -try: - if public_key: - # Prefer passing public_key; avoid secret_key kwarg for compatibility - try: - fuse_client = get_client(public_key=public_key, host=host) if host else get_client(public_key=public_key) - except TypeError: - # Fallback: set env vars and call get_client() - os.environ.setdefault('LANGFUSE_PUBLIC_KEY', public_key) - if host: - os.environ.setdefault('LANGFUSE_BASE_URL', host) - fuse_client = get_client() - try: - langfuse_handler = CallbackHandler(client=fuse_client) - except TypeError: - langfuse_handler = CallbackHandler() - logging.info("Langfuse handler initialized.") - else: - logging.info("LANGFUSE_PUBLIC_KEY not set; Langfuse tracing disabled.") -except Exception as e: - logging.exception("Failed to initialize Langfuse handler: %s", e) - langfuse_handler = None + +# Langfuse/tracing setup +langfuse_handler, fuse_client, conversation_id, lf_run, langfuse_user, FLUSH_PER_QUERY = setup_langfuse(config) mlflow_tools = data_access.get_all_tools() +checkpointer = InMemorySaver() + + agent = create_agent( model=llm, tools=mlflow_tools, + checkpointer=checkpointer, + response_format=ToolStrategy(schema=BLOCK_RESPONSE_SCHEMA), system_prompt=( "You are a precise MLflow experiment assistant. RULES:\n" - "1) Always use the provided tools to fetch or query MLflow data. Do not invent or guess run IDs, experiment IDs, metrics, parameters, or artifact locations.\n" - "2) When a user requests data (experiments, runs, metrics, params, artifacts), CALL the appropriate tool and DO NOT embed raw data in your assistant message. The tool's structured output will be rendered by the UI.\n" - "3) After a tool call, provide a short natural-language summary (no tables, no code blocks) of <=2 sentences describing the high-level result and next steps.\n" - "4) If asked to return data directly, return valid JSON only (array or object), no Markdown, no ASCII tables.\n" - "5) On tool errors, return a JSON object: {\"error\": , \"message\": }. Do not raise exceptions.\n" - "6) For any action that may be destructive, ask for explicit confirmation before proceeding.\n" - "7) Keep responses concise and focused on user's goal.\n" - "Adhere strictly to these rules." + "1) Always use the provided tools to fetch or query MLflow data. Do not invent or guess information.\n" + "2) Keep responses concise and focused on the user's MLflow goals.\n" + "3) If a tool encounters an error, explain the issue in a text block. Do not raise exceptions.\n" + "4) For destructive actions, ask for explicit confirmation in a text block before proceeding.\n" + "\n" + "OUTPUT SCHEMA (edit this section if needed):\n" + "- Return JSON only, in this exact shape:\n" + " {\"blocks\": [{\"type\": \"text\", \"markdown\": \"...\"} | {\"type\": \"table\", \"markdown\": \"|h|...\"}]}\n" + "- Use type=\"text\" for analysis, summaries, and next steps.\n" + "- Use type=\"table\" only for clean Markdown pipe tables when comparisons are requested or helpful.\n" + "- Do not use TextBlock/TableBlock or any other keys; only blocks/type/markdown are allowed." ), + middleware=[handle_tool_errors], + ) @@ -93,7 +84,15 @@ def run_query(user_query: str): config_kwargs = {} if langfuse_handler is not None: logging.info("Attaching Langfuse handler to agent invocation.") + # attach handler config_kwargs['callbacks'] = [langfuse_handler] + config_kwargs['configurable'] = {'thread_id': conversation_id or "default_thread"} + # include conversation metadata if agent/client supports it + if conversation_id is not None: + meta = config_kwargs.setdefault('metadata', {}) + meta['conversation_id'] = conversation_id + if langfuse_user: + meta['user'] = langfuse_user # Only pass config when non-empty to avoid passing None handlers if config_kwargs: result = agent.invoke({"messages": messages}, config=config_kwargs) @@ -102,6 +101,22 @@ def run_query(user_query: str): return result +def _print_result(result): + try: + import console_ui as ui + ui.print_result(result) + return + except Exception: + pass + try: + print(f"\n{result['messages'][-1].content}") + except Exception: + try: + print(f"\n{result}") + except Exception: + pass + + def loading_animation(message, duration=3): spinner = ['|', '/', '-', '\\'] @@ -118,30 +133,106 @@ def main(): print("==============================") print("Initializing agent and loading tools...") loading_animation("Starting up, please wait...", duration=3) + print("Agent is ready! Type 'exit' to quit.") # Try to show a colorful welcome banner (non-fatal) try: import console_ui as ui ui.print_welcome() except Exception: - logging.debug("console_ui not available for welcome banner") + logging.info("console_ui not available for welcome banner") + # Print tracing info if enabled + if fuse_client is not None: + print(f"Langfuse tracing enabled. See at {os.getenv('LANGFUSE_BASE_URL', 'your Langfuse dashboard')}") + + # Enter session-level attribute propagation for grouping traces/observations + if conversation_id is not None and fuse_client is not None: + session_prop = propagate_attributes(session_id=conversation_id, user_id=langfuse_user) + else: + session_prop = None + + if session_prop is not None: + with session_prop: + _interactive_loop(fuse_client) + else: + _interactive_loop(fuse_client) + + +def _get_user_input(): + """Get user input using prompt_toolkit if available, otherwise fall back. + + Tries prompt_toolkit (rich features), then Rich Console.input, then builtin input. + Raises EOFError up to the caller to handle termination. + """ + try: + from prompt_toolkit import PromptSession + from prompt_toolkit.formatted_text import HTML + session = PromptSession() + prompt_html = HTML('You ') + return session.prompt(prompt_html) + except Exception: + try: + import console_ui as ui + return ui.console.input("\n[bold blue]You[/bold blue] ") + except Exception: + # last fallback to builtin input; let EOFError bubble up + return input("\n> ") + + +def _interactive_loop(fuse_client_local): while True: - user_query = input("\n> ") + try: + user_query = _get_user_input() + except EOFError: + print("\nGoodbye!") + break + except KeyboardInterrupt: + # User pressed Ctrl-C; continue the loop to allow graceful exit + print("\nInterrupted. Goodbye!") + continue if user_query.strip().lower() in {"exit", "quit"}: + if fuse_client_local is not None and not FLUSH_PER_QUERY: + try: + fuse_client_local.flush() + except Exception: + pass print("Goodbye!") break - result = run_query(user_query) - # Render result using rich console UI if available - try: - import console_ui as ui - ui.print_result(result) - except Exception: - logging.warning("console_ui not available or failed to render result, falling back to plain print.") - # fallback: print last message or raw + + # Create a root observation/span for this query so trace-level IO is populated + if fuse_client_local is not None and hasattr(fuse_client_local, 'start_as_current_observation'): try: - print(f"\n{result['messages'][-1].content}") + with fuse_client_local.start_as_current_observation(as_type="span", name="langchain-call") as obs_ctx: + try: + obs_ctx.update(input={"query": user_query}) + except Exception: + pass + result = run_query(user_query) + _print_result(result) + # Try to extract a readable output snippet + output_snippet = None + try: + output_snippet = result.get('messages')[-1].content + except Exception: + try: + output_snippet = str(result) + except Exception: + output_snippet = None + if output_snippet is not None: + try: + obs_ctx.update(output={"result": output_snippet}) + except Exception: + pass + if FLUSH_PER_QUERY: + try: + fuse_client_local.flush() + except Exception: + pass except Exception: - print(f"\n{result}") - -if __name__ == "__main__": - main() + # Fallback to running without explicit observation context + result = run_query(user_query) + _print_result(result) + else: + # No Langfuse client or observation support; just run the query + result = run_query(user_query) + _print_result(result) diff --git a/src/app.py b/src/app.py new file mode 100644 index 0000000..b90d613 --- /dev/null +++ b/src/app.py @@ -0,0 +1,21 @@ + +import sys +sys.path.append('./src/agent') # Add agent directory to path for imports + +from langchain_agent import main, lf_run, fuse_client, conversation_id + +if __name__ == "__main__": + try: + main() + except KeyboardInterrupt: + print("\nInterrupted. Finishing Langfuse run if present...") + try: + if lf_run is not None: + if hasattr(lf_run, 'finish'): + lf_run.finish() + elif fuse_client is not None and hasattr(fuse_client, 'finish_run'): + fuse_client.finish_run(conversation_id) + if fuse_client is not None: + fuse_client.flush() + except Exception: + pass diff --git a/src/llm/tracing.py b/src/llm/tracing.py new file mode 100644 index 0000000..5e92218 --- /dev/null +++ b/src/llm/tracing.py @@ -0,0 +1,75 @@ +""" +Langfuse tracing setup utilities for MLflow agent +""" + +import os +import time +import logging + +# Optionally import langfuse if available +try: + from langfuse import get_client, propagate_attributes + from langfuse.langchain import CallbackHandler +except ImportError: + get_client = propagate_attributes = CallbackHandler = None + logging.warning("Langfuse not installed; tracing will be disabled.") + +# Exposed module-level variable for user metadata +langfuse_user = None + + +def setup_langfuse(config): + """ + Set up Langfuse tracing based on config and environment variables. + Returns: langfuse_handler, fuse_client, conversation_id, lf_run, FLUSH_PER_QUERY + """ + global langfuse_user + langfuse_handler = None + langfuse_user = config.get('langfuse', {}).get('user', 'unknown_user') + public_key = os.getenv('LANGFUSE_PUBLIC_KEY') + host = os.getenv('LANGFUSE_BASE_URL') + fuse_client = None + conversation_id = None + lf_run = None + FLUSH_PER_QUERY = os.getenv('LANGFUSE_FLUSH_PER_QUERY', 'false').lower() == 'true' + if get_client is None: + return None, None, None, None, None, FLUSH_PER_QUERY + try: + if public_key: + try: + fuse_client = get_client(public_key=public_key, host=host) if host else get_client(public_key=public_key) + except TypeError: + os.environ.setdefault('LANGFUSE_PUBLIC_KEY', public_key) + if host: + os.environ.setdefault('LANGFUSE_BASE_URL', host) + fuse_client = get_client() + try: + langfuse_handler = CallbackHandler(client=fuse_client) + except TypeError: + langfuse_handler = CallbackHandler() + try: + import uuid + conversation_id = f"cli-{uuid.uuid4().hex[:8]}" + except Exception: + conversation_id = f"cli-{int(time.time())}" + try: + run_kwargs = {} + if langfuse_user: + run_kwargs['user'] = langfuse_user + if hasattr(fuse_client, 'start_run'): + lf_run = fuse_client.start_run(name=conversation_id, **run_kwargs) + elif hasattr(fuse_client, 'runs') and hasattr(fuse_client.runs, 'create'): + lf_run = fuse_client.runs.create(name=conversation_id, **run_kwargs) + except Exception: + lf_run = None + logging.info("Langfuse handler initialized. conversation_id=%s", conversation_id) + else: + logging.info("LANGFUSE_PUBLIC_KEY not set; Langfuse tracing disabled.") + except Exception as e: + logging.exception("Failed to initialize Langfuse handler: %s", e) + langfuse_handler = None + + # Return the initialized objects (may be None if disabled) + return langfuse_handler, fuse_client, conversation_id, lf_run, langfuse_user, FLUSH_PER_QUERY +# Re-export propagate_attributes for convenience +__all__ = ["setup_langfuse", "propagate_attributes", "langfuse_user"] diff --git a/src/mlflow_tools/data_access.py b/src/mlflow_tools/data_access.py index fa3728d..5bb3bab 100644 --- a/src/mlflow_tools/data_access.py +++ b/src/mlflow_tools/data_access.py @@ -12,6 +12,7 @@ from typing import Any, Dict, List, Optional # mlflow is optional for unit tests in minimal environments try: + logging.getLogger("mlflow").setLevel(logging.WARNING) import mlflow from mlflow.tracking import MlflowClient from mlflow.exceptions import MlflowException @@ -31,7 +32,8 @@ def _decorator(f): from . import schemas # Keep MLflow logs quieter by default -os.environ.setdefault("MLFLOW_LOGGING_LEVEL", "WARNING") +# os.environ.setdefault("MLFLOW_LOGGING_LEVEL", "WARNING") +# logging.getLogger("mlflow").setLevel(logging.WARNING) # Load mlruns_dir from global config CONFIG_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../config.json')) @@ -83,40 +85,59 @@ def raw_list_experiments(include_deleted: bool = False, max_results: int = 100) def raw_list_runs( - experiment_ids: List[str], + experiment_id: str, status: Optional[List[str]] = None, start_time: Optional[int] = None, end_time: Optional[int] = None, order_by: Optional[str] = None, max_results: int = 100, + include_metrics: bool = False, ) -> List[Dict[str, Any]]: """Return summarized runs for given experiments. Note: MLflowClient.search_runs accepts experiment_ids and order_by; more complex filtering can be added later. """ + if experiment_id.lower() in ["all","*"]: + raise ValueError("Listing runs across all experiments is not supported for token economy. Please specify a single experiment ID.") try: - runs = client.search_runs(experiment_ids, order_by=[order_by] if order_by else None, max_results=max_results) + runs = client.search_runs([experiment_id], order_by=[order_by] if order_by else None, max_results=max_results) except MlflowException as e: logging.error("Error listing runs: %s", e) raise out = [] for run in runs: - metrics = dict(getattr(run, 'data', SimpleNamespace()).metrics) if hasattr(run, 'data') else {} - params = dict(getattr(run, 'data', SimpleNamespace()).params) if hasattr(run, 'data') else {} - out.append({ + run_info = { 'run_id': run.info.run_id, 'run_name': getattr(run.info, 'run_name', None), 'status': getattr(run.info, 'status', None), 'start_time_iso': _iso_from_epoch_ms(getattr(run.info, 'start_time', None)), 'end_time_iso': _iso_from_epoch_ms(getattr(run.info, 'end_time', None)), - 'metrics_preview': {k: metrics[k] for i, k in enumerate(metrics) if i < 5}, - 'params_preview': {k: params[k] for i, k in enumerate(params) if i < 10}, - }) + } + if include_metrics: + metrics = dict(getattr(run, 'data', SimpleNamespace()).metrics) if hasattr(run, 'data') else {} + params = dict(getattr(run, 'data', SimpleNamespace()).params) if hasattr(run, 'data') else {} + run_info['metrics_preview'] = {k: metrics[k] for i, k in enumerate(metrics) if i < 5} + run_info['params_preview'] = {k: params[k] for i, k in enumerate(params) if i < 10} + out.append(run_info) return out +# New: Count runs per experiment ID +def raw_count_runs_per_experiment(experiment_ids: List[str]) -> Dict[str, int]: + """Return a dict of experiment_id -> number of runs.""" + counts = {} + for exp_id in experiment_ids: + try: + runs = client.search_runs([exp_id], max_results=50000) + counts[exp_id] = len(runs) + except MlflowException as e: + logging.error(f"Error counting runs for experiment {exp_id}: {e}") + counts[exp_id] = -1 + return counts + + def raw_get_run_metrics(run_id: str) -> Dict[str, float]: """Return a dict of metric_name -> latest_value for the run.""" try: @@ -145,6 +166,11 @@ def raw_find_best_runs_by_metric( filter_params: Optional[Dict[str, str]] = None, ) -> List[Dict[str, Any]]: """Return top-k runs ordered by metric (max or min).""" + + if isinstance(experiment_ids, str): + experiment_ids = [experiment_ids] + if experiment_ids and isinstance(experiment_ids, list) and isinstance(experiment_ids[0], str) and experiment_ids[0].lower() in ["all","*"]: + raise ValueError("To search across all experiments, use a list of all experiment IDs obtained from list_experiments.") order = f"metrics.{metric} DESC" if mode == 'max' else f"metrics.{metric} ASC" try: runs = client.search_runs(experiment_ids, order_by=[order], max_results=top_k) @@ -193,7 +219,7 @@ def raw_check_experiment_generalization( raise experiment_id = getattr(exp, 'experiment_id', None) - runs = raw_list_runs([experiment_id], max_results=1000) + runs = raw_list_runs(experiment_id=experiment_id, max_results=1000) failing = [] for r in runs: @@ -219,32 +245,68 @@ def raw_check_experiment_generalization( @tool(description="List MLflow experiments (tool wrapper).", args_schema=schemas.ListExperimentsParams) def list_experiments_tool(include_deleted: bool = False, max_results: int = 100): - return raw_list_experiments(include_deleted=include_deleted, max_results=max_results) + result = raw_list_experiments(include_deleted=include_deleted, max_results=max_results) + try: + return json.dumps(result, default=str) + except Exception: + return str(result) + + + +# Update tool wrapper for new signature +@tool(description="List MLflow runs for a single experiment (tool wrapper).", args_schema=schemas.ListRunsParams) +def list_runs_tool(experiment_id: str, status: Optional[List[str]] = None, start_time: Optional[int] = None, end_time: Optional[int] = None, order_by: Optional[str] = None, max_results: int = 100, include_metrics: bool = False): + result = raw_list_runs(experiment_id=experiment_id, status=status, start_time=start_time, end_time=end_time, order_by=order_by, max_results=max_results, include_metrics=include_metrics) + try: + return json.dumps(result, default=str) + except Exception: + return str(result) -@tool(description="List MLflow runs (tool wrapper).", args_schema=schemas.ListRunsParams) -def list_runs_tool(experiment_ids: List[str], status: Optional[List[str]] = None, start_time: Optional[int] = None, end_time: Optional[int] = None, order_by: Optional[str] = None, max_results: int = 100): - return raw_list_runs(experiment_ids=experiment_ids, status=status, start_time=start_time, end_time=end_time, order_by=order_by, max_results=max_results) +# New tool wrapper for counting runs +@tool(description="Count MLflow runs per experiment (tool wrapper).", args_schema=schemas.CountRunsPerExperimentParams) +def count_runs_per_experiment_tool(experiment_ids: List[str]): + result = raw_count_runs_per_experiment(experiment_ids) + try: + return json.dumps(result, default=str) + except Exception: + return str(result) @tool(description="Get run metrics (tool wrapper).", args_schema=schemas.GetRunMetricsParams) def get_run_metrics_tool(run_id: str): - return raw_get_run_metrics(run_id) + result = raw_get_run_metrics(run_id) + try: + return json.dumps(result, default=str) + except Exception: + return str(result) @tool(description="Get run params (tool wrapper).", args_schema=schemas.GetRunParamsParams) def get_run_params_tool(run_id: str): - return raw_get_run_params(run_id) + result = raw_get_run_params(run_id) + try: + return json.dumps(result, default=str) + except Exception: + return str(result) @tool(description="Find top runs by metric (tool wrapper).", args_schema=schemas.FindBestRunByMetricParams) def find_best_runs_by_metric_tool(experiment_ids: List[str], metric: str, mode: str = 'max', top_k: int = 1): - return raw_find_best_runs_by_metric(experiment_ids=experiment_ids, metric=metric, mode=mode, top_k=top_k) + result = raw_find_best_runs_by_metric(experiment_ids=experiment_ids, metric=metric, mode=mode, top_k=top_k) + try: + return json.dumps(result, default=str) + except Exception: + return str(result) @tool(description="Check experiment generalization (tool wrapper).", args_schema=schemas.CheckExperimentGeneralizationParams) def check_experiment_generalization_tool(experiment_id_or_name: str, train_metric: str = 'train_loss', test_metric: str = 'test_loss', threshold_abs: Optional[float] = None, threshold_rel: Optional[float] = 0.2): - return raw_check_experiment_generalization(experiment_id_or_name=experiment_id_or_name, train_metric=train_metric, test_metric=test_metric, threshold_abs=threshold_abs, threshold_rel=threshold_rel) + result = raw_check_experiment_generalization(experiment_id_or_name=experiment_id_or_name, train_metric=train_metric, test_metric=test_metric, threshold_abs=threshold_abs, threshold_rel=threshold_rel) + try: + return json.dumps(result, default=str) + except Exception: + return str(result) def get_all_tools(): @@ -252,6 +314,7 @@ def get_all_tools(): return [ list_experiments_tool, list_runs_tool, + count_runs_per_experiment_tool, get_run_metrics_tool, get_run_params_tool, find_best_runs_by_metric_tool, diff --git a/src/mlflow_tools/schemas.py b/src/mlflow_tools/schemas.py index 50ac1f2..f6595fc 100644 --- a/src/mlflow_tools/schemas.py +++ b/src/mlflow_tools/schemas.py @@ -9,12 +9,18 @@ class ListExperimentsParams(BaseModel): max_results: int = Field(100, description="Maximum number of experiments to return") class ListRunsParams(BaseModel): - experiment_ids: List[str] = Field(..., description="IDs of the experiments to list runs for.") + experiment_id: str = Field(..., description="ID of the experiment to list runs for.") status: Optional[List[str]] = Field(None, description="Filter by run status, e.g. ['FINISHED']") start_time: Optional[int] = Field(None, description="Only runs started after this epoch ms") end_time: Optional[int] = Field(None, description="Only runs started before this epoch ms") - order_by: Optional[str] = Field(None, description="order_by clause for MLflow search_runs") + order_by: Optional[str] = Field(None, description="order_by clause for MLflow search_runs e.g. 'metrics.accuracy DESC', 'attributes.start_time ASC'") max_results: int = Field(100, description="Maximum number of runs to return.") + include_metrics: bool = Field(False, description="Whether to include metrics in the response.") + + +# New: schema for counting runs per experiment +class CountRunsPerExperimentParams(BaseModel): + experiment_ids: List[str] = Field(..., description="IDs of the experiments to count runs for.") class GetRunMetricsParams(BaseModel): run_id: str = Field(..., description="ID of the run to get metrics for.")