# typing, Hugging Face, configuration, and ML context dependencies from __future__ import annotations import json import re from typing import Any from huggingface_hub import InferenceClient from config import DEFAULT_MAX_TOKENS, HF_MODEL_ID, HF_PROVIDER, HF_TOKEN, MAX_AGENT_STEPS from ml_engine import MLContext # Define the system instructions that constrain the ML agent's analysis and final response format SYSTEM_PROMPT = '''You are a senior Machine Learning Engineering Agent. Use tools before making claims about the uploaded dataset or model performance. Frame the supervised-learning problem, inspect target behavior, recommend appropriate algorithms, train or compare candidates when useful, and produce reproducible scikit-learn pipeline code grounded in the current dataset. Do not invent metrics. Only cite metrics returned by tools. Distinguish model evaluation from business impact. Flag leakage, class imbalance, target quality, high-cardinality categoricals, tiny samples, or metric-selection risks when relevant. Final answer must contain exactly these sections: ## ML assessment ## Dataset and target findings ## Model strategy ## Evaluation ## Risks and next steps ## Recommended Pipeline ```python ... ``` ''' # declare the function-tool schemas exposed to the Hugging Face chat-completion agent TOOL_SCHEMAS = [ # register the dataset-inspection tool used to summarize dataset structure and sample characteristics { "type": "function", "function": { "name": "inspect_dataset", "description": "Inspect dataset rows, columns, types, nulls, cardinality, and examples.", "parameters": {"type": "object", "properties": {}}, }, }, # register the modeling-setup tool used to infer the task type and suitable candidate algorithms { "type": "function", "function": { "name": "recommend_modeling_setup", "description": "Infer or validate the ML problem type for the selected target and recommend candidate algorithms.", "parameters": {"type": "object", "properties": {}}, }, }, # register the candidate-training tool used to fit and evaluate the selected algorithm { "type": "function", "function": { "name": "train_candidate_model", "description": "Train and evaluate the currently selected candidate algorithm on a holdout split.", "parameters": {"type": "object", "properties": {}}, }, }, # register the comparison tool used to benchmark a compact set of baseline algorithms { "type": "function", "function": { "name": "compare_algorithms", "description": "Run a compact deterministic comparison of suitable baseline algorithms.", "parameters": {"type": "object", "properties": {}}, }, }, # pipeline-generation tool used to produce reproducible scikit-learn code { "type": "function", "function": { "name": "generate_pipeline_code", "description": "Generate reproducible scikit-learn pipeline code for the selected target and algorithm.", "parameters": {"type": "object", "properties": {}}, }, }, ] # Extract the first fenced Python code block from a model response, returning an empty string when absent def _extract_code(text: str, language: str = "python") -> str: # Accept common Python fence labels when searching the response text aliases = [language, "py", "python"] match = re.search(rf"```(?:{'|'.join(re.escape(x) for x in aliases)})\s*(.*?)```", text or "", re.I | re.S) return match.group(1).strip() if match else "" # Normalize a model tool-call object into the dictionary structure expected by the chat message history def _serialize_tool_call(call: Any) -> dict[str, Any]: # Read the raw tool arguments so string-encoded JSON can be normalized before serialization args = call.function.arguments # Parse JSON argument strings into dictionaries while safely falling back on malformed input if isinstance(args, str): try: args = json.loads(args) except json.JSONDecodeError: args = {} # Return a stable function-call payload with an ID, type, name, and JSON-encoded arguments. return { "id": getattr(call, "id", None) or f"call_{call.function.name}", "type": "function", "function": {"name": call.function.name, "arguments": json.dumps(args or {})}, } # Coordinate ML-context tools with an optional Hugging Face model to produce a grounded agent response class MachineLearningAgent: # Initialize the agent with dataset context, modeling choices, evaluation settings, and generation controls def __init__( self, context: MLContext, target: str, algorithm: str = "Auto", problem_type: str = "Auto", test_size: float = 0.2, temperature: float = 0.15, max_tokens: int = DEFAULT_MAX_TOKENS, ): # Store normalized runtime settings and initialize the trace used to record tool execution history self.context = context self.target = target self.algorithm = algorithm or "Auto" self.problem_type = problem_type or "Auto" self.test_size = float(test_size) self.temperature = float(temperature) self.max_tokens = int(max_tokens) self.trace: list[dict[str, Any]] = [] # Dispatch a named agent tool to the matching MLContext operation and record its result def _tool(self, name: str, arguments: dict[str, Any] | None = None) -> dict[str, Any]: # Normalize missing tool arguments to an empty dictionary before dispatch. arguments = arguments or {} # Route each supported tool name to the corresponding deterministic ML operation if name == "inspect_dataset": result = self.context.compact_profile() elif name == "recommend_modeling_setup": result = self.context.modeling_recommendation(self.target, self.problem_type) elif name == "train_candidate_model": result = self.context.train_candidate( self.target, self.algorithm, self.problem_type, self.test_size, ) elif name == "compare_algorithms": result = self.context.compare_algorithms(self.target, self.problem_type, self.test_size) elif name == "generate_pipeline_code": # Generate reproducible pipeline code and package it with the selected target and algorithm metadata code = self.context.generate_pipeline_code(self.target, self.algorithm, self.problem_type, self.test_size) result = {"target": self.target, "algorithm": self.algorithm, "pipeline_code": code} # Return a structured error payload when the model requests an unsupported tool name else: result = {"error": f"Unknown tool: {name}"} # Append every tool invocation and result to the trace for downstream inspection and reproducibility self.trace.append({"tool": name, "arguments": arguments, "result": result}) return result # Build a deterministic response from local ML tooling when model synthesis is unavailable or fails def _fallback(self, task: str, error: str | None = None): # Collect dataset profiling, modeling recommendations, training results, and pipeline code for the fallback profile = self._tool("inspect_dataset") setup = self._tool("recommend_modeling_setup") train = self._tool("train_candidate_model") code = self.context.generate_pipeline_code(self.target, self.algorithm, self.problem_type, self.test_size) self.trace.append({"tool": "generate_pipeline_code", "arguments": {}, "result": {"pipeline_code": code}}) # Explain whether fallback mode was triggered by a missing token or by a failed model request if not HF_TOKEN: note = "HF model synthesis is disabled because `HF_TOKEN` is not configured, so deterministic ML tooling produced this response." else: note = f"The HF model request failed, so deterministic ML tooling produced this response. Error: `{error}`" # Format returned training metrics and assemble the required structured fallback answer metrics = "\n".join(f"- **{k}**: `{v}`" for k, v in train.get("metrics", {}).items()) or "- No metrics were returned." answer = f'''## ML assessment {note} Task: {task} Loaded `{profile['source']}` with **{profile['rows']:,} rows** and **{profile['columns']:,} columns**. The selected target is `{self.target}` and the resolved problem type is **{setup['problem_type']}**. ## Dataset and target findings - Target unique values: **{setup.get('target_unique', 'n/a')}** - Target nulls: **{setup.get('target_nulls', 'n/a')}** - Candidate algorithms: {', '.join(setup.get('recommended_algorithms', []))} ## Model strategy Trained **{train.get('algorithm', self.algorithm)}** using a reusable preprocessing + estimator pipeline with train/test separation. ## Evaluation {metrics} ## Risks and next steps - Check for target leakage and time-dependent leakage before production use. - Match the primary metric to the business cost of false positives, false negatives, or prediction error. - Add cross-validation and hyperparameter optimization after establishing the baseline. - Validate drift, calibration, fairness, and operational latency where relevant. ## Recommended Pipeline ```python {code} ``` ''' # Return the synthesized answer, reusable pipeline code, and complete execution trace together return answer, code, self.trace # Execute the agent workflow, using deterministic fallback immediately when no Hugging Face token is configured def run(self, task: str): if not HF_TOKEN: return self._fallback(task) # HF inference client and seed the conversation with system and task context client = InferenceClient(token=HF_TOKEN, provider=HF_PROVIDER, timeout=120) messages = [ {"role": "system", "content": SYSTEM_PROMPT}, { "role": "user", "content": ( f"Task: {task}\n" f"Selected target: {self.target}\n" f"Problem type setting: {self.problem_type}\n" f"Selected algorithm: {self.algorithm}\n" f"Test size: {self.test_size}\n" "Use tools first." ), }, ] # Run bounded model/tool interaction steps and fall back safely if any exception escapes the workflow try: # Track the final natural-language response once the model stops requesting tools final = "" # Limit iterative model/tool exchanges to the configured maximum number of agent steps for _ in range(MAX_AGENT_STEPS): # Request the next assistant message with tool definitions and the current conversation state response = client.chat_completion( model=HF_MODEL_ID, messages=messages, tools=TOOL_SCHEMAS, tool_choice="auto", temperature=self.temperature, max_tokens=self.max_tokens, ) # Inspect the returned assistant message for function calls or a completed final response message = response.choices[0].message calls = getattr(message, "tool_calls", None) or [] # Stop the loop when the model returns content without requesting any additional tools if not calls: final = (message.content or "").strip() break # Preserve the assistant's tool-call message in history before executing requested functions messages.append( { "role": "assistant", "content": message.content or "", "tool_calls": [_serialize_tool_call(call) for call in calls], } ) # Execute each requested tool, normalize its arguments, and append its result as a tool message for call in calls: args = call.function.arguments # Parse string-encoded tool arguments while tolerating malformed JSON from the model if isinstance(args, str): try: args = json.loads(args) except json.JSONDecodeError: args = {} # Dispatch the tool through the local ML context using the normalized argument dictionary result = self._tool(call.function.name, args or {}) # Add the serialized tool output to the conversation so the model can use it on the next step messages.append( { "role": "tool", "tool_call_id": getattr(call, "id", None) or f"call_{call.function.name}", "name": call.function.name, "content": json.dumps(result, default=str)[:30000], } ) # Force a final synthesis request if the bounded loop ends without producing final content if not final: response = client.chat_completion( model=HF_MODEL_ID, messages=messages + [{"role": "user", "content": "Synthesize the final response now. Do not call more tools."}], temperature=self.temperature, max_tokens=self.max_tokens, ) final = (response.choices[0].message.content or "").strip() # Reuse Python code from the final response when present, otherwise generate pipeline code deterministically code = _extract_code(final) or self.context.generate_pipeline_code( self.target, self.algorithm, self.problem_type, self.test_size ) # Return the final response, extracted or generated pipeline code, and accumulated tool trace return final, code, self.trace # Recover from inference or orchestration errors by returning the deterministic fallback response except Exception as exc: return self._fallback(task, str(exc))