Aniruddha7/QueryLens-Text2SQL_DocVQA-V2

1from fastapi import FastAPI, HTTPException2from pydantic import BaseModel3from typing import Any, Dict, Optional, List4import asyncio5import importlib6import uuid7import os8import re9 10# Attempt to import agent workflow for schema/inspector & utilities (optional)11try:12    import Agent.agentic_workflow as aw13    print(f"[MCP STARTUP] Successfully imported Agent.agentic_workflow")14except Exception as e:15    print(f"[MCP STARTUP] Failed to import Agent.agentic_workflow: {e}")16    aw = None17 18app = FastAPI(title="Minimal MCP Server (dev)")19 20# Simple tool registry21TOOLS: Dict[str, Any] = {}22 23 24class CompleteRequest(BaseModel):25    prompt: str26    context: Optional[Dict[str, Any]] = None27    timeout: Optional[float] = 30.028 29 30def register_tool(name: str, fn: Any):31    TOOLS[name] = fn32 33 34async def _call_tool(name: str, args: Dict[str, Any]) -> Any:35    if name not in TOOLS:36        raise KeyError(f"Tool '{name}' not registered")37    fn = TOOLS[name]38    if asyncio.iscoroutinefunction(fn):39        return await fn(**(args or {}))40    else:41        # run sync in threadpool42        loop = asyncio.get_running_loop()43        return await loop.run_in_executor(None, lambda: fn(**(args or {})))44 45 46@app.on_event("startup")47async def _startup():48    # Register built-in/mock tools lazily to avoid importing heavy modules at import time49    # Try to register a real document scanner implementation; otherwise register a mock50    ds_registered = None51    try:52        from .tools import document_scanner as ds53        register_tool("document_scanner.process", ds.process_image)54        ds_registered = f"Agent.tools.document_scanner.process_image (relative import)"55    except Exception:56        # If relative import fails (running as script), try absolute import57        try:58            import Agent.tools.document_scanner as ds259            register_tool("document_scanner.process", ds2.process_image)60            ds_registered = f"Agent.tools.document_scanner.process_image (absolute import)"61        except Exception:62            # Register a fallback mock tool63            def _mock_doc(**kwargs):64                # Accept arbitrary keyword args (image_url, image_bytes, etc.) to be tolerant65                return {"doc_id": str(uuid.uuid4()), "text": "(mock OCR output)", "metadata": {"pages": 1, "confidence": 0.0}, "received_args": kwargs}66            register_tool("document_scanner.process", _mock_doc)67            ds_registered = "mock(document_scanner.process)"68 69    # Log which document scanner implementation was registered for easier debugging70    try:71        print(f"[MCP STARTUP] Registered document_scanner.process -> {ds_registered}")72    except Exception:73        pass74 75    # Register Granite Vision direct visual Q&A tool76    gv_registered = None77    try:78        from .tools import granite_vision as gv79        register_tool("granite_vision.qa", gv.qa)80        gv_registered = "Agent.tools.granite_vision.qa (relative import)"81    except Exception as e1:82        print(f"[MCP STARTUP] granite_vision relative import FAILED: {type(e1).__name__}: {e1}")83        try:84            import Agent.tools.granite_vision as gv285            register_tool("granite_vision.qa", gv2.qa)86            gv_registered = "Agent.tools.granite_vision.qa (absolute import)"87        except Exception as e2:88            print(f"[MCP STARTUP] granite_vision absolute import FAILED: {type(e2).__name__}: {e2}")89            def _mock_gv(**kwargs):90                return {"answer": "(granite_vision unavailable - mock)", "model": "mock", "received_args": kwargs}91            register_tool("granite_vision.qa", _mock_gv)92            gv_registered = "mock(granite_vision.qa)"93    print(f"[MCP STARTUP] Registered granite_vision.qa -> {gv_registered}")94 95 96    # Register OCR QA composite tool: runs OCR then runs the agent OCR QA helper97    try:98        # Use the global aw variable instead of re-importing99        if not aw:100            raise ImportError("Agent workflow not available")101 102        async def _ocr_qa_tool(image_url: Optional[str] = None, image_bytes: Optional[bytes] = None, question: Optional[str] = None, **kwargs):103            # Call the document scanner tool first (use the registered one to allow mocks)104            ds_fn = TOOLS.get("document_scanner.process")105            if not ds_fn:106                raise RuntimeError("document_scanner.process not registered")107 108            loop = asyncio.get_running_loop()109            try:110                if asyncio.iscoroutinefunction(ds_fn):111                    tool_res = await ds_fn(image_url=image_url, image_bytes=image_bytes, options=kwargs)112                else:113                    tool_res = await loop.run_in_executor(None, lambda: ds_fn(image_url=image_url, image_bytes=image_bytes, options=kwargs))114            except Exception as e:115                return {"error": f"OCR tool failed: {e}"}116 117            doc_id = tool_res.get("doc_id") if isinstance(tool_res, dict) else None118            ocr_text = (tool_res.get("text") if isinstance(tool_res, dict) else str(tool_res)) or ""119 120            # Persist OCR text so downstream callers can retrieve it via doc_id121            try:122                if doc_id and ocr_text:123                    docs_dir = os.path.join("chat_store", "docs")124                    os.makedirs(docs_dir, exist_ok=True)125                    doc_path = os.path.join(docs_dir, f"{doc_id}.txt")126                    # Write file (overwrite if exists)127                    with open(doc_path, "w", encoding="utf-8") as df:128                        df.write(ocr_text)129                    # Also add a short history entry if agent chat store available130                    try:131                        if hasattr(aw, 'chat_store_private') and hasattr(aw, 'ChatMessage') and hasattr(aw, 'MessageRole'):132                            user_msg = aw.ChatMessage(role=aw.MessageRole.USER, content=f"Uploaded document {doc_id}")133                            assistant_msg = aw.ChatMessage(role=aw.MessageRole.ASSISTANT, content=f"OCR stored: {ocr_text[:400]}")134                            aw.chat_store_private.add_message(key="conversation", message=user_msg)135                            aw.chat_store_private.add_message(key="conversation", message=assistant_msg)136                            try:137                                aw.chat_store_private.persist(str(aw.private_store_path))138                            except Exception:139                                pass140                    except Exception:141                        pass142            except Exception:143                # Non-fatal: continue even if persistence fails144                pass145 146            # If a question was provided, answer it using the agent helper (run in executor to avoid blocking)147            if question:148                try:149                    # Use executor if helper is sync to avoid blocking event loop150                    if callable(getattr(aw, 'ocr_agent_qa', None)):151                        answer = await loop.run_in_executor(None, lambda: aw.ocr_agent_qa(question, doc_id))152                        return {"doc_id": doc_id, "answer": answer}153                    else:154                        return {"doc_id": doc_id, "answer": "ocr_agent_qa not available"}155                except Exception as e:156                    return {"error": f"ocr_agent_qa failed: {e}", "tool_result": tool_res}157 158            # No question: return the OCR result and doc_id159            return {"doc_id": doc_id, "ocr_text": ocr_text}160 161        register_tool("ocr_qa", _ocr_qa_tool)162        print(f"[MCP STARTUP] Registered ocr_qa -> real _ocr_qa_tool (agent available)")163    except Exception:164        # If agent import fails, register a tolerant mock tool165        def _ocr_qa_mock(**kwargs):166            return {"doc_id": str(uuid.uuid4()), "answer": "(ocr_qa unavailable - mock response)", "received_args": kwargs}167        register_tool("ocr_qa", _ocr_qa_mock)168        print(f"[MCP STARTUP] Registered ocr_qa -> mock(ocr_qa)")169 170    # ------------------------------------------------------------------171    # Deterministic schema-aware SQL helper tools (agentic compliant)172    # ------------------------------------------------------------------173    ENABLE_FAST_SQL = os.environ.get("ENABLE_MCP_FAST_TOOLS", "1") in ("1","true","True")174    if not ENABLE_FAST_SQL:175        print("[MCP STARTUP] Fast SQL tools disabled via ENABLE_MCP_FAST_TOOLS=0")176        return177 178    # Initialize database connection if not already done179    if not aw:180        print("[MCP STARTUP][FAST] Agent workflow not available; skipping fast SQL tool registration.")181        return182        183    if not getattr(aw, 'inspector', None) or not getattr(aw, 'available_tables', None):184        print("[MCP STARTUP][FAST] Database not initialized; attempting to initialize...")185        try:186            # Initialize database connection - use await since we're in async context187            db_uri = os.environ.get("DB_CONNECTION_URL")188            if not db_uri:189                print("[MCP STARTUP][FAST] DB_CONNECTION_URL not set; skipping fast SQL tool registration.")190                return191            db_result = await aw.initialize_database(db_uri)192            if db_result.get("tables"):193                print(f"[MCP STARTUP][FAST] Database initialized with {len(db_result['tables'])} tables")194            else:195                print("[MCP STARTUP][FAST] Database initialization failed; skipping fast SQL tool registration.")196                return197        except Exception as db_err:198            print(f"[MCP STARTUP][FAST] Database initialization error: {db_err}; skipping fast SQL tool registration.")199            return200 201    inspector = aw.inspector202    available_tables = aw.available_tables203 204    def _colnames(table: str) -> List[str]:205        try:206            return [c['name'] for c in inspector.get_columns(table_name=table)]207        except Exception:208            return []209 210    def _has_tables(*tables: str) -> bool:211        at_lower = {t.lower() for t in available_tables}212        return all(t.lower() in at_lower for t in tables)213 214    # Simple validation & sanitization re-use215    validate_sql_query = getattr(aw, 'validate_sql_query', lambda q: {"valid": True})216    fix_common_sql_errors = getattr(aw, 'fix_common_sql_errors', lambda q: q)217 218    async def _verify_sql(sql: str, question: str, tag: str) -> str:219        """Optional light LLM verification (short prompt) if LLM available."""220        if not aw or not getattr(aw, 'Settings', None) or not getattr(aw.Settings, 'llm', None):221            return sql222        try:223            prompt = (224                f"{aw.EMPLOYEE_DB_SCHEMA}\n\nYou are an SQL verifier. Tag={tag}."225                " If the candidate SQL is valid and uses only existing tables/columns keep it."226                " If SMALL fixes (alias, LIMIT 100, semicolon) needed, fix them. Never invent tables/columns."227                f"\nQuestion: {question}\nCandidate SQL:\n{sql}\nReturn ONLY final SQL." )228            resp = aw.Settings.llm.complete(prompt=prompt)229            txt = getattr(resp, 'text', str(resp)).strip()230            if '```' in txt:231                parts = re.split(r"```(?:sql)?", txt, flags=re.IGNORECASE)232                txt = '\n'.join(p for p in parts if 'select' in p.lower() or 'with' in p.lower())233            if txt.lower().startswith('sql'):234                txt = txt[3:].strip()235            txt = fix_common_sql_errors(txt)236            v = validate_sql_query(txt)237            if v.get('valid'):238                return txt239        except Exception as ve:240            print(f"[FAST][VERIFY][WARN] Verification failed: {ve}; using original SQL")241        return sql242 243    def _build_sql(select_cols: List[str], base_table: str, joins: Optional[List[str]] = None,244                   where: Optional[List[str]] = None, group_by: Optional[List[str]] = None,245                   order_by: Optional[str] = None, limit: Optional[int] = 100) -> str:246        parts = ["SELECT " + ", ".join(select_cols), "FROM " + base_table]247        for j in (joins or []):248            parts.append(j)249        if where:250            parts.append("WHERE " + " AND ".join(where))251        if group_by:252            parts.append("GROUP BY " + ", ".join(group_by))253        if order_by:254            parts.append("ORDER BY " + order_by)255        if limit is not None:256            parts.append(f"LIMIT {limit}")257        sql = " ".join(parts)258        if not sql.strip().endswith(';'):259            sql += ';'260        return sql261 262    async def tool_list_departments(question: str = ""):263        if not _has_tables('department'):264            return {"error": "department table not available"}265        cols = _colnames('department')266        name_col = 'dept_name' if 'dept_name' in cols else (cols[0] if cols else 'dept_name')267        sql = _build_sql([f"d.{name_col} AS department_name"], 'department d', order_by='department_name ASC')268        sql = fix_common_sql_errors(sql)269        sql = await _verify_sql(sql, question or 'List departments', 'list_departments')270        return {"sql": sql}271 272    async def tool_count_employees(question: str = ""):273        if not _has_tables('employee'):274            return {"error": "employee table not available"}275        sql = _build_sql(["COUNT(*) AS total_employees"], 'employee e', limit=None)276        sql = fix_common_sql_errors(sql)277        sql = await _verify_sql(sql, question or 'Total employees', 'count_employees')278        return {"sql": sql}279 280    async def tool_employee_count_by_department(question: str = ""):281        if not _has_tables('department','dept_emp'):282            return {"error": "required tables missing"}283        dcols = _colnames('department'); decols = _colnames('dept_emp')284        dept_name = 'dept_name' if 'dept_name' in dcols else (dcols[0] if dcols else 'dept_name')285        dept_no = 'dept_no' if 'dept_no' in dcols or 'dept_no' in decols else 'dept_no'286        emp_no = 'emp_no' if 'emp_no' in decols else 'emp_no'287        where = []288        if 'to_date' in decols:289            where.append("de.to_date = '9999-01-01'")290        sql = _build_sql([f"d.{dept_name} AS department_name", f"COUNT(DISTINCT de.{emp_no}) AS employee_count"],291                         'department d', joins=[f"JOIN dept_emp de ON d.{dept_no} = de.{dept_no}"],292                         where=where, group_by=[f"d.{dept_name}"], order_by='employee_count DESC')293        sql = fix_common_sql_errors(sql)294        sql = await _verify_sql(sql, question or 'Employees per department', 'employee_count_by_department')295        return {"sql": sql}296 297    async def tool_gender_pay_gap(question: str = ""):298        if not _has_tables('department','dept_emp','employee','salary'):299            return {"error": "required tables missing"}300        scolumns = _colnames('salary')301        sval = 'salary' if 'salary' in scolumns else ('amount' if 'amount' in scolumns else (scolumns[0] if scolumns else 'salary'))302        where = ["de.to_date = '9999-01-01'", "s.to_date = '9999-01-01'"]303        sql = _build_sql([304            "d.dept_name",305            f"AVG(CASE WHEN e.gender='M' THEN s.{sval} END) AS avg_male_salary",306            f"AVG(CASE WHEN e.gender='F' THEN s.{sval} END) AS avg_female_salary",307            f"(AVG(CASE WHEN e.gender='M' THEN s.{sval} END) - AVG(CASE WHEN e.gender='F' THEN s.{sval} END)) AS gap",308            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"309        ], 'department d',310            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'],311            where=where, group_by=['d.dept_name'], order_by='gap DESC')312        sql = fix_common_sql_errors(sql)313        sql = await _verify_sql(sql, question or 'Gender pay gap', 'gender_pay_gap')314        return {"sql": sql}315 316    async def tool_salary_extremes(question: str = ""):317        if not _has_tables('salary'):318            return {"error": "salary table missing"}319        scolumns = _colnames('salary'); sval = 'salary' if 'salary' in scolumns else ('amount' if 'amount' in scolumns else (scolumns[0] if scolumns else 'salary'))320        where = ["s.to_date = '9999-01-01'"] if 'to_date' in scolumns else []321        sql = _build_sql([f"MAX(s.{sval}) AS highest_salary", f"MIN(s.{sval}) AS lowest_salary"], 'salary s', where=where, limit=None)322        sql = fix_common_sql_errors(sql)323        sql = await _verify_sql(sql, question or 'Salary extremes', 'salary_extremes')324        return {"sql": sql}325 326    async def tool_salary_range_department(department_name: str = "Development", question: str = ""):327        if not _has_tables('department','dept_emp','salary'):328            return {"error": "required tables missing"}329        scolumns = _colnames('salary'); sval = 'salary' if 'salary' in scolumns else ('amount' if 'amount' in scolumns else (scolumns[0] if scolumns else 'salary'))330        # Escape single quotes in department_name for SQL safety (basic)331        escaped_dept = department_name.replace("'", "''")332        where = ["de.to_date = '9999-01-01'", "s.to_date = '9999-01-01'", f"d.dept_name = '{escaped_dept}'"]333        sql = _build_sql(["d.dept_name", f"MAX(s.{sval}) AS max_salary", f"MIN(s.{sval}) AS min_salary"], 'department d',334                         joins=['JOIN dept_emp de ON d.dept_no = de.dept_no','JOIN salary s ON de.emp_no = s.emp_no'],335                         where=where, group_by=['d.dept_name'], limit=None)336        sql = fix_common_sql_errors(sql)337        sql = await _verify_sql(sql, question or 'Department salary range', 'salary_range_department')338        return {"sql": sql, "department": department_name}339 340    async def tool_top_paid_employees(top_n: int = 10, question: str = ""):341        if not _has_tables('employee','salary'):342            return {"error": "required tables missing"}343        scolumns = _colnames('salary'); sval = 'salary' if 'salary' in scolumns else (scolumns[0] if scolumns else 'salary')344        where = ["s.to_date = '9999-01-01'"] if 'to_date' in scolumns else []345        sql = _build_sql(["e.emp_no","e.first_name","e.last_name", f"s.{sval} AS salary"], 'employee e',346                         joins=['JOIN salary s ON e.emp_no = s.emp_no'], where=where,347                         order_by=f"s.{sval} DESC", limit=top_n)348        sql = fix_common_sql_errors(sql)349        sql = await _verify_sql(sql, question or 'Top paid employees', 'top_paid_employees')350        return {"sql": sql, "top_n": top_n}351 352    async def tool_department_highest_avg_salary(question: str = ""):353        if not _has_tables('department','dept_emp','salary'):354            return {"error": "required tables missing"}355        scolumns = _colnames('salary'); sval = 'salary' if 'salary' in scolumns else (scolumns[0] if scolumns else 'salary')356        where = ["de.to_date = '9999-01-01'", "s.to_date = '9999-01-01'"]357        sql = _build_sql(["d.dept_name", f"AVG(s.{sval}) AS avg_salary"], 'department d',358                         joins=['JOIN dept_emp de ON d.dept_no = de.dept_no','JOIN salary s ON de.emp_no = s.emp_no'],359                         where=where, group_by=['d.dept_name'], order_by='avg_salary DESC', limit=1)360        sql = fix_common_sql_errors(sql)361        sql = await _verify_sql(sql, question or 'Department highest average salary', 'dept_highest_avg_salary')362        return {"sql": sql}363 364    async def tool_department_manager_listing(question: str = ""):365        if not _has_tables('department','dept_manager','employee'):366            return {"error": "required tables missing"}367        where = ["dm.to_date = '9999-01-01'"]368        sql = _build_sql(["d.dept_name AS department_name","e.first_name","e.last_name"], 'department d',369                         joins=['JOIN dept_manager dm ON d.dept_no = dm.dept_no','JOIN employee e ON dm.emp_no = e.emp_no'],370                         where=where, order_by='department_name ASC')371        sql = fix_common_sql_errors(sql)372        sql = await _verify_sql(sql, question or 'Department managers', 'department_manager_listing')373        return {"sql": sql}374 375    # NEW: Complex analytical tools for business intelligence376    async def tool_gender_distribution_by_dept(question: str = ""):377        if not _has_tables('department','dept_emp','employee'):378            return {"error": "required tables missing"}379        where = ["de.to_date = '9999-01-01'"]380        sql = _build_sql([381            "d.dept_name AS department_name",382            "e.gender",383            "COUNT(*) AS count",384            "ROUND(COUNT(*) * 100.0 / SUM(COUNT(*)) OVER (PARTITION BY d.dept_name), 2) AS percentage"385        ], 'department d',386            joins=['JOIN dept_emp de ON d.dept_no = de.dept_no','JOIN employee e ON de.emp_no = e.emp_no'],387            where=where, group_by=['d.dept_name', 'e.gender'], order_by='department_name, gender')388        sql = fix_common_sql_errors(sql)389        sql = await _verify_sql(sql, question or 'Gender distribution by department', 'gender_distribution_by_dept')390        return {"sql": sql}391 392    async def tool_hiring_trend_by_year(question: str = ""):393        if not _has_tables('employee'):394            return {"error": "required table employee missing"}395        sql = _build_sql([396            "EXTRACT(YEAR FROM e.hire_date) AS hire_year",397            "COUNT(*) AS employees_hired",398            "COUNT(CASE WHEN e.gender='M' THEN 1 END) AS male_hires",399            "COUNT(CASE WHEN e.gender='F' THEN 1 END) AS female_hires"400        ], 'employee e', group_by=['EXTRACT(YEAR FROM e.hire_date)'], order_by='hire_year ASC')401        sql = fix_common_sql_errors(sql)402        sql = await _verify_sql(sql, question or 'Hiring trends by year', 'hiring_trend_by_year')403        return {"sql": sql}404 405    async def tool_employee_tenure_by_dept(question: str = ""):406        if not _has_tables('department','dept_emp','employee'):407            return {"error": "required tables missing"}408        where = ["de.to_date = '9999-01-01'"]409        sql = _build_sql([410            "d.dept_name AS department_name",411            "AVG(EXTRACT(DAYS FROM (CURRENT_DATE - e.hire_date))/365.25) AS avg_tenure_years",412            "MIN(EXTRACT(DAYS FROM (CURRENT_DATE - e.hire_date))/365.25) AS min_tenure_years", 413            "MAX(EXTRACT(DAYS FROM (CURRENT_DATE - e.hire_date))/365.25) AS max_tenure_years"414        ], 'department d',415            joins=['JOIN dept_emp de ON d.dept_no = de.dept_no','JOIN employee e ON de.emp_no = e.emp_no'],416            where=where, group_by=['d.dept_name'], order_by='avg_tenure_years DESC')417        sql = fix_common_sql_errors(sql)418        sql = await _verify_sql(sql, question or 'Employee tenure by department', 'employee_tenure_by_dept')419        return {"sql": sql}420 421    async def tool_female_managers_count(question: str = ""):422        if not _has_tables('dept_manager','employee'):423            return {"error": "required tables missing"}424        where = ["dm.to_date = '9999-01-01'", "e.gender = 'F'"]425        sql = _build_sql([426            "COUNT(*) AS female_managers_count",427            "COUNT(*) * 100.0 / (SELECT COUNT(*) FROM dept_manager dm2 WHERE dm2.to_date = '9999-01-01') AS percentage_female_managers"428        ], 'dept_manager dm',429            joins=['JOIN employee e ON dm.emp_no = e.emp_no'],430            where=where, limit=None)431        sql = fix_common_sql_errors(sql)432        sql = await _verify_sql(sql, question or 'Female managers count', 'female_managers_count')433        return {"sql": sql}434 435    async def tool_avg_salary_by_title_dept(question: str = ""):436        """Get average salary by title and department."""437        if not _has_tables('department', 'dept_emp', 'salary', 'title', 'employee'):438            return {"error": "required tables missing"}439        where = ["de.to_date = '9999-01-01'", "s.to_date = '9999-01-01'", "t.to_date = '9999-01-01'"]440        sql = _build_sql(["d.dept_name AS department_name", "t.title", "AVG(s.salary) AS avg_salary"], 'department d',441                         joins=['JOIN dept_emp de ON d.dept_no = de.dept_no',442                               'JOIN employee e ON de.emp_no = e.emp_no',443                               'JOIN salary s ON e.emp_no = s.emp_no',444                               'JOIN title t ON e.emp_no = t.emp_no'],445                         where=where, group_by=['d.dept_name', 't.title'], 446                         order_by='d.dept_name, avg_salary DESC', limit=100)447        sql = fix_common_sql_errors(sql)448        sql = await _verify_sql(sql, question or 'Average salary by title and department', 'avg_salary_by_title_dept')449        return {"sql": sql}450 451    register_tool('fast_sql.list_departments', tool_list_departments)452    register_tool('fast_sql.count_employees', tool_count_employees)453    register_tool('fast_sql.employee_count_by_department', tool_employee_count_by_department)454    register_tool('fast_sql.gender_pay_gap', tool_gender_pay_gap)  # RE-ENABLED: Complex analytical queries455    register_tool('fast_sql.salary_extremes', tool_salary_extremes)456    register_tool('fast_sql.salary_range_department', tool_salary_range_department)457    register_tool('fast_sql.top_paid_employees', tool_top_paid_employees)458    register_tool('fast_sql.department_highest_avg_salary', tool_department_highest_avg_salary)459    register_tool('fast_sql.department_manager_listing', tool_department_manager_listing)460    461    # Register new complex analytical tools462    register_tool('fast_sql.gender_distribution_by_dept', tool_gender_distribution_by_dept)463    register_tool('fast_sql.hiring_trend_by_year', tool_hiring_trend_by_year)464    register_tool('fast_sql.employee_tenure_by_dept', tool_employee_tenure_by_dept)465    register_tool('fast_sql.female_managers_count', tool_female_managers_count)466    register_tool('fast_sql.avg_salary_by_title_dept', tool_avg_salary_by_title_dept)467    print('[MCP STARTUP] Registered fast SQL tools: ' + ', '.join([k for k in TOOLS if k.startswith('fast_sql.')]))468 469 470@app.post("/complete")471async def complete(req: CompleteRequest):472    """Minimal completion endpoint:473    - If context contains {'tool_call': {'name': ..., 'args': {...}}}, the server will run the tool and return its result.474    - Otherwise it returns a simple echo of the prompt under 'text'.475    This scaffold is purposely small so you can extend tool orchestration later.476    """477    ctx = req.context or {}478    # Tool invocation flow479    if 'tool_call' in ctx:480        tc = ctx['tool_call']481        name = tc.get('name')482        args = tc.get('args', {})483        try:484            tool_res = await _call_tool(name, args)485            return {"text": f"Tool '{name}' executed.", "tool_result": tool_res}486        except KeyError as ke:487            raise HTTPException(status_code=404, detail=str(ke))488        except Exception as e:489            raise HTTPException(status_code=500, detail=f"Tool call failed: {e}")490 491    # No tool requested: return a simple completion echo for now492    # Future: integrate a local model here and tool loop orchestration493    return {"text": req.prompt}494 495 496@app.get("/tools")497def list_tools():498    return {"tools": list(TOOLS.keys())}499 500 501@app.get("/health")502def health():503    return {"status": "ok", "tools_registered": len(TOOLS)}504