kumar6591/data-quality-env

1"""2Chat-style AI auditor for DataQualityEnv.3 4This wrapper now behaves like a modern assistant stack:5- planner produces hypotheses and safe probe ideas6- executor runs OpenEnv tool calls7- critic normalizes/repairs the final report8- memory influences future turns9"""10 11from __future__ import annotations12 13import argparse14import json15import os16from typing import Any17 18import requests19from openai import OpenAI20 21from env.agent_memory import MemoryStore22from env.multi_agent_orchestrator import MultiAgentOrchestrator23 24API_BASE_URL = os.environ.get("API_BASE_URL", "")25MODEL_NAME = os.environ.get("MODEL_NAME", "")26API_KEY = os.environ.get("HF_TOKEN") or os.environ.get("OPENAI_API_KEY", "")27ENV_URL = os.environ.get("ENV_URL", "http://localhost:7860")28MEMORY_PATH = os.environ.get("AGENT_MEMORY_PATH", "outputs/agent_memory.json")29 30 31SYSTEM_PROMPT = """You are a data quality auditing assistant.32You can investigate data via SQL and then submit a final JSON report.33 34Return valid JSON only in this schema:35{36  "assistant_message": "short natural language reply",37  "action": {38    "action_type": "query" | "submit_report",39    "sql": "... optional when query ...",40    "report": {41      "null_issues": {"col": 0},42      "duplicate_row_count": 0,43      "schema_violations": [],44      "drifted_columns": [],45      "drift_details": {},46      "recommended_fixes": []47    }48  }49}50 51Rules:52- If user asks to inspect, use action_type=query with safe SELECT/WITH SQL.53- If enough evidence exists or user asks to finalize, use action_type=submit_report.54- Keep assistant_message concise and helpful.55"""56 57 58class ChatAuditor:59    def __init__(self, task_id: int, seed: int) -> None:60        if not API_BASE_URL or not MODEL_NAME or not API_KEY:61            raise RuntimeError("Set API_BASE_URL, MODEL_NAME, and HF_TOKEN/OPENAI_API_KEY.")62        self.client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)63        self.memory = MemoryStore(MEMORY_PATH)64        self.orchestrator = MultiAgentOrchestrator(memory=self.memory)65        self.task_id = task_id66        self.seed = seed67        self.history: list[dict[str, Any]] = []68        self.obs = self.call_env("reset", {"task_id": task_id, "seed": seed})69 70    def call_env(self, endpoint: str, payload: dict | None = None, method: str = "POST") -> dict:71        url = f"{ENV_URL}/{endpoint}"72        if method == "POST":73            r = requests.post(url, json=payload or {}, timeout=30)74        else:75            r = requests.get(url, timeout=30)76        r.raise_for_status()77        return r.json()78 79    def build_user_payload(self, user_text: str) -> str:80        view = {81            "user_request": user_text,82            "task_id": self.obs.get("task_id"),83            "task_description": self.obs.get("task_description"),84            "table_name": self.obs.get("table_name"),85            "schema": self.obs.get("schema"),86            "row_count": self.obs.get("row_count"),87            "step": self.obs.get("step"),88            "max_steps": self.obs.get("max_steps"),89            "last_query_result": (self.obs.get("last_query_result") or [])[:5],90            "last_action_error": self.obs.get("last_action_error"),91            "recent_history": self.history[-6:],92        }93        return json.dumps(view)94 95    def decide(self, user_text: str) -> dict:96        base_queries = [97            f"SELECT COUNT(*) AS n FROM {self.obs['table_name']}",98            f"SELECT * FROM {self.obs['table_name']} LIMIT 5",99        ]100        plan = self.orchestrator.build_chat_response(101            user_text=user_text,102            obs=self.obs,103            task_id=self.task_id,104            base_queries=base_queries,105            reasoning_hints=[],106        )107        return {108            "assistant_message": plan.assistant_message,109            "action": plan.action,110            "hypotheses": plan.hypotheses,111            "selected_queries": plan.selected_queries,112        }113 114    def step(self, user_text: str) -> tuple[str, dict]:115        decision = self.decide(user_text)116        assistant_message = str(decision.get("assistant_message", ""))117        action = decision.get("action", {"action_type": "query", "sql": f"SELECT COUNT(*) FROM {self.obs['table_name']}"})118 119        out = self.call_env("step", {"action": action})120        self.obs = out.get("observation", self.obs)121        reward = out.get("reward", {})122 123        self.history.append(124            {125                "user": user_text,126                "assistant_message": assistant_message,127                "action_type": action.get("action_type"),128                "reward": reward.get("value", 0.0),129                "done": reward.get("done", False),130                "selected_queries": decision.get("selected_queries", []),131            }132        )133        self.memory.save()134        return assistant_message, out135 136 137def main() -> None:138    parser = argparse.ArgumentParser(description="Chat-like AI auditor for DataQualityEnv")139    parser.add_argument("--task-id", type=int, default=1, choices=[1, 2, 3])140    parser.add_argument("--seed", type=int, default=42)141    args = parser.parse_args()142 143    auditor = ChatAuditor(task_id=args.task_id, seed=args.seed)144    print(f"Chat auditor ready for task {args.task_id}. Type 'finalize' to submit, 'exit' to quit.")145 146    while True:147        user_text = input("you> ").strip()148        if user_text.lower() in {"exit", "quit"}:149            break150        if user_text.lower() == "finalize":151            user_text = "Finalize and submit the best report now."152 153        msg, result = auditor.step(user_text)154        reward = result.get("reward", {})155        print(f"agent> {msg}")156        print(f"reward={reward.get('value', 0.0)} done={reward.get('done', False)}")157        if reward.get("done"):158            print("Episode complete.")159            break160 161 162if __name__ == "__main__":163    main()164