Spaces:
Running on Zero
Running on Zero
Fix: generic schema tools always registered + live state (fixes Level-3 404 on non-employee DBs)
1140593 verified | from fastapi import FastAPI, HTTPException | |
| from pydantic import BaseModel | |
| from typing import Any, Dict, Optional, List | |
| import asyncio | |
| import importlib | |
| import uuid | |
| import os | |
| import re | |
| # Attempt to import agent workflow for schema/inspector & utilities (optional) | |
| try: | |
| import Agent.agentic_workflow as aw | |
| print(f"[MCP STARTUP] Successfully imported Agent.agentic_workflow") | |
| except Exception as e: | |
| print(f"[MCP STARTUP] Failed to import Agent.agentic_workflow: {e}") | |
| aw = None | |
| app = FastAPI(title="Minimal MCP Server (dev)") | |
| # Simple tool registry | |
| TOOLS: Dict[str, Any] = {} | |
| class CompleteRequest(BaseModel): | |
| prompt: str | |
| context: Optional[Dict[str, Any]] = None | |
| timeout: Optional[float] = 30.0 | |
| def register_tool(name: str, fn: Any): | |
| TOOLS[name] = fn | |
| async def _call_tool(name: str, args: Dict[str, Any]) -> Any: | |
| if name not in TOOLS: | |
| raise KeyError(f"Tool '{name}' not registered") | |
| fn = TOOLS[name] | |
| if asyncio.iscoroutinefunction(fn): | |
| return await fn(**(args or {})) | |
| else: | |
| # run sync in threadpool | |
| loop = asyncio.get_running_loop() | |
| return await loop.run_in_executor(None, lambda: fn(**(args or {}))) | |
| async def _startup(): | |
| # Register built-in/mock tools lazily to avoid importing heavy modules at import time | |
| # Try to register a real document scanner implementation; otherwise register a mock | |
| ds_registered = None | |
| try: | |
| from .tools import document_scanner as ds | |
| register_tool("document_scanner.process", ds.process_image) | |
| ds_registered = f"Agent.tools.document_scanner.process_image (relative import)" | |
| except Exception: | |
| # If relative import fails (running as script), try absolute import | |
| try: | |
| import Agent.tools.document_scanner as ds2 | |
| register_tool("document_scanner.process", ds2.process_image) | |
| ds_registered = f"Agent.tools.document_scanner.process_image (absolute import)" | |
| except Exception: | |
| # Register a fallback mock tool | |
| def _mock_doc(**kwargs): | |
| # Accept arbitrary keyword args (image_url, image_bytes, etc.) to be tolerant | |
| return {"doc_id": str(uuid.uuid4()), "text": "(mock OCR output)", "metadata": {"pages": 1, "confidence": 0.0}, "received_args": kwargs} | |
| register_tool("document_scanner.process", _mock_doc) | |
| ds_registered = "mock(document_scanner.process)" | |
| # Log which document scanner implementation was registered for easier debugging | |
| try: | |
| print(f"[MCP STARTUP] Registered document_scanner.process -> {ds_registered}") | |
| except Exception: | |
| pass | |
| # Optional Groq Vision health check in cloud mode | |
| try: | |
| from .tools import groq_vision_health | |
| hc = groq_vision_health.health_check() | |
| if not hc.get("ok", True): | |
| print(f"[MCP STARTUP] Groq Vision health warning: {hc.get('error')} (model={hc.get('model')})") | |
| else: | |
| print(f"[MCP STARTUP] Groq Vision health: ok (model={hc.get('model')})") | |
| except Exception as e: | |
| print(f"[MCP STARTUP] Groq Vision health check skipped: {e}") | |
| # Register Granite Vision direct visual Q&A tool | |
| gv_registered = None | |
| try: | |
| from .tools import granite_vision as gv | |
| register_tool("granite_vision.qa", gv.qa) | |
| gv_registered = "Agent.tools.granite_vision.qa (relative import)" | |
| except Exception as e1: | |
| print(f"[MCP STARTUP] granite_vision relative import FAILED: {type(e1).__name__}: {e1}") | |
| try: | |
| import Agent.tools.granite_vision as gv2 | |
| register_tool("granite_vision.qa", gv2.qa) | |
| gv_registered = "Agent.tools.granite_vision.qa (absolute import)" | |
| except Exception as e2: | |
| print(f"[MCP STARTUP] granite_vision absolute import FAILED: {type(e2).__name__}: {e2}") | |
| def _mock_gv(**kwargs): | |
| return {"answer": "(granite_vision unavailable - mock)", "model": "mock", "received_args": kwargs} | |
| register_tool("granite_vision.qa", _mock_gv) | |
| gv_registered = "mock(granite_vision.qa)" | |
| print(f"[MCP STARTUP] Registered granite_vision.qa -> {gv_registered}") | |
| # Register OCR QA composite tool: runs OCR then runs the agent OCR QA helper | |
| try: | |
| # Use the global aw variable instead of re-importing | |
| if not aw: | |
| raise ImportError("Agent workflow not available") | |
| async def _ocr_qa_tool(image_url: Optional[str] = None, image_bytes: Optional[bytes] = None, question: Optional[str] = None, **kwargs): | |
| # Call the document scanner tool first (use the registered one to allow mocks) | |
| ds_fn = TOOLS.get("document_scanner.process") | |
| if not ds_fn: | |
| raise RuntimeError("document_scanner.process not registered") | |
| loop = asyncio.get_running_loop() | |
| try: | |
| if asyncio.iscoroutinefunction(ds_fn): | |
| tool_res = await ds_fn(image_url=image_url, image_bytes=image_bytes, options=kwargs) | |
| else: | |
| tool_res = await loop.run_in_executor(None, lambda: ds_fn(image_url=image_url, image_bytes=image_bytes, options=kwargs)) | |
| except Exception as e: | |
| return {"error": f"OCR tool failed: {e}"} | |
| doc_id = tool_res.get("doc_id") if isinstance(tool_res, dict) else None | |
| ocr_text = (tool_res.get("text") if isinstance(tool_res, dict) else str(tool_res)) or "" | |
| # Persist OCR text so downstream callers can retrieve it via doc_id | |
| try: | |
| if doc_id and ocr_text: | |
| docs_dir = os.path.join("chat_store", "docs") | |
| os.makedirs(docs_dir, exist_ok=True) | |
| doc_path = os.path.join(docs_dir, f"{doc_id}.txt") | |
| # Write file (overwrite if exists) | |
| with open(doc_path, "w", encoding="utf-8") as df: | |
| df.write(ocr_text) | |
| # Also add a short history entry if agent chat store available | |
| try: | |
| if hasattr(aw, 'chat_store_private') and hasattr(aw, 'ChatMessage') and hasattr(aw, 'MessageRole'): | |
| user_msg = aw.ChatMessage(role=aw.MessageRole.USER, content=f"Uploaded document {doc_id}") | |
| assistant_msg = aw.ChatMessage(role=aw.MessageRole.ASSISTANT, content=f"OCR stored: {ocr_text[:400]}") | |
| aw.chat_store_private.add_message(key="conversation", message=user_msg) | |
| aw.chat_store_private.add_message(key="conversation", message=assistant_msg) | |
| try: | |
| aw.chat_store_private.persist(str(aw.private_store_path)) | |
| except Exception: | |
| pass | |
| except Exception: | |
| pass | |
| except Exception: | |
| # Non-fatal: continue even if persistence fails | |
| pass | |
| # If a question was provided, answer it using the agent helper (run in executor to avoid blocking) | |
| if question: | |
| try: | |
| # Use executor if helper is sync to avoid blocking event loop | |
| if callable(getattr(aw, 'ocr_agent_qa', None)): | |
| answer = await loop.run_in_executor(None, lambda: aw.ocr_agent_qa(question, doc_id)) | |
| return {"doc_id": doc_id, "answer": answer} | |
| else: | |
| return {"doc_id": doc_id, "answer": "ocr_agent_qa not available"} | |
| except Exception as e: | |
| return {"error": f"ocr_agent_qa failed: {e}", "tool_result": tool_res} | |
| # No question: return the OCR result and doc_id | |
| return {"doc_id": doc_id, "ocr_text": ocr_text} | |
| register_tool("ocr_qa", _ocr_qa_tool) | |
| print(f"[MCP STARTUP] Registered ocr_qa -> real _ocr_qa_tool (agent available)") | |
| except Exception: | |
| # If agent import fails, register a tolerant mock tool | |
| def _ocr_qa_mock(**kwargs): | |
| return {"doc_id": str(uuid.uuid4()), "answer": "(ocr_qa unavailable - mock response)", "received_args": kwargs} | |
| register_tool("ocr_qa", _ocr_qa_mock) | |
| print(f"[MCP STARTUP] Registered ocr_qa -> mock(ocr_qa)") | |
| # ------------------------------------------------------------------ | |
| # Deterministic schema-aware SQL helper tools (agentic compliant) | |
| # ------------------------------------------------------------------ | |
| ENABLE_FAST_SQL = os.environ.get("ENABLE_MCP_FAST_TOOLS", "1") in ("1","true","True") | |
| if not ENABLE_FAST_SQL: | |
| print("[MCP STARTUP] Fast SQL tools disabled via ENABLE_MCP_FAST_TOOLS=0") | |
| return | |
| # Initialize database connection if not already done | |
| if not aw: | |
| print("[MCP STARTUP][FAST] Agent workflow not available; skipping fast SQL tool registration.") | |
| return | |
| if not getattr(aw, 'inspector', None) or not getattr(aw, 'available_tables', None): | |
| print("[MCP STARTUP][FAST] Database not initialized; attempting to initialize...") | |
| try: | |
| # Initialize database connection - use await since we're in async context. | |
| # Priority chain mirrors worker startup: 1) env DB 2) active schema | |
| # profile from the registry 3) bundled demo employee SQLite file. | |
| db_uri = os.environ.get("DB_CONNECTION_URL") | |
| if not db_uri: | |
| demo_db_path = os.path.join( | |
| os.path.dirname(os.path.dirname(os.path.abspath(__file__))), | |
| "employee-sample-database", "sqlite", "dataset_small", "employee.db", | |
| ) | |
| try: | |
| from .schema_profile import profile_registry as _reg | |
| _active = _reg.active() or {} | |
| db_uri = _active.get("db_url") | |
| except Exception: | |
| db_uri = None | |
| if not db_uri and os.path.isfile(demo_db_path): | |
| db_uri = f"sqlite:///{demo_db_path.replace(os.sep, '/')}" | |
| print("[MCP STARTUP][FAST] Using bundled demo employee database.") | |
| if not db_uri: | |
| print("[MCP STARTUP][FAST] No DB env, active profile, or bundled demo found; skipping fast SQL tool registration.") | |
| return | |
| db_result = await aw.initialize_database(db_uri) | |
| if db_result.get("tables"): | |
| print(f"[MCP STARTUP][FAST] Database initialized with {len(db_result['tables'])} tables") | |
| else: | |
| print("[MCP STARTUP][FAST] Database initialization failed; skipping fast SQL tool registration.") | |
| return | |
| except Exception as db_err: | |
| print(f"[MCP STARTUP][FAST] Database initialization error: {db_err}; skipping fast SQL tool registration.") | |
| return | |
| inspector = aw.inspector | |
| available_tables = aw.available_tables | |
| def _live_tables(): | |
| """Live table list — aw.available_tables is REBOUND on every profile | |
| switch, so a boot-time snapshot would describe the wrong database.""" | |
| return list(getattr(aw, 'available_tables', None) or available_tables or []) | |
| def _colnames(table: str) -> List[str]: | |
| try: | |
| # Live lookup: aw.inspector is rebound on every profile switch. | |
| insp = getattr(aw, 'inspector', None) or inspector | |
| return [c['name'] for c in insp.get_columns(table_name=table)] | |
| except Exception: | |
| return [] | |
| def _has_tables(*tables: str) -> bool: | |
| at_lower = {t.lower() for t in _live_tables()} | |
| return all(t.lower() in at_lower for t in tables) | |
| # Simple validation & sanitization re-use | |
| validate_sql_query = getattr(aw, 'validate_sql_query', lambda q: {"valid": True}) | |
| fix_common_sql_errors = getattr(aw, 'fix_common_sql_errors', lambda q: q) | |
| async def _verify_sql(sql: str, question: str, tag: str) -> str: | |
| """Optional light LLM verification (short prompt) if LLM available.""" | |
| if not aw or not getattr(aw, 'Settings', None) or not getattr(aw.Settings, 'llm', None): | |
| return sql | |
| try: | |
| _schema_block = "" | |
| try: | |
| if aw.ACTIVE_SCHEMA_PROFILE.is_employee_db: | |
| _schema_block = aw.EMPLOYEE_DB_SCHEMA | |
| else: | |
| _schema_block = aw.ACTIVE_SCHEMA_PROFILE.build_digest(max_tables=10) | |
| except Exception: | |
| _schema_block = "" | |
| prompt = ( | |
| (f"{_schema_block}\n\n" if _schema_block else "") | |
| + "You are an SQL verifier." | |
| + f" Tag={tag}." | |
| + " If the candidate SQL is valid and uses only existing tables/columns keep it." | |
| " If SMALL fixes (alias, LIMIT 100, semicolon) needed, fix them. Never invent tables/columns." | |
| f"\nQuestion: {question}\nCandidate SQL:\n{sql}\nReturn ONLY final SQL." ) | |
| resp = aw.Settings.llm.complete(prompt=prompt) | |
| txt = getattr(resp, 'text', str(resp)).strip() | |
| if '```' in txt: | |
| parts = re.split(r"```(?:sql)?", txt, flags=re.IGNORECASE) | |
| txt = '\n'.join(p for p in parts if 'select' in p.lower() or 'with' in p.lower()) | |
| if txt.lower().startswith('sql'): | |
| txt = txt[3:].strip() | |
| txt = fix_common_sql_errors(txt) | |
| v = validate_sql_query(txt) | |
| if v.get('valid'): | |
| return txt | |
| except Exception as ve: | |
| print(f"[FAST][VERIFY][WARN] Verification failed: {ve}; using original SQL") | |
| return sql | |
| def _build_sql(select_cols: List[str], base_table: str, joins: Optional[List[str]] = None, | |
| where: Optional[List[str]] = None, group_by: Optional[List[str]] = None, | |
| order_by: Optional[str] = None, limit: Optional[int] = 100) -> str: | |
| parts = ["SELECT " + ", ".join(select_cols), "FROM " + base_table] | |
| for j in (joins or []): | |
| parts.append(j) | |
| if where: | |
| parts.append("WHERE " + " AND ".join(where)) | |
| if group_by: | |
| parts.append("GROUP BY " + ", ".join(group_by)) | |
| if order_by: | |
| parts.append("ORDER BY " + order_by) | |
| if limit is not None: | |
| parts.append(f"LIMIT {limit}") | |
| sql = " ".join(parts) | |
| if not sql.strip().endswith(';'): | |
| sql += ';' | |
| return sql | |
| async def tool_list_departments(question: str = ""): | |
| if not _has_tables('department'): | |
| return {"error": "department table not available"} | |
| cols = _colnames('department') | |
| name_col = 'dept_name' if 'dept_name' in cols else (cols[0] if cols else 'dept_name') | |
| sql = _build_sql([f"d.{name_col} AS department_name"], 'department d', order_by='department_name ASC') | |
| sql = fix_common_sql_errors(sql) | |
| sql = await _verify_sql(sql, question or 'List departments', 'list_departments') | |
| return {"sql": sql} | |
| async def tool_count_employees(question: str = ""): | |
| if not _has_tables('employee'): | |
| return {"error": "employee table not available"} | |
| sql = _build_sql(["COUNT(*) AS total_employees"], 'employee e', limit=None) | |
| sql = fix_common_sql_errors(sql) | |
| sql = await _verify_sql(sql, question or 'Total employees', 'count_employees') | |
| return {"sql": sql} | |
| async def tool_employee_count_by_department(question: str = ""): | |
| if not _has_tables('department','dept_emp'): | |
| return {"error": "required tables missing"} | |
| dcols = _colnames('department'); decols = _colnames('dept_emp') | |
| dept_name = 'dept_name' if 'dept_name' in dcols else (dcols[0] if dcols else 'dept_name') | |
| dept_no = 'dept_no' if 'dept_no' in dcols or 'dept_no' in decols else 'dept_no' | |
| emp_no = 'emp_no' if 'emp_no' in decols else 'emp_no' | |
| where = [] | |
| if 'to_date' in decols: | |
| where.append("de.to_date = '9999-01-01'") | |
| sql = _build_sql([f"d.{dept_name} AS department_name", f"COUNT(DISTINCT de.{emp_no}) AS employee_count"], | |
| 'department d', joins=[f"JOIN dept_emp de ON d.{dept_no} = de.{dept_no}"], | |
| where=where, group_by=[f"d.{dept_name}"], order_by='employee_count DESC') | |
| sql = fix_common_sql_errors(sql) | |
| sql = await _verify_sql(sql, question or 'Employees per department', 'employee_count_by_department') | |
| return {"sql": sql} | |
| async def tool_gender_pay_gap(question: str = ""): | |
| if not _has_tables('department','dept_emp','employee','salary'): | |
| return {"error": "required tables missing"} | |
| scolumns = _colnames('salary') | |
| sval = 'salary' if 'salary' in scolumns else ('amount' if 'amount' in scolumns else (scolumns[0] if scolumns else 'salary')) | |
| where = ["de.to_date = '9999-01-01'", "s.to_date = '9999-01-01'"] | |
| sql = _build_sql([ | |
| "d.dept_name", | |
| f"AVG(CASE WHEN e.gender='M' THEN s.{sval} END) AS avg_male_salary", | |
| f"AVG(CASE WHEN e.gender='F' THEN s.{sval} END) AS avg_female_salary", | |
| f"(AVG(CASE WHEN e.gender='M' THEN s.{sval} END) - AVG(CASE WHEN e.gender='F' THEN s.{sval} END)) AS gap", | |
| f"CASE WHEN AVG(CASE WHEN e.gender='M' THEN s.{sval} END)=0 THEN NULL ELSE ((AVG(CASE WHEN e.gender='M' THEN s.{sval} END) - AVG(CASE WHEN e.gender='F' THEN s.{sval} END))/NULLIF(AVG(CASE WHEN e.gender='M' THEN s.{sval} END),0))*100 END AS pct_gap" | |
| ], 'department d', | |
| joins=['JOIN dept_emp de ON d.dept_no = de.dept_no','JOIN employee e ON de.emp_no = e.emp_no','JOIN salary s ON e.emp_no = s.emp_no'], | |
| where=where, group_by=['d.dept_name'], order_by='gap DESC') | |
| sql = fix_common_sql_errors(sql) | |
| sql = await _verify_sql(sql, question or 'Gender pay gap', 'gender_pay_gap') | |
| return {"sql": sql} | |
| async def tool_salary_extremes(question: str = ""): | |
| if not _has_tables('salary'): | |
| return {"error": "salary table missing"} | |
| scolumns = _colnames('salary'); sval = 'salary' if 'salary' in scolumns else ('amount' if 'amount' in scolumns else (scolumns[0] if scolumns else 'salary')) | |
| where = ["s.to_date = '9999-01-01'"] if 'to_date' in scolumns else [] | |
| sql = _build_sql([f"MAX(s.{sval}) AS highest_salary", f"MIN(s.{sval}) AS lowest_salary"], 'salary s', where=where, limit=None) | |
| sql = fix_common_sql_errors(sql) | |
| sql = await _verify_sql(sql, question or 'Salary extremes', 'salary_extremes') | |
| return {"sql": sql} | |
| async def tool_salary_range_department(department_name: str = "Development", question: str = ""): | |
| if not _has_tables('department','dept_emp','salary'): | |
| return {"error": "required tables missing"} | |
| scolumns = _colnames('salary'); sval = 'salary' if 'salary' in scolumns else ('amount' if 'amount' in scolumns else (scolumns[0] if scolumns else 'salary')) | |
| # Escape single quotes in department_name for SQL safety (basic) | |
| escaped_dept = department_name.replace("'", "''") | |
| where = ["de.to_date = '9999-01-01'", "s.to_date = '9999-01-01'", f"d.dept_name = '{escaped_dept}'"] | |
| sql = _build_sql(["d.dept_name", f"MAX(s.{sval}) AS max_salary", f"MIN(s.{sval}) AS min_salary"], 'department d', | |
| joins=['JOIN dept_emp de ON d.dept_no = de.dept_no','JOIN salary s ON de.emp_no = s.emp_no'], | |
| where=where, group_by=['d.dept_name'], limit=None) | |
| sql = fix_common_sql_errors(sql) | |
| sql = await _verify_sql(sql, question or 'Department salary range', 'salary_range_department') | |
| return {"sql": sql, "department": department_name} | |
| async def tool_top_paid_employees(top_n: int = 10, question: str = ""): | |
| if not _has_tables('employee','salary'): | |
| return {"error": "required tables missing"} | |
| scolumns = _colnames('salary'); sval = 'salary' if 'salary' in scolumns else (scolumns[0] if scolumns else 'salary') | |
| where = ["s.to_date = '9999-01-01'"] if 'to_date' in scolumns else [] | |
| sql = _build_sql(["e.emp_no","e.first_name","e.last_name", f"s.{sval} AS salary"], 'employee e', | |
| joins=['JOIN salary s ON e.emp_no = s.emp_no'], where=where, | |
| order_by=f"s.{sval} DESC", limit=top_n) | |
| sql = fix_common_sql_errors(sql) | |
| sql = await _verify_sql(sql, question or 'Top paid employees', 'top_paid_employees') | |
| return {"sql": sql, "top_n": top_n} | |
| async def tool_department_highest_avg_salary(question: str = ""): | |
| if not _has_tables('department','dept_emp','salary'): | |
| return {"error": "required tables missing"} | |
| scolumns = _colnames('salary'); sval = 'salary' if 'salary' in scolumns else (scolumns[0] if scolumns else 'salary') | |
| where = ["de.to_date = '9999-01-01'", "s.to_date = '9999-01-01'"] | |
| sql = _build_sql(["d.dept_name", f"AVG(s.{sval}) AS avg_salary"], 'department d', | |
| joins=['JOIN dept_emp de ON d.dept_no = de.dept_no','JOIN salary s ON de.emp_no = s.emp_no'], | |
| where=where, group_by=['d.dept_name'], order_by='avg_salary DESC', limit=1) | |
| sql = fix_common_sql_errors(sql) | |
| sql = await _verify_sql(sql, question or 'Department highest average salary', 'dept_highest_avg_salary') | |
| return {"sql": sql} | |
| async def tool_department_manager_listing(question: str = ""): | |
| if not _has_tables('department','dept_manager','employee'): | |
| return {"error": "required tables missing"} | |
| where = ["dm.to_date = '9999-01-01'"] | |
| sql = _build_sql(["d.dept_name AS department_name","e.first_name","e.last_name"], 'department d', | |
| joins=['JOIN dept_manager dm ON d.dept_no = dm.dept_no','JOIN employee e ON dm.emp_no = e.emp_no'], | |
| where=where, order_by='department_name ASC') | |
| sql = fix_common_sql_errors(sql) | |
| sql = await _verify_sql(sql, question or 'Department managers', 'department_manager_listing') | |
| return {"sql": sql} | |
| # NEW: Complex analytical tools for business intelligence | |
| async def tool_gender_distribution_by_dept(question: str = ""): | |
| if not _has_tables('department','dept_emp','employee'): | |
| return {"error": "required tables missing"} | |
| where = ["de.to_date = '9999-01-01'"] | |
| sql = _build_sql([ | |
| "d.dept_name AS department_name", | |
| "e.gender", | |
| "COUNT(*) AS count", | |
| "ROUND(COUNT(*) * 100.0 / SUM(COUNT(*)) OVER (PARTITION BY d.dept_name), 2) AS percentage" | |
| ], 'department d', | |
| joins=['JOIN dept_emp de ON d.dept_no = de.dept_no','JOIN employee e ON de.emp_no = e.emp_no'], | |
| where=where, group_by=['d.dept_name', 'e.gender'], order_by='department_name, gender') | |
| sql = fix_common_sql_errors(sql) | |
| sql = await _verify_sql(sql, question or 'Gender distribution by department', 'gender_distribution_by_dept') | |
| return {"sql": sql} | |
| async def tool_hiring_trend_by_year(question: str = ""): | |
| if not _has_tables('employee'): | |
| return {"error": "required table employee missing"} | |
| sql = _build_sql([ | |
| "EXTRACT(YEAR FROM e.hire_date) AS hire_year", | |
| "COUNT(*) AS employees_hired", | |
| "COUNT(CASE WHEN e.gender='M' THEN 1 END) AS male_hires", | |
| "COUNT(CASE WHEN e.gender='F' THEN 1 END) AS female_hires" | |
| ], 'employee e', group_by=['EXTRACT(YEAR FROM e.hire_date)'], order_by='hire_year ASC') | |
| sql = fix_common_sql_errors(sql) | |
| sql = await _verify_sql(sql, question or 'Hiring trends by year', 'hiring_trend_by_year') | |
| return {"sql": sql} | |
| async def tool_employee_tenure_by_dept(question: str = ""): | |
| if not _has_tables('department','dept_emp','employee'): | |
| return {"error": "required tables missing"} | |
| where = ["de.to_date = '9999-01-01'"] | |
| sql = _build_sql([ | |
| "d.dept_name AS department_name", | |
| "AVG(EXTRACT(DAYS FROM (CURRENT_DATE - e.hire_date))/365.25) AS avg_tenure_years", | |
| "MIN(EXTRACT(DAYS FROM (CURRENT_DATE - e.hire_date))/365.25) AS min_tenure_years", | |
| "MAX(EXTRACT(DAYS FROM (CURRENT_DATE - e.hire_date))/365.25) AS max_tenure_years" | |
| ], 'department d', | |
| joins=['JOIN dept_emp de ON d.dept_no = de.dept_no','JOIN employee e ON de.emp_no = e.emp_no'], | |
| where=where, group_by=['d.dept_name'], order_by='avg_tenure_years DESC') | |
| sql = fix_common_sql_errors(sql) | |
| sql = await _verify_sql(sql, question or 'Employee tenure by department', 'employee_tenure_by_dept') | |
| return {"sql": sql} | |
| async def tool_female_managers_count(question: str = ""): | |
| if not _has_tables('dept_manager','employee'): | |
| return {"error": "required tables missing"} | |
| where = ["dm.to_date = '9999-01-01'", "e.gender = 'F'"] | |
| sql = _build_sql([ | |
| "COUNT(*) AS female_managers_count", | |
| "COUNT(*) * 100.0 / (SELECT COUNT(*) FROM dept_manager dm2 WHERE dm2.to_date = '9999-01-01') AS percentage_female_managers" | |
| ], 'dept_manager dm', | |
| joins=['JOIN employee e ON dm.emp_no = e.emp_no'], | |
| where=where, limit=None) | |
| sql = fix_common_sql_errors(sql) | |
| sql = await _verify_sql(sql, question or 'Female managers count', 'female_managers_count') | |
| return {"sql": sql} | |
| async def tool_avg_salary_by_title_dept(question: str = ""): | |
| """Get average salary by title and department.""" | |
| if not _has_tables('department', 'dept_emp', 'salary', 'title', 'employee'): | |
| return {"error": "required tables missing"} | |
| where = ["de.to_date = '9999-01-01'", "s.to_date = '9999-01-01'", "t.to_date = '9999-01-01'"] | |
| sql = _build_sql(["d.dept_name AS department_name", "t.title", "AVG(s.salary) AS avg_salary"], 'department d', | |
| joins=['JOIN dept_emp de ON d.dept_no = de.dept_no', | |
| 'JOIN employee e ON de.emp_no = e.emp_no', | |
| 'JOIN salary s ON e.emp_no = s.emp_no', | |
| 'JOIN title t ON e.emp_no = t.emp_no'], | |
| where=where, group_by=['d.dept_name', 't.title'], | |
| order_by='d.dept_name, avg_salary DESC', limit=100) | |
| sql = fix_common_sql_errors(sql) | |
| sql = await _verify_sql(sql, question or 'Average salary by title and department', 'avg_salary_by_title_dept') | |
| return {"sql": sql} | |
| # ---- Registration ---- | |
| # Dynamic schema support: the 14 fast_sql.* tools are purpose-built SQL | |
| # skeletons for the employee demo schema — register them ADDITIVELY when | |
| # the active profile is the employee demo. The GENERIC schema-grounded | |
| # tools (schema.list_tables / describe_table / sample_data / sql.build_select) | |
| # work on ANY database and are ALWAYS registered: they read the live | |
| # profile state at call time, so they survive profile switches. | |
| # (Previously they existed only in the non-employee boot branch — an | |
| # employee-profile boot left them unregistered, and every other database | |
| # then lost Level-3 hints to a 404 (KeyError -> HTTPException) on | |
| # schema.list_tables.) | |
| is_employee_profile = bool( | |
| getattr(aw, 'ACTIVE_SCHEMA_PROFILE', None) | |
| and aw.ACTIVE_SCHEMA_PROFILE.is_employee_db | |
| ) | |
| # ---------- Generic schema-grounded tools (any database) ---------- | |
| def tool_schema_list_tables(**kwargs): | |
| return {"tables": _live_tables()} | |
| def tool_schema_describe_table(table: str = "", **kwargs): | |
| if not table: | |
| return {"error": "table is required", "tables": _live_tables()} | |
| cols = _colnames(table) | |
| if not cols: | |
| # Case-insensitive retry | |
| lower_map = {t.lower(): t for t in _live_tables()} | |
| table = lower_map.get(table.lower()) | |
| if table: | |
| cols = _colnames(table) | |
| return {"table": table, "exists": bool(cols), "columns": cols} | |
| def tool_schema_sample_data(table: str = "", limit: int = 3, **kwargs): | |
| if not table: | |
| return {"error": "table is required"} | |
| try: | |
| from sqlalchemy import text as _sa_text | |
| safe_limit = max(1, min(int(limit), 10)) | |
| with aw.sql_database.connect() as conn: | |
| rows = conn.execute( | |
| _sa_text(f'SELECT * FROM "{table}" LIMIT {safe_limit}') | |
| ).mappings().all() | |
| return {"table": table, "rows": [dict(r) for r in rows]} | |
| except Exception: | |
| try: | |
| from sqlalchemy import text as _sa_text | |
| safe_limit = max(1, min(int(limit), 10)) | |
| with aw.sql_database.connect() as conn: | |
| rows = conn.execute( | |
| _sa_text(f"SELECT * FROM {table} LIMIT {safe_limit}") | |
| ).mappings().all() | |
| return {"table": table, "rows": [dict(r) for r in rows]} | |
| except Exception as e2: | |
| return {"table": table, "error": str(e2)} | |
| async def tool_sql_build_select(select_cols: Optional[List[str]] = None, base_table: str = "", | |
| joins: Optional[List[str]] = None, where: Optional[List[str]] = None, | |
| group_by: Optional[List[str]] = None, order_by: Optional[str] = None, | |
| limit: Optional[int] = 100, question: str = "", **kwargs): | |
| """Deterministic schema-grounded SELECT builder: validates requested | |
| columns against the live table, uses SQLAlchemy inspector, then an | |
| optional light LLM refinement. Returns a validated SQL hint.""" | |
| cols = _colnames(base_table) | |
| if not cols: | |
| lower_map = {t.lower(): t for t in _live_tables()} | |
| base_table = lower_map.get((base_table or "").lower(), base_table) | |
| cols = _colnames(base_table) | |
| if not cols: | |
| return {"error": f"Unknown table '{base_table}'", | |
| "tables": _live_tables()} | |
| _cols_lower = {str(c).lower(): str(c) for c in cols} | |
| for raw in (select_cols or []): | |
| # tokens like alias.col, col, COUNT(col) etc. | |
| for tok in re.findall(r"[A-Za-z_][\w]*(?:\.[A-Za-z_]\w*)?", str(raw)): | |
| t_, sep, c_ = tok.partition('.') | |
| if sep and c_.lower() not in _cols_lower and c_.lower() != '1': | |
| return {"error": f"Unknown column '{tok}' in table '{base_table}'", | |
| "columns": cols, | |
| "hint": "select base_table.column for aggregate args, e.g. COUNT(col)"} | |
| sql = _build_sql(select_cols=select_cols, base_table=base_table, joins=joins, | |
| where=where, group_by=group_by, order_by=order_by, limit=limit) | |
| sql = fix_common_sql_errors(sql) | |
| sql = await _verify_sql(sql, question or f"SELECT from {base_table}", f"build_select_{base_table}") | |
| return {"sql": sql, "table": base_table, "columns": cols} | |
| register_tool('schema.list_tables', tool_schema_list_tables) | |
| register_tool('schema.describe_table', tool_schema_describe_table) | |
| register_tool('schema.sample_data', tool_schema_sample_data) | |
| register_tool('sql.build_select', tool_sql_build_select) | |
| if is_employee_profile: | |
| register_tool('fast_sql.list_departments', tool_list_departments) | |
| register_tool('fast_sql.count_employees', tool_count_employees) | |
| register_tool('fast_sql.employee_count_by_department', tool_employee_count_by_department) | |
| register_tool('fast_sql.gender_pay_gap', tool_gender_pay_gap) # RE-ENABLED: Complex analytical queries | |
| register_tool('fast_sql.salary_extremes', tool_salary_extremes) | |
| register_tool('fast_sql.salary_range_department', tool_salary_range_department) | |
| register_tool('fast_sql.top_paid_employees', tool_top_paid_employees) | |
| register_tool('fast_sql.department_highest_avg_salary', tool_department_highest_avg_salary) | |
| register_tool('fast_sql.department_manager_listing', tool_department_manager_listing) | |
| # Register new complex analytical tools | |
| register_tool('fast_sql.gender_distribution_by_dept', tool_gender_distribution_by_dept) | |
| register_tool('fast_sql.hiring_trend_by_year', tool_hiring_trend_by_year) | |
| register_tool('fast_sql.employee_tenure_by_dept', tool_employee_tenure_by_dept) | |
| register_tool('fast_sql.female_managers_count', tool_female_managers_count) | |
| register_tool('fast_sql.avg_salary_by_title_dept', tool_avg_salary_by_title_dept) | |
| print('[MCP STARTUP] Registered GENERIC schema tools + employee-demo fast SQL tools: ' | |
| + ', '.join([k for k in TOOLS if k.startswith('fast_sql.')])) | |
| else: | |
| print('[MCP STARTUP] Registered GENERIC schema tools: ' | |
| + ', '.join([k for k in TOOLS if k.startswith('schema.') or k.startswith('sql.')]) | |
| + ' (non-employee schema active)') | |
| async def complete(req: CompleteRequest): | |
| """Minimal completion endpoint: | |
| - If context contains {'tool_call': {'name': ..., 'args': {...}}}, the server will run the tool and return its result. | |
| - Otherwise it returns a simple echo of the prompt under 'text'. | |
| This scaffold is purposely small so you can extend tool orchestration later. | |
| """ | |
| ctx = req.context or {} | |
| # Tool invocation flow | |
| if 'tool_call' in ctx: | |
| tc = ctx['tool_call'] | |
| name = tc.get('name') | |
| args = tc.get('args', {}) | |
| try: | |
| tool_res = await _call_tool(name, args) | |
| return {"text": f"Tool '{name}' executed.", "tool_result": tool_res} | |
| except KeyError as ke: | |
| raise HTTPException(status_code=404, detail=str(ke)) | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Tool call failed: {e}") | |
| # No tool requested: return a simple completion echo for now | |
| # Future: integrate a local model here and tool loop orchestration | |
| return {"text": req.prompt} | |
| def list_tools(): | |
| return {"tools": list(TOOLS.keys())} | |
| def health(): | |
| return {"status": "ok", "tools_registered": len(TOOLS)} | |