ArchCoder/llm-excel-plotter-agent
0
1import ast2import difflib3import json4import logging5import os6import re7import time8 9from dotenv import load_dotenv10 11from chart_generator import ChartGenerator12from data_processor import DataProcessor13 14load_dotenv()15 16logger = logging.getLogger(__name__)17 18# ---------------------------------------------------------------------------19# Model IDs (downloaded at Docker build, cached in HF_HOME)20# ---------------------------------------------------------------------------21QWEN_MODEL_ID = os.getenv("QWEN_MODEL_ID", "Qwen/Qwen2.5-Coder-0.5B-Instruct")22BART_MODEL_ID = os.getenv("BART_MODEL_ID", "ArchCoder/fine-tuned-bart-large")23 24# ---------------------------------------------------------------------------25# Prompt templates with few-shot examples26# ---------------------------------------------------------------------------27 28_SYSTEM_PROMPT = """\29You are a data visualization expert. Given the user request and dataset schema, \30output ONLY a valid JSON object. No explanation, no markdown fences, no extra text.31 32Required JSON keys:33 "x" : string — exact column name for the x-axis34 "y" : array — one or more exact column names for the y-axis35 "chart_type" : string — one of: line, bar, scatter, pie, histogram, box, area36 "color" : string or null — optional CSS color like "red", "#4f8cff"37 38Rules:39- Use ONLY column names from the schema. Never invent names.40- For pie charts: y must contain exactly one column.41- For histogram/box: x may equal the first element of y.42- Default to "line" if chart type is ambiguous.43 44### Examples45 46Example 1:47Schema: Year (integer), Sales (float), Profit (float)48User: "plot sales over the years with a red line"49Output: {"x": "Year", "y": ["Sales"], "chart_type": "line", "color": "red"}50 51Example 2:52Schema: Month (string), Revenue (float), Expenses (float)53User: "bar chart comparing revenue and expenses by month"54Output: {"x": "Month", "y": ["Revenue", "Expenses"], "chart_type": "bar", "color": null}55 56Example 3:57Schema: Category (string), Count (integer)58User: "pie chart of count by category"59Output: {"x": "Category", "y": ["Count"], "chart_type": "pie", "color": null}60 61Example 4:62Schema: Date (string), Temperature (float), Humidity (float)63User: "scatter plot of temperature vs humidity in blue"64Output: {"x": "Temperature", "y": ["Humidity"], "chart_type": "scatter", "color": "blue"}65 66Example 5:67Schema: Year (integer), Sales (float), Employee expense (float), Marketing expense (float)68User: "show me an area chart of sales and marketing expense over years"69Output: {"x": "Year", "y": ["Sales", "Marketing expense"], "chart_type": "area", "color": null}70"""71 72 73def _user_message(query: str, columns: list, dtypes: dict, sample_rows: list) -> str:74 schema = "\n".join(f" - {c} ({dtypes.get(c, 'unknown')})" for c in columns)75 samples = "".join(f" {json.dumps(r)}\n" for r in sample_rows[:3])76 return (77 f"Schema:\n{schema}\n\n"78 f"Sample rows:\n{samples}\n"79 f"User: \"{query}\"\n"80 f"Output:"81 )82 83 84# ---------------------------------------------------------------------------85# Output parsing & validation86# ---------------------------------------------------------------------------87 88def _parse_output(text: str):89 text = text.strip()90 if "```" in text:91 for part in text.split("```"):92 part = part.strip().lstrip("json").strip()93 if part.startswith("{"):94 text = part95 break96 try:97 return json.loads(text)98 except json.JSONDecodeError:99 pass100 try:101 return ast.literal_eval(text)102 except (SyntaxError, ValueError):103 pass104 return None105 106 107def _validate(args: dict, columns: list):108 if not isinstance(args, dict):109 return None110 if not all(k in args for k in ("x", "y", "chart_type")):111 return None112 if isinstance(args["y"], str):113 args["y"] = [args["y"]]114 valid = {"line", "bar", "scatter", "pie", "histogram", "box", "area"}115 if args["chart_type"] not in valid:116 args["chart_type"] = "line"117 if args["x"] not in columns:118 return None119 if not all(c in columns for c in args["y"]):120 return None121 return args122 123 124def _pick_chart_type(query: str) -> str:125 lowered = query.lower()126 aliases = {127 "scatter": ["scatter", "scatterplot"],128 "bar": ["bar", "column"],129 "pie": ["pie", "donut"],130 "histogram": ["histogram", "distribution"],131 "box": ["box", "boxplot"],132 "area": ["area"],133 "line": ["line", "trend", "over time", "over the years"],134 }135 for chart_type, keywords in aliases.items():136 if any(keyword in lowered for keyword in keywords):137 return chart_type138 return "line"139 140 141def _pick_color(query: str):142 lowered = query.lower()143 colors = [144 "red", "blue", "green", "yellow", "orange", "purple", "pink",145 "black", "white", "gray", "grey", "cyan", "teal", "indigo",146 ]147 for color in colors:148 if re.search(rf"\b{re.escape(color)}\b", lowered):149 return color150 return None151 152 153def _pick_columns(query: str, columns: list, dtypes: dict):154 lowered = query.lower()155 query_tokens = re.findall(r"[a-zA-Z0-9_]+", lowered)156 157 def score_column(column: str) -> float:158 col_lower = column.lower()159 score = 0.0160 if col_lower in lowered:161 score += 10.0162 for token in query_tokens:163 if token and token in col_lower:164 score += 2.0165 score += difflib.SequenceMatcher(None, lowered, col_lower).ratio()166 return score167 168 sorted_columns = sorted(columns, key=score_column, reverse=True)169 numeric_columns = [col for col in columns if dtypes.get(col) in {"integer", "float"}]170 temporal_columns = [col for col in columns if dtypes.get(col) == "datetime"]171 year_like = [col for col in columns if "year" in col.lower() or "date" in col.lower() or "month" in col.lower()]172 173 x_col = None174 for candidate in year_like + temporal_columns + sorted_columns:175 if candidate in columns:176 x_col = candidate177 break178 if x_col is None and columns:179 x_col = columns[0]180 181 y_candidates = [col for col in sorted_columns if col != x_col and col in numeric_columns]182 if not y_candidates:183 y_candidates = [col for col in numeric_columns if col != x_col]184 if not y_candidates:185 y_candidates = [col for col in columns if col != x_col]186 187 return x_col, y_candidates[:1]188 189 190def _heuristic_plot_args(query: str, columns: list, dtypes: dict) -> dict:191 x_col, y_cols = _pick_columns(query, columns, dtypes)192 if not x_col:193 x_col = "Year"194 if not y_cols:195 fallback_y = next((col for col in columns if col != x_col), columns[:1])196 y_cols = list(fallback_y) if isinstance(fallback_y, tuple) else fallback_y197 if isinstance(y_cols, str):198 y_cols = [y_cols]199 return {200 "x": x_col,201 "y": y_cols,202 "chart_type": _pick_chart_type(query),203 "color": _pick_color(query),204 }205 206 207# ---------------------------------------------------------------------------208# Agent209# ---------------------------------------------------------------------------210 211class LLM_Agent:212 def __init__(self, data_path=None):213 logger.info("Initializing LLM_Agent")214 self.data_processor = DataProcessor(data_path)215 self.chart_generator = ChartGenerator(self.data_processor.data)216 self._bart_tokenizer = None217 self._bart_model = None218 self._qwen_tokenizer = None219 self._qwen_model = None220 221 # -- model runners -------------------------------------------------------222 223 def _run_qwen(self, user_msg: str) -> str:224 """Qwen2.5-Coder-0.5B-Instruct — fast structured-JSON generation."""225 if self._qwen_model is None:226 from transformers import AutoModelForCausalLM, AutoTokenizer227 logger.info(f"Loading Qwen model: {QWEN_MODEL_ID}")228 self._qwen_tokenizer = AutoTokenizer.from_pretrained(QWEN_MODEL_ID)229 self._qwen_model = AutoModelForCausalLM.from_pretrained(QWEN_MODEL_ID)230 logger.info("Qwen model loaded.")231 messages = [232 {"role": "system", "content": _SYSTEM_PROMPT},233 {"role": "user", "content": user_msg},234 ]235 text = self._qwen_tokenizer.apply_chat_template(236 messages, tokenize=False, add_generation_prompt=True237 )238 inputs = self._qwen_tokenizer(text, return_tensors="pt")239 outputs = self._qwen_model.generate(240 **inputs, max_new_tokens=256, temperature=0.1, do_sample=True241 )242 return self._qwen_tokenizer.decode(243 outputs[0][inputs.input_ids.shape[-1]:], skip_special_tokens=True244 )245 246 def _run_gemini(self, user_msg: str) -> str:247 import google.generativeai as genai248 api_key = os.getenv("GEMINI_API_KEY")249 if not api_key:250 raise ValueError("GEMINI_API_KEY is not set")251 genai.configure(api_key=api_key)252 model = genai.GenerativeModel(253 "gemini-2.0-flash",254 system_instruction=_SYSTEM_PROMPT,255 )256 return model.generate_content(user_msg).text257 258 def _run_grok(self, user_msg: str) -> str:259 from openai import OpenAI260 api_key = os.getenv("GROK_API_KEY")261 if not api_key:262 raise ValueError("GROK_API_KEY is not set")263 client = OpenAI(api_key=api_key, base_url="https://api.x.ai/v1")264 resp = client.chat.completions.create(265 model="grok-3-mini",266 messages=[267 {"role": "system", "content": _SYSTEM_PROMPT},268 {"role": "user", "content": user_msg},269 ],270 max_tokens=256,271 temperature=0.1,272 )273 return resp.choices[0].message.content274 275 def _run_bart(self, query: str) -> str:276 """ArchCoder/fine-tuned-bart-large — lightweight Seq2Seq fallback."""277 if self._bart_model is None:278 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer279 logger.info(f"Loading BART model: {BART_MODEL_ID}")280 self._bart_tokenizer = AutoTokenizer.from_pretrained(BART_MODEL_ID)281 self._bart_model = AutoModelForSeq2SeqLM.from_pretrained(BART_MODEL_ID)282 logger.info("BART model loaded.")283 inputs = self._bart_tokenizer(284 query, return_tensors="pt", max_length=512, truncation=True285 )286 outputs = self._bart_model.generate(**inputs, max_length=100)287 return self._bart_tokenizer.decode(outputs[0], skip_special_tokens=True)288 289 # -- main entry point ----------------------------------------------------290 291 def process_request(self, data: dict) -> dict:292 t0 = time.time()293 query = data.get("query", "")294 data_path = data.get("file_path")295 model = data.get("model", "qwen")296 297 if data_path and os.path.exists(data_path):298 self.data_processor = DataProcessor(data_path)299 self.chart_generator = ChartGenerator(self.data_processor.data)300 301 columns = self.data_processor.get_columns()302 dtypes = self.data_processor.get_dtypes()303 sample_rows = self.data_processor.preview(3)304 305 default_args = {306 "x": columns[0] if columns else "Year",307 "y": [columns[1]] if len(columns) > 1 else ["Sales"],308 "chart_type": "line",309 }310 311 raw_text = ""312 plot_args = None313 try:314 user_msg = _user_message(query, columns, dtypes, sample_rows)315 if model == "gemini": raw_text = self._run_gemini(user_msg)316 elif model == "grok": raw_text = self._run_grok(user_msg)317 elif model == "bart": raw_text = self._run_bart(query)318 elif model == "qwen":319 try:320 raw_text = self._run_qwen(user_msg)321 except Exception as qwen_exc:322 logger.warning(f"Qwen failed, falling back to BART: {qwen_exc}")323 raw_text = self._run_bart(query)324 else:325 raw_text = self._run_qwen(user_msg)326 327 logger.info(f"LLM [{model}] output: {raw_text}")328 parsed = _parse_output(raw_text)329 plot_args = _validate(parsed, columns) if parsed else None330 except Exception as exc:331 logger.error(f"LLM error [{model}]: {exc}")332 raw_text = str(exc)333 334 if not plot_args:335 logger.warning("Falling back to heuristic plot args")336 plot_args = _validate(_heuristic_plot_args(query, columns, dtypes), columns) or default_args337 338 try:339 chart_result = self.chart_generator.generate_chart(plot_args)340 chart_path = chart_result["chart_path"]341 chart_spec = chart_result["chart_spec"]342 except Exception as exc:343 logger.error(f"Chart generation error: {exc}")344 return {345 "response": f"Chart generation failed: {exc}",346 "chart_path": "",347 "chart_spec": None,348 "verified": False,349 "plot_args": plot_args,350 }351 352 logger.info(f"Request processed in {time.time() - t0:.2f}s")353 return {354 "response": json.dumps(plot_args),355 "chart_path": chart_path,356 "chart_spec": chart_spec,357 "verified": True,358 "plot_args": plot_args,359 }360