agent.py

# 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))