From 0a63becac3a31cc8cf971d4e8296b44c42aab958 Mon Sep 17 00:00:00 2001 From: Saumit Paul Date: Tue, 5 May 2026 01:32:27 -0400 Subject: [PATCH 1/6] Enhance Langfuse integration and add conversation tracking in langchain_agent --- diary.md | 8 ++- src/agent/langchain_agent.py | 110 +++++++++++++++++++++++++++++------ 2 files changed, 99 insertions(+), 19 deletions(-) diff --git a/diary.md b/diary.md index 39afce4..2dd728a 100644 --- a/diary.md +++ b/diary.md @@ -30,4 +30,10 @@ 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. \ No newline at end of file diff --git a/src/agent/langchain_agent.py b/src/agent/langchain_agent.py index 4396267..29272b3 100644 --- a/src/agent/langchain_agent.py +++ b/src/agent/langchain_agent.py @@ -12,8 +12,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 +# import langfuse +from langfuse import get_client, propagate_attributes from langfuse.langchain import CallbackHandler +from langgraph.store.memory import InMemoryStore from dotenv import load_dotenv @@ -45,8 +47,13 @@ # Langfuse setup — use environment variables only for simplicity 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') +# conversation/run tracking +fuse_client = None +conversation_id = None +lf_run = None try: if public_key: # Prefer passing public_key; avoid secret_key kwarg for compatibility @@ -62,7 +69,26 @@ langfuse_handler = CallbackHandler(client=fuse_client) except TypeError: langfuse_handler = CallbackHandler() - logging.info("Langfuse handler initialized.") + # create a stable conversation id for grouping traces + try: + import uuid + conversation_id = f"cli-{uuid.uuid4().hex[:8]}" + except Exception: + conversation_id = f"cli-{int(time.time())}" + # try to start a run (optional; SDKs vary) + try: + # include user metadata when creating the run if supported + 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: @@ -70,10 +96,12 @@ langfuse_handler = None mlflow_tools = data_access.get_all_tools() +store = InMemoryStore() agent = create_agent( model=llm, tools=mlflow_tools, + store=store, 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" @@ -93,7 +121,14 @@ 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] + # 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) @@ -126,22 +161,61 @@ def main(): except Exception: logging.debug("console_ui not available for welcome banner") while True: - user_query = input("\n> ") - if user_query.strip().lower() in {"exit", "quit"}: - print("Goodbye!") - break - result = run_query(user_query) - # Render result using rich console UI if available + # Guard Langfuse context: only use if client is available + if fuse_client is not None: + obs_ctx = fuse_client.start_as_current_observation(as_type="span", name="langchain-call") if hasattr(fuse_client, 'start_as_current_observation') else None + else: + obs_ctx = None + if obs_ctx is not None: + obs_enter = obs_ctx.__enter__ + obs_exit = obs_ctx.__exit__ + obs_enter() + # propagate attributes for grouping (no-op if not configured) + # propagate conversation id and user so Langfuse groups events + prop_ctx = propagate_attributes(session_id=conversation_id, user_id=langfuse_user) if conversation_id is not None else None 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 - try: - print(f"\n{result['messages'][-1].content}") - except Exception: - print(f"\n{result}") + if prop_ctx is not None: + prop_ctx.__enter__() + user_query = input("\n> ") + if user_query.strip().lower() in {"exit", "quit"}: + fuse_client.flush() + print("Goodbye!") + break + result = run_query(user_query) + finally: + if prop_ctx is not None: + try: + prop_ctx.__exit__(None, None, None) + except Exception: + pass + if obs_ctx is not None: + try: + obs_exit(None, None, None) + except Exception: + pass + # 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 + try: + print(f"\n{result['messages'][-1].content}") + except Exception: + print(f"\n{result}") if __name__ == "__main__": - 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) + fuse_client.flush() + except Exception: + pass From bf356343bacce905d28af23fae8aef57cebb9ea3 Mon Sep 17 00:00:00 2001 From: Saumit Paul Date: Sat, 9 May 2026 20:11:04 -0400 Subject: [PATCH 2/6] Added memory and fixed Langfuse sessions - Langfuse sessions now correctly show root input and output - Simplified the Langfuse session code - Session (short term) memory works now, and agent can answer questions about previous queries too. --- .gitignore | 2 + src/agent/langchain_agent.py | 139 +++++++++++++++++++++++++---------- 2 files changed, 102 insertions(+), 39 deletions(-) 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/src/agent/langchain_agent.py b/src/agent/langchain_agent.py index 29272b3..e14034c 100644 --- a/src/agent/langchain_agent.py +++ b/src/agent/langchain_agent.py @@ -15,7 +15,8 @@ # import langfuse from langfuse import get_client, propagate_attributes from langfuse.langchain import CallbackHandler -from langgraph.store.memory import InMemoryStore +from langgraph.graph import StateGraph +from langgraph.checkpoint.memory import InMemorySaver from dotenv import load_dotenv @@ -54,6 +55,7 @@ fuse_client = None conversation_id = None lf_run = None +FLUSH_PER_QUERY = os.getenv('LANGFUSE_FLUSH_PER_QUERY', 'false').lower() == 'true' try: if public_key: # Prefer passing public_key; avoid secret_key kwarg for compatibility @@ -85,6 +87,21 @@ 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) + # normalize run id extraction helper + def _extract_run_id(run_obj): + if run_obj is None: + return None + try: + if hasattr(run_obj, 'id'): + return getattr(run_obj, 'id') + if isinstance(run_obj, dict): + return run_obj.get('id') or run_obj.get('run_id') or run_obj.get('name') + if isinstance(run_obj, str): + return run_obj + except Exception: + return None + return None + lf_run_id = _extract_run_id(lf_run) except Exception: lf_run = None @@ -96,12 +113,12 @@ langfuse_handler = None mlflow_tools = data_access.get_all_tools() -store = InMemoryStore() +checkpointer = InMemorySaver() agent = create_agent( model=llm, tools=mlflow_tools, - store=store, + checkpointer=checkpointer, 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" @@ -123,6 +140,7 @@ def run_query(user_query: str): 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', {}) @@ -137,6 +155,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 = ['|', '/', '-', '\\'] @@ -153,6 +187,7 @@ 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: @@ -160,50 +195,76 @@ def main(): ui.print_welcome() except Exception: logging.debug("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 _interactive_loop(fuse_client_local): while True: - # Guard Langfuse context: only use if client is available - if fuse_client is not None: - obs_ctx = fuse_client.start_as_current_observation(as_type="span", name="langchain-call") if hasattr(fuse_client, 'start_as_current_observation') else None - else: - obs_ctx = None - if obs_ctx is not None: - obs_enter = obs_ctx.__enter__ - obs_exit = obs_ctx.__exit__ - obs_enter() - # propagate attributes for grouping (no-op if not configured) - # propagate conversation id and user so Langfuse groups events - prop_ctx = propagate_attributes(session_id=conversation_id, user_id=langfuse_user) if conversation_id is not None else None try: - if prop_ctx is not None: - prop_ctx.__enter__() user_query = input("\n> ") - if user_query.strip().lower() in {"exit", "quit"}: - fuse_client.flush() - print("Goodbye!") - break - result = run_query(user_query) - finally: - if prop_ctx is not None: - try: - prop_ctx.__exit__(None, None, None) - except Exception: - pass - if obs_ctx is not None: + except EOFError: + print("\nGoodbye!") + break + if user_query.strip().lower() in {"exit", "quit"}: + if fuse_client_local is not None and not FLUSH_PER_QUERY: try: - obs_exit(None, None, None) + fuse_client_local.flush() except Exception: pass - # 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 + print("Goodbye!") + break + + # 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: + with fuse_client_local.start_as_current_observation(as_type="span", name="langchain-call") as obs_ctx: try: - print(f"\n{result['messages'][-1].content}") + obs_ctx.update(input={"query": user_query}) except Exception: - print(f"\n{result}") + 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: + # 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) if __name__ == "__main__": try: From 97d455d1473a354719396dd77269e74f4cc9e390 Mon Sep 17 00:00:00 2001 From: Saumit Paul Date: Sat, 9 May 2026 22:08:26 -0400 Subject: [PATCH 3/6] Refactor raw_list_runs to accept a single experiment ID and add include_metrics option; introduce raw_count_runs_per_experiment and corresponding tool wrapper --- src/mlflow_tools/data_access.py | 48 +++++++++++++++++++++++++-------- src/mlflow_tools/schemas.py | 7 ++++- 2 files changed, 43 insertions(+), 12 deletions(-) diff --git a/src/mlflow_tools/data_access.py b/src/mlflow_tools/data_access.py index fa3728d..64b2cdd 100644 --- a/src/mlflow_tools/data_access.py +++ b/src/mlflow_tools/data_access.py @@ -83,12 +83,13 @@ 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. @@ -96,27 +97,43 @@ def raw_list_runs( filtering can be added later. """ 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: @@ -222,9 +239,17 @@ def list_experiments_tool(include_deleted: bool = False, max_results: int = 100) return raw_list_experiments(include_deleted=include_deleted, max_results=max_results) -@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) + +# 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): + return 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) + + +# 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]): + return raw_count_runs_per_experiment(experiment_ids) @tool(description="Get run metrics (tool wrapper).", args_schema=schemas.GetRunMetricsParams) @@ -252,6 +277,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..ba12b38 100644 --- a/src/mlflow_tools/schemas.py +++ b/src/mlflow_tools/schemas.py @@ -9,13 +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") max_results: int = Field(100, description="Maximum number of runs to return.") + +# 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.") From b3d530aa639faaffe83b077d6bc5db0881b27a58 Mon Sep 17 00:00:00 2001 From: Saumit Paul Date: Tue, 12 May 2026 19:48:02 -0400 Subject: [PATCH 4/6] Refactor agent execution and enhance Langfuse tracing setup; add new app entry point and improve error handling in tool wrappers --- run_agent.sh | 2 +- src/agent/langchain_agent.py | 105 +++++++------------------------- src/app.py | 21 +++++++ src/llm/tracing.py | 75 +++++++++++++++++++++++ src/mlflow_tools/data_access.py | 52 +++++++++++++--- src/mlflow_tools/schemas.py | 2 +- 6 files changed, 163 insertions(+), 94 deletions(-) create mode 100644 src/app.py create mode 100644 src/llm/tracing.py 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/langchain_agent.py b/src/agent/langchain_agent.py index e14034c..9317bc9 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,11 +13,12 @@ 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 -# import langfuse -from langfuse import get_client, propagate_attributes -from langfuse.langchain import CallbackHandler +from llm.tracing import setup_langfuse, propagate_attributes +from llm.tracing import setup_langfuse, propagate_attributes, langfuse_user from langgraph.graph import StateGraph from langgraph.checkpoint.memory import InMemorySaver +from langchain.agents.middleware import wrap_tool_call +from langchain.messages import ToolMessage from dotenv import load_dotenv @@ -46,75 +48,25 @@ 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 -langfuse_user = config.get('langfuse', {}).get('user', 'unknown_user') -public_key = os.getenv('LANGFUSE_PUBLIC_KEY') -host = os.getenv('LANGFUSE_BASE_URL') -# conversation/run tracking -fuse_client = None -conversation_id = None -lf_run = None -FLUSH_PER_QUERY = os.getenv('LANGFUSE_FLUSH_PER_QUERY', 'false').lower() == 'true' -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() - # create a stable conversation id for grouping traces - try: - import uuid - conversation_id = f"cli-{uuid.uuid4().hex[:8]}" - except Exception: - conversation_id = f"cli-{int(time.time())}" - # try to start a run (optional; SDKs vary) - try: - # include user metadata when creating the run if supported - 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) - # normalize run id extraction helper - def _extract_run_id(run_obj): - if run_obj is None: - return None - try: - if hasattr(run_obj, 'id'): - return getattr(run_obj, 'id') - if isinstance(run_obj, dict): - return run_obj.get('id') or run_obj.get('run_id') or run_obj.get('name') - if isinstance(run_obj, str): - return run_obj - except Exception: - return None - return None - lf_run_id = _extract_run_id(lf_run) - 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 +# Langfuse/tracing setup +langfuse_handler, fuse_client, conversation_id, lf_run, FLUSH_PER_QUERY = setup_langfuse(config) mlflow_tools = data_access.get_all_tools() checkpointer = InMemorySaver() +@wrap_tool_call +def handle_tool_errors(request, handler): + """Handle tool execution errors with custom messages.""" + try: + return handler(request) + except Exception as e: + # Return a custom error message to the model + return ToolMessage( + content=f"Tool error: Please check your input and try again. ({str(e)})", + tool_call_id=request.tool_call["id"] + ) + agent = create_agent( model=llm, tools=mlflow_tools, @@ -130,6 +82,7 @@ def _extract_run_id(run_obj): "7) Keep responses concise and focused on user's goal.\n" "Adhere strictly to these rules." ), + middleware=[handle_tool_errors] ) @@ -145,7 +98,7 @@ def run_query(user_query: str): if conversation_id is not None: meta = config_kwargs.setdefault('metadata', {}) meta['conversation_id'] = conversation_id - if langfuse_user: + if langfuse_user: meta['user'] = langfuse_user # Only pass config when non-empty to avoid passing None handlers if config_kwargs: @@ -194,7 +147,7 @@ def main(): 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')}") @@ -266,17 +219,3 @@ def _interactive_loop(fuse_client_local): result = run_query(user_query) _print_result(result) -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) - fuse_client.flush() - except Exception: - pass 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..f026c91 --- /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, 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, 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 64b2cdd..0f6b71c 100644 --- a/src/mlflow_tools/data_access.py +++ b/src/mlflow_tools/data_access.py @@ -31,7 +31,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')) @@ -96,6 +97,8 @@ def raw_list_runs( 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_id], order_by=[order_by] if order_by else None, max_results=max_results) except MlflowException as e: @@ -162,6 +165,9 @@ 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 experiment_ids.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) @@ -210,7 +216,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: @@ -236,40 +242,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): - return 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) + 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) + try: + return json.dumps(result, default=str) + except Exception: + return str(result) # 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]): - return raw_count_runs_per_experiment(experiment_ids) + 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(): diff --git a/src/mlflow_tools/schemas.py b/src/mlflow_tools/schemas.py index ba12b38..988d183 100644 --- a/src/mlflow_tools/schemas.py +++ b/src/mlflow_tools/schemas.py @@ -13,7 +13,7 @@ class ListRunsParams(BaseModel): 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.") From 8bc3a6dcb997b9f3e0263b7aafef73268f248ec4 Mon Sep 17 00:00:00 2001 From: Saumit Paul Date: Thu, 21 May 2026 21:51:38 -0400 Subject: [PATCH 5/6] Enhance Langfuse integration by fixing session trace outputs, updating list_runs_tool to include metrics, and adding error handling middleware; introduce manual test checklist for MLflow agent workflows. --- diary.md | 8 +- ...flow-agent-manual-test-checklist-design.md | 232 ++++++++++++++++++ src/agent/__init__.py | 0 src/agent/agent_middleware.py | 102 ++++++++ src/agent/langchain_agent.py | 22 +- src/llm/tracing.py | 4 +- src/mlflow_tools/data_access.py | 4 +- src/mlflow_tools/schemas.py | 1 + 8 files changed, 351 insertions(+), 22 deletions(-) create mode 100644 docs/superpowers/specs/2026-05-21-mlflow-agent-manual-test-checklist-design.md create mode 100644 src/agent/__init__.py create mode 100644 src/agent/agent_middleware.py diff --git a/diary.md b/diary.md index 2dd728a..a9b959f 100644 --- a/diary.md +++ b/diary.md @@ -36,4 +36,10 @@ Added a pretty banner! 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. \ No newline at end of file +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/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..248a112 --- /dev/null +++ b/src/agent/agent_middleware.py @@ -0,0 +1,102 @@ +from langchain.agents.middleware import wrap_tool_call +from langchain.agents.structured_output import ProviderStrategy +from langchain.messages import ToolMessage +from langchain.chat_models import init_chat_model +from pydantic import BaseModel, Field +from typing import Dict, Literal, Type +import os +import logging + + +# Output schema constants (exported for tests) +class SimpleResponse(BaseModel): + message: str = Field(description="Short natural language response") + + +class TableResponse(BaseModel): + columns: list[str] = Field(description="Column headers for the table") + rows: list[list[str]] = Field(description="Table rows as lists of string values") + + +SIMPLE_RESPONSE_SCHEMA = {"type": "simple", "description": "Short natural language response"} +TABLE_RESPONSE_SCHEMA = {"type": "table", "description": "Tabular data response"} + +SCHEMA_REGISTRY: Dict[str, Type[BaseModel]] = { + "simple": SimpleResponse, + "table": TableResponse, +} + + +class IntentRouter(BaseModel): + """Select the output schema required for the user's request.""" + selected_schema: Literal["simple", "table"] = Field( + description=( + "Choose 'table' for rankings, lists, comparisons, or any response that should be displayed" + " as rows/columns. Choose 'simple' for direct explanations or short answers." + ) + ) + +_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 classify_and_set_schema(request, handler): + """Classify user intent and set an output schema on the tool_call metadata. + + Uses a small Groq model to decide whether the user's query expects a + tabular (table) response or a simple natural-language response. The + decision is recorded at `request.tool_call['metadata']['output_schema']`. + """ + try: + messages = getattr(request, "messages", None) + if not messages: + return handler(request) + + last = messages[-1] + role = last.get("role") if isinstance(last, dict) else getattr(last, "role", None) + if role != "human": + return handler(request) + + groq_api_key = os.getenv("GROQ_API_KEY") + selected_key = "simple" + if groq_api_key: + try: + router_llm = init_chat_model("allam-2-7b", model_provider="groq", temperature=0) + structured_router = router_llm.with_structured_output(IntentRouter) + routing_decision = structured_router.invoke(messages) + selected_key = routing_decision.selected_schema + print(f"Intent classification result: {selected_key}") + except Exception as e: + _logger.warning("Groq classification failed: %s", e) + + tool_call = getattr(request, "tool_call", None) + if isinstance(tool_call, dict): + meta = tool_call.setdefault("metadata", {}) + meta["output_schema"] = TABLE_RESPONSE_SCHEMA if selected_key == "table" else SIMPLE_RESPONSE_SCHEMA + + target_schema = SCHEMA_REGISTRY[selected_key] + request.response_format = ProviderStrategy(schema=target_schema) + except Exception as e: + _logger.exception("Failed to classify intent: %s", e) + + return handler(request) + + diff --git a/src/agent/langchain_agent.py b/src/agent/langchain_agent.py index 9317bc9..806714b 100644 --- a/src/agent/langchain_agent.py +++ b/src/agent/langchain_agent.py @@ -14,11 +14,8 @@ from llm.inference_engine import GroqEngine, get_llm_from_config from mlflow_tools import data_access from llm.tracing import setup_langfuse, propagate_attributes -from llm.tracing import setup_langfuse, propagate_attributes, langfuse_user -from langgraph.graph import StateGraph from langgraph.checkpoint.memory import InMemorySaver -from langchain.agents.middleware import wrap_tool_call -from langchain.messages import ToolMessage +from agent.agent_middleware import handle_tool_errors, classify_and_set_schema from dotenv import load_dotenv @@ -50,22 +47,12 @@ # Langfuse/tracing setup -langfuse_handler, fuse_client, conversation_id, lf_run, FLUSH_PER_QUERY = setup_langfuse(config) +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() -@wrap_tool_call -def handle_tool_errors(request, handler): - """Handle tool execution errors with custom messages.""" - try: - return handler(request) - except Exception as e: - # Return a custom error message to the model - return ToolMessage( - content=f"Tool error: Please check your input and try again. ({str(e)})", - tool_call_id=request.tool_call["id"] - ) + agent = create_agent( model=llm, @@ -82,7 +69,8 @@ def handle_tool_errors(request, handler): "7) Keep responses concise and focused on user's goal.\n" "Adhere strictly to these rules." ), - middleware=[handle_tool_errors] + middleware=[classify_and_set_schema, handle_tool_errors], + ) diff --git a/src/llm/tracing.py b/src/llm/tracing.py index f026c91..5e92218 100644 --- a/src/llm/tracing.py +++ b/src/llm/tracing.py @@ -33,7 +33,7 @@ def setup_langfuse(config): 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, FLUSH_PER_QUERY + return None, None, None, None, None, FLUSH_PER_QUERY try: if public_key: try: @@ -70,6 +70,6 @@ def setup_langfuse(config): langfuse_handler = None # Return the initialized objects (may be None if disabled) - return langfuse_handler, fuse_client, conversation_id, lf_run, FLUSH_PER_QUERY + 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 0f6b71c..d42b2c7 100644 --- a/src/mlflow_tools/data_access.py +++ b/src/mlflow_tools/data_access.py @@ -252,8 +252,8 @@ def list_experiments_tool(include_deleted: bool = False, max_results: int = 100) # 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): - 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) +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: diff --git a/src/mlflow_tools/schemas.py b/src/mlflow_tools/schemas.py index 988d183..f6595fc 100644 --- a/src/mlflow_tools/schemas.py +++ b/src/mlflow_tools/schemas.py @@ -15,6 +15,7 @@ class ListRunsParams(BaseModel): 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 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 From 937660f08611419794b5bf8027104e575b43ee42 Mon Sep 17 00:00:00 2001 From: Saumit Paul Date: Fri, 22 May 2026 21:15:43 -0400 Subject: [PATCH 6/6] Structured Outputs functional --- README.md | 2 +- assets/floki-banner-2.png | Bin 0 -> 12694 bytes src/agent/agent_middleware.py | 108 ++++++--------- src/agent/console_ui.py | 231 ++++++++++++++++++-------------- src/agent/langchain_agent.py | 53 ++++++-- src/mlflow_tools/data_access.py | 7 +- 6 files changed, 216 insertions(+), 185 deletions(-) create mode 100644 assets/floki-banner-2.png 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 0000000000000000000000000000000000000000..3799d1660cbd113a0bf30da7cf762701f0c20450 GIT binary patch literal 12694 zcmch8X;>54x^=p3+GDFt(}<{mQBgnAObRrQBaUE0y2gqwjLE3l~KmD zj50>%IT3|K<_2VtDM*+@5<>`?ldsA?=id8#=lt@!&-a5TsY=z}yWX|lwcfpJ^X%#s z^PNBK`w0Sp?7Vc*%oYOqfdYYir|`EQfir&|KKB9m_%7Ji{O^#;e&rS5;0Ld>R%aoQ zs>JQ0TQb1$wtE+ygCUTg--CbOk&&-&Lm)bum(0#y3%^01M%}mPq7MjQyDuI%eE65h z18;A#>+g85-|RHF5}-VE#l_1o#JG|+bK7BYv1s`I9kn*^liukl6fHG0X{8`+q9`#- z&kemWvRD5`Uzt|mzR7k25>ozz?Bq@MuLpj)c=2K@1G=2O>{n1DtL(4(Y{09cX7>_z zQrLjSx81_>>1eg%{;naWz%U^XDt{#pCKYkX(@i<5g+;+0yn-d^D5z(J%2ug(rU{t!K@1b*H&`?3@KJb8!* zVg%VM7qtcWco6qD!q+}8fB-@u2UH(^ZRWXLyFPgQzkU13Ti%nU>?WA+{+VC0_PZb& zW!~B^E4iLlPp&)$z3ZA|xzEO1oW)_0n`@N<5VchR4BL#Ti1*ZW;j)A_Ju}bBt(`g(WwXq@Afny zh?ZV7B}jKA1|`Zup4Z)u>g)ctn8>HqR~X@h)r+5ZR;|3^Ja&b|ou!Od8)lr!GySvP zC}X3(d5GRJM26BxC}})rGk#OxhVC#P@+={qsUq%rdWS;k{db(r%{g_NrVrP zTu~5+w*5(h=fbzL@CU zJMkm>ZqMTFp=bfEex$3QEJmU=5Y8l4%@E3IxVmzodf((U*0H$_0=d-1%S`>YO!-S( zxOe*AE4dG!qit`1?0IZjdnfp$_jUf)Zf%axBv3jXHU@M(2zH%b7c_~hYiPlmqp%Q& zBktRo_z?N_^~Ni$tXK@>NpkCr!$A1Hg#P+%O7I`a0JNQaNp%%J36P=kn4sYdwM|Nj z6joPaN2?jQQ^+u>+sfwmki$o)z7F3eJQ;6l{YUsy>3XurXe-$Y*wptocQm6HzM?i5 zT$X&vQ9at6@O>@QEoH~tj7(}IxxaC8i~q89(U@V{bErA_1Oj0)j2$x)zhRN*9U6qNap z|Mk8B>wJ)O_?t=ojr#mAZ~ylknf)m|WTB?3z{X2FgiDBy-iepkos2#B(eX-LmQoQ( zerp1GXvyJ74Nwb^9)H!B+0y&c?yVkU%(Xr&1Hr!;lzdIPGL4IqZZNFXl40u&P=QO7 z6dquymT1Amjs1dX_O_6@@>VaKd_{K>97H@( z{O4F+T|j71a-Iv#lZ-g?B|Tl`L2n5@dympgn4M^Vdpwp4Yu4HACQNWT?G+-LdPifg z^pqWQ7Y;HSn{N38h5f5{EW?j}YGz z_2LNX8Q1F(O>7SqqmxD*$Q9-DDv%hX-Fn*3DB%NMNm)B-l2lm|JZj!m$15^Ib8uBt zoumtGJmOGWOUmfPC6;JPwAB#YNgJXlG0$f~H86BWs=#ckH>~u19e+gEXHM06&2l#o z^b-Nw2;Wg>(zf8|xA-Hlj2>`E^J?z}v0Ei4iPdXZF;wbw|BiFO89{lkvow|dBjah`cK7g!sg|x2`)rF!aQ+SDd=*8#r^)%L9?k3=@uMO)&MDwbcD~;ubWYZF|n_1W@&A`@ocpD zQ&m)d1q(ss3gzTP87+d2Ov1>1ra} zVuLW$A<6O24o^rde_6x?Sx*xc}G)p_(3At?k)$qIbf=-1>r9$3JS+J-U*P1ahat%-;*gNM(4!J{WO5=#-vbC%u8P-y%(P zTZR(I1Z$4UL%?L?y}=OLAZ>@s!fEtnbjRD4sUL;Qj&Nla6Z1LZG?#oH`}}hd_@GXK z^wMxjJY30GR(hM`WY!T242?_K_+3=m5TU!+*+ocZy_;FuUM~bwhC@E2J!YY@rM;oQ zc-`B|%1S4U_0Xx_&-t2z!y-2Pvn822_zA1_UvM<8=Byd-4eLe7bMa9pK)o|X4=X8A zr(?&+_$hu>BX;2t3+X7aZEfW84`Q`A>C#oog>#%<-c*}Z3l|D-DG-p$eE(4^xh+JD zQKQIK^gb2~wa82(ts6FwbWyq5AA~byQf0fcBGR)NDtI+@DNXwKC@lHOL;q2jc;#y{ zr5S3jYf^)PI6WgJKx5!#u%j0niiOuief^sTHA{)^?>=omTemV?Y#tI4a_ZEn4&xbC z!d)Mqm~CzN&iMo+8Rz#2cl>(@vMQ9`PW77D!5C-aMlNW=WvTY z2^&r(Q){?U+M6K+1*jA!eRY$P4Q4M>+`}Yg(im}1v71CQvyQDu>Aff*u{yVuVGP}8 zpF&Mh153Z6cBi8@HYQfsbIjbMqQa4{IhKV#lZmJFzn`SaCyHF!Pb=LJ7l*UQa{-)r zd1RYl;K^2gw+{kVP&2gh zGxeX96A`?b`eo&C$|UG^u>k+~DE>9K(wMYRQK#0MBcc<}RO#l1i}}3IY($(Ff)8bN zsI#HY3qfH-(21NFU-ad1gGkHsU<@ov!w@kKgL6f5?Kl~zFgz&uLW#gN zz^YBhWGu|6pTkv8-n)<_{qDK2z^&y+%{!Ru+dNhT!Obnk!|dX5zH>J()sAS}O&sc` zLm_)$KTsantlSTClnk)pyw6ehN=usz4;>ODDt9f(7D_+7SPu8`xPPWjCqr*FTjk@Z ztqN+f8RyJ1$lh_>RZy+z&zwZs_ z+Y$HIU>0*=-cUQW^ydd;Y zdSdAc07P0z$Pcs{DF8Y-+iMsM!4y0q_FSYnTHYJn72?Qo*i$eve8U5zqOs*cZo){+bV0;Sx^c-^PxjO zOJ{|AYOyx~h5m?l$srmtW|~ZHNQV&f?BQ8;}@C!+xt` z*QBswQv?cjm-(HQvffEc%)%Q`^HMD2X8C2(Sb5Zz3x&he(}7lZ?wFZf)9a0!*&@up zMQ?s-;O`=lOIvIiGoRk+AIRHENkk`_{_54z;yxplFn>L=6Ytsn?j;9_Iei66j|P8x ziv1#fC<39;i`~7@gLAU8yD9-eoLflGLoIFIZS?TJZ`xvbB`oYby;|weg+DM0%kZGX z0;7bEk2L&1vECFxt)IhD_(9fZw^j@xVLKf@Nl~oePGt>f7E211;T70nWyu&BZ0t25 z^ac<)6%}joiKa_7eKhQ`QhW~lmmR!xAi@T~DR1#`_*jd3kcBVi{GZRX@Md@H>`bvh zCYzCmtrw)N%t7-!s(IErzsbBfc=B-S5Ef+e3PV0{wuKTi^Jr@xW@PoD$o&cfqndDr zAHIg#c=45CoI1A0=72b={d8wH5={|q7lgT6S+x@iHBQ5aNyI7Er5@YrR_?9m#Rd;qE2f#))t7?lQN6eu3M0?VV8!~>P)~}!?CZPH+SGD0i;%BU z#e_;631}cm@+E1msu%u{c$2MZlv6n{!Y{=wE%~LlJe4EXS$DC1=Qj3_L%mRuNm)N1 zjS%JPa}VnuJywBL8?lVt)3}WeXqz%`Q-AxSBRJ}h?OXWpurqa4H|;M=-{ucD`6nfD zprz6u6Oq!v1dGUyGGUz?Ke7&9e^8JRYG!{9s8l<<@jzZBw0)Vn!$lVK4%D(>VNpZt z7J5AN-C%iU)$viRH(Rv=?JZ8Ot7XG&N%|G&1po)2rNAg!+Dn4F{i!w~!=TA4?EVfH zotL4QqIg8pP&aHraMdzMYGH!K-v^aa6)>60VsZi~hOd1B5`x9zL2Rw+v8>FzyrR

^{2x=daM*su0ZtQ0dP;=@2WwXw0w zn(}iRyJo)=MMWDXv2}5#y}Wc|Q*VA@cg5-_?y_p&S_bq;4b0U=81#gOD=*BEUg~0L z)uGeNQr3gpM>Ee`kfvJJy4y?-bGJgezlH-%rvqh8(;bHmKnQfzuXGe-W&$j&o_!+e zNYj$hZxHTBA8roygI*g-4LL#^6Hinc438*;^dO%Bh8h2?t}bw`(4vKbgwFB@~|~F^K%0<(VekqPjYT`Bc>&Nso*LC_eiyAS=8$j4_Xugg}S!$hwh-?o8(<@ zJOl7{2es~0xtk*B`v)D1<^vAx8?B55jD3B@TpunEIJ*a9i|mE>6Cs09T*1SzlOgm8 zIY3L#v!`jGmAppO8F79LdSgp#SeD1d)j!3;GtL|`LSPOS^=Cmx*o)=^4E>}}HhTH5 z{y?{F-9!XQciXM;yu~OWxYB$jpccCTXzov&c;dpS$mavlvo?RRe5r5RrsZBBW{no0 zv7@{|1#pM2#op+^R5C8izQ!ew``c%a{w6MR=ZtnSAc1$CTC|~WFE18z5QavGsDog( zPRIN6>+SpQe&CuPZ4L`aY*s>cbCL9S7N}dO{~;r(lW^{((VCa;|J0C4PZw8L%N%{U zk0OF+Klb9#y}rmOyQHmX=RA7@*p5^m)OJo0~hxaUtBF0{WJP`*y>!| zU(3`NDyFEmKoSiCiP2b8S~@$}0G1D69rC(`RrC_5A#~84rOPyfGq3-#-6$7xY&GC+ z>?iJwBB!CUw>XTz`+KdF9hM*Zc@DgG@Xcm_^x4E-&}_`P{W*iY$}sX8iWxg>Zai0k z#~vYKGi_;Ml|*=r`Fk)9&<^F*^m(9SCD&%9f4ve=wLl&M+Ow4EB6W`eBP)xRR=s{* z9+L%^KK@-n!}zC7q=>>u0O*kisC+MhORyjqP734A=kAA%J&vgI3E5Dau}%vW)k+Nc z0-8<*{qHm27<7SBV-^b4(#QugJCC|jocB4QPDeN+bSKe$*V#yce*@xWI*{biAJUVP zN3jd5;ewLw#CwrUw2g_h+2G4G^8EQbBhSxfzuLi|JgEyKv7p4Hsl)bF{EX9*g#>FC zOi7S%$+RnSx1HZ#{_FmSh>)NfSX5RvukbiMZ~UKT9HgL03|j03ixM!>QEyhq*#o57 z`t9w@8lOdRDJ-RJo!v0508oU$!vb|0e;Wi@{O{AWJ;RHpUv63kWP}XnSAcVa=-<=d zczHKc6HkLPi6Tn#$s3Z7c=ejb?`JCcm+GGwOXHd2bArD?($&7c`n##@3-A&|H2++= zQH1gJ73vnp1U3{h`+d9qgFqi3wMgYycFiDb!)}oEk#sw><1GqRZQIQTKNYvJ^Ijm!MLa2)VjHG3Wa@J78XvCJLe-|F8BMGyi%$5od>7ibsZUQtk!>9KWda zjOyrUeMwE{VsJ3^ytsc{CwGZ^FdXIXStk+z8W*kWLb^&tG%yH{q-CHZSB!T8vY; zSoAWCzBXR%Z(iGwI0ZZVbNyUe{p(y?75LPKrKwhhba?%!DRDUNu|%sFo!@npjbkBO zdnu>R^)X*G5_ZvuFEpSxk6-*}(W;cLmB5`rk3GjtI%r!5V7d}Q_a|0NhIWkl8z(s< z%N7?&jlq5XUGs9uA$&QTQ(w|XSojWJpSSKTGbR6ExvwY@V9cTOMR>3UGotrx|-Xfh3YThbP@<8#~t`(P8&PLdZj%+moDC=N266Z`4K}{0hRr-4#lC zKXdT*Yd)eBct1#6H%|s)R`o@|myz;+<3GuyI3p0%+R&S|>*GgHtWG!j1zMI>G%Z7; zb$)q!*YD+i1qFq{PZ4M8{U&XGQNGlZ`1^>ND~w(B)O|r6+R_;KY3}9p#IMrzae28d zWUEMVyuQa6P4QtYNFa5Tb!SJ6ue;~_xMdnG(J_pWhHZT_wO&?M+yRUgJRcpsyke+; z;xdT^=EsupnEU!haEp+0~Pn!nf&W`TF|0QzVZ9tBcha`*P&%?Ng65 z)TQv<$}xtWg0M!K-BBt|jV{_hMIb#{k&|OhD4~Ap$Jek zkWTHfx6C%UEiFt%^a4*5b(L6OuXY1V?Yrw85K~%uNha})q7trlvv{Gnj@TS|xmL?n znX^SRQYaAQ6Nzvjg9fB%yK+ot%1LQY4$+&MUgRX*!Qq;ptQeQryc6~BWeAcSFh|E4 zG;A>i8?d)d&*SjU{2$|-^|f7<<980THJ={SEOU>TtWIM0*az@I)|nvbJ<tvW9NkDqLlPkwOBGGQgFO|=fR#ZGA5C}a!Qc*;3USX3C+!QSq%TR-l z(s9PcdI=%>i3M05m%}l)2p)c}spq+4>gvy!e{h_ZSOC3Xuip=Pkn@`(gH}}wo0QUv ziZZ&&59TbdCV1EB$!%eZx~4yjjaczUzN&CSgoS7_D_foYS3 ze*-T%I&G13O9-nFovN$T3;aWTRe8l=AGt2fiSESOnUsA%>=KY4z2~S5cf=` zDU>TGL>olS-!=V`=0J46oFiT6r?`JbZDV8Of}>+t@5>uF3Zbsh5Que_EJy4!o2jm1peKox3Ug3`jNet?CY8K0q7j>x(@w$EF45H1@=mP+YGx zf+=~VIL+JK*dT4H;1_!ErKQLDV_L)Az?YakJagWR7L#OSTD04YZ$`h$WRkAc6VjS7@t6=Xb<0Z;An|Db_hg3T|%GI$?~p zha0azQR|3_xxN0pKiecr8q7&q1>_=?nNIlKXk*8jeJM)s8(fdFRp}y)G~8+UPy%HJ zV;}%B0ARt|4BK;UrbZtCin{R!!3wA$ibr5pRbP|P^UKr+j$M$cYyhPzcjlq_;U0ek zfJP3*oA1qr`%7yjm&&k*7m?M{p>FG8KwZGcCJ-dcD=YESsLkE2LQVgup$S|XZd3w1 zAfVtH#6Hjl4A}3EHEX;A(#{A>yG|0!%0;S7KeNA1Px1R`Il1AYkzM7HS0(}06<^YX z+g2jIFRvgTR9ygGxYPh)1oFvRBFq?83{Vb|LM#E;Mi5h^GikP#A)NV_Cd*sp9ArgT%>$w? z`%9+`D*_9scP}Ld|2zr=>a+Tb(@{X^%sgKi3DZGFI`oQX+sQr51lQ3tG0&3<ryx8Y;dp#kg6nrP_N6fb_2G=Z+7!eH-c-MSMhCD zS;=JzFvlQmv&-Uc9TyTYd$Ta!dV-&mhAWy953j=a?fRFTOAECo^MJHvT(2<#|BfB2 zY+5F##f;!zzQ~jUTbFDLI0KKz4rIp&n!WjuDaKi5vc)9vdb*3`>R7`m=S|q#RB#2| zxr^&zm3{wK=T!#O%bUN@-7G4&&R2(h9?Y{{1Ek;v=8ms2!(g}85rXdVAM0vLst#tP zmNqr%06g^I=IXN(D|cY#y)~MsDFP>i^}yqF`cn4*Iw{dp!B-l>UP^0Sttc%$!AouZ zw7k07FEtzstI+D>*RF6vZ0zkL2Kx}!)<|XAnHyg)xB{T)IL9!~2BgCPwdvQJtz&X! zGdBic%eCYb$m`p zcxi)B)b+B2@q~pR#jvr36kIEz@Ms+tNmzLT=25JKAhRbPSi#q!T_g_zUaQO-sYrV{ zJh3n5(Tx}&--9>TJ^}Eaeq-g$RDJeqH(-?Mv( zR)2feQT#L1aG%bX+{O`2?K3{9nneFkb6~LVE1T>FIt={s=L?`#OM-c<%*lZD>gLa^ z@cvil%wq_&#pa0M?fbFE&jW~ADuo~1EU_!jE3H0x?AWn&u`$pWIJS?OC30K*_e~KE z-*d1f{BuoHd>x($u$I&C%Z3=nkF|O3?!DgrP32j{+u|PwnC=`Uzeb5wi)mpz_UIF1 zqJ@o*@QH2r|BzXYA$Z|O;;#3p@Lz{W>J{Qn081=GoZ3V_n>@4%Y;NQu7_8W%g;@7l z=Hu1Sanpb37 zCpEO-#dDc4FLj0QPA98b#i6%e7Irq)u3P7<#yV?RRox!s=v36J86~L10yzlQaqD-o z&wP@XSGS7uU8=nhZ0BK_dtJdiF@&@-?B&i_bJn@_PS3{VPEc=#iRKl+Z~<5F`*nEzcX-sz6+>P;|)SeVZi)ecd5UV!^Q*dLom=Wo2?vYX;zki z^~K!)pSW%EheCmd_U$Voe^yoDZ{mzb6o-9$*>tQwH#bTc{6vYEPd4&cpD(@ZlfAdb zB@>icWs}I`(i_>eo0K>AZ8kPGe^!d}_3?`aWbWn2(5^Z?JL=o^b5-c2$A~%O!4VZb zyR!U1&vOKkU9rG@iPl2BNPY})a|{c+=IFSzWQ{~3K@-79jyGa%D4Pz+sRM=!RLs|n zy+SRsonBb?WZXa~#_qnMh+K>C9Y^GxeI1ObkSJSr>`9k$c&EnP{&>@iWESO7&)!wA~Ns8_qeuo#(g#? z(9vUd=N`(}rGpQjff>o&mpUeg4)y0ArRUs&oR|P!CC9m=RW~eRoF{>)b)X`oIfYy7 zc9dF0ft{b=zetmfGX+G5B~U_Mu(0!(dgB)e{?(q+@62u3?5rHyj%DqFkbVLdJb;(p zZ~u1b>)HVD_5b$<@PBPl`+wZw_W$Shi)X)7pePBz#{KYnzXHL6Tzyzl5t|R3dRY5# zb`sq5zv)IaFxUqErqRCEg3(FvpLBZe@JrOBkmu+=^ikj>iY5BXD1#TI-1MQ%fyXMz zMWf^yGST+A{@3mBGDqcf{`>z5zP|@veASXmPzm&#B*lJ@)VcoG-hB-m!WjJLi3f2a Z>D-yk!#|||h6b;_bpDE2<=<~U{69jsR0aS5 literal 0 HcmV?d00001 diff --git a/src/agent/agent_middleware.py b/src/agent/agent_middleware.py index 248a112..17d139f 100644 --- a/src/agent/agent_middleware.py +++ b/src/agent/agent_middleware.py @@ -1,41 +1,45 @@ from langchain.agents.middleware import wrap_tool_call from langchain.agents.structured_output import ProviderStrategy from langchain.messages import ToolMessage -from langchain.chat_models import init_chat_model -from pydantic import BaseModel, Field -from typing import Dict, Literal, Type -import os import logging -# Output schema constants (exported for tests) -class SimpleResponse(BaseModel): - message: str = Field(description="Short natural language response") - - -class TableResponse(BaseModel): - columns: list[str] = Field(description="Column headers for the table") - rows: list[list[str]] = Field(description="Table rows as lists of string values") - - -SIMPLE_RESPONSE_SCHEMA = {"type": "simple", "description": "Short natural language response"} -TABLE_RESPONSE_SCHEMA = {"type": "table", "description": "Tabular data response"} - -SCHEMA_REGISTRY: Dict[str, Type[BaseModel]] = { - "simple": SimpleResponse, - "table": TableResponse, +# 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, } - -class IntentRouter(BaseModel): - """Select the output schema required for the user's request.""" - selected_schema: Literal["simple", "table"] = Field( - description=( - "Choose 'table' for rankings, lists, comparisons, or any response that should be displayed" - " as rows/columns. Choose 'simple' for direct explanations or short answers." - ) - ) - _logger = logging.getLogger(__name__) @@ -54,49 +58,13 @@ def handle_tool_errors(request, handler): content=f"Tool error: Please check your input and try again. ({str(e)})", tool_call_id=request.tool_call["id"] ) - @wrap_tool_call -def classify_and_set_schema(request, handler): - """Classify user intent and set an output schema on the tool_call metadata. - - Uses a small Groq model to decide whether the user's query expects a - tabular (table) response or a simple natural-language response. The - decision is recorded at `request.tool_call['metadata']['output_schema']`. - """ +def set_block_response_schema(request, handler): + """Force the agent to use the BlockResponse JSON schema for structured output.""" try: - messages = getattr(request, "messages", None) - if not messages: - return handler(request) - - last = messages[-1] - role = last.get("role") if isinstance(last, dict) else getattr(last, "role", None) - if role != "human": - return handler(request) - - groq_api_key = os.getenv("GROQ_API_KEY") - selected_key = "simple" - if groq_api_key: - try: - router_llm = init_chat_model("allam-2-7b", model_provider="groq", temperature=0) - structured_router = router_llm.with_structured_output(IntentRouter) - routing_decision = structured_router.invoke(messages) - selected_key = routing_decision.selected_schema - print(f"Intent classification result: {selected_key}") - except Exception as e: - _logger.warning("Groq classification failed: %s", e) - - tool_call = getattr(request, "tool_call", None) - if isinstance(tool_call, dict): - meta = tool_call.setdefault("metadata", {}) - meta["output_schema"] = TABLE_RESPONSE_SCHEMA if selected_key == "table" else SIMPLE_RESPONSE_SCHEMA - - target_schema = SCHEMA_REGISTRY[selected_key] - request.response_format = ProviderStrategy(schema=target_schema) + request.response_format = ProviderStrategy(schema=BLOCK_RESPONSE_SCHEMA) except Exception as e: - _logger.exception("Failed to classify intent: %s", 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 806714b..7a9be5e 100644 --- a/src/agent/langchain_agent.py +++ b/src/agent/langchain_agent.py @@ -15,7 +15,8 @@ from mlflow_tools import data_access from llm.tracing import setup_langfuse, propagate_attributes from langgraph.checkpoint.memory import InMemorySaver -from agent.agent_middleware import handle_tool_errors, classify_and_set_schema +from langchain.agents.structured_output import ToolStrategy +from agent.agent_middleware import handle_tool_errors, BLOCK_RESPONSE_SCHEMA from dotenv import load_dotenv @@ -58,18 +59,22 @@ 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=[classify_and_set_schema, handle_tool_errors], + middleware=[handle_tool_errors], ) @@ -153,13 +158,38 @@ def main(): _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: try: - user_query = input("\n> ") + 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: @@ -206,4 +236,3 @@ def _interactive_loop(fuse_client_local): # No Langfuse client or observation support; just run the query result = run_query(user_query) _print_result(result) - diff --git a/src/mlflow_tools/data_access.py b/src/mlflow_tools/data_access.py index d42b2c7..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 @@ -32,7 +33,7 @@ def _decorator(f): # Keep MLflow logs quieter by default # os.environ.setdefault("MLFLOW_LOGGING_LEVEL", "WARNING") -logging.getLogger("mlflow").setLevel(logging.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')) @@ -166,7 +167,9 @@ def raw_find_best_runs_by_metric( ) -> List[Dict[str, Any]]: """Return top-k runs ordered by metric (max or min).""" - if experiment_ids.lower() in ["all","*"]: + 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: