Aniruddha7/QueryLens-Text2SQL_DocVQA-V2

1"""Minimal MCP client for local MCP server.2 3Provides mcp_complete(prompt, timeout, context) -> str4Tries to POST to MCP server and return the completion text. Raises on errors so caller can fallback.5"""6import os7import json8import requests9from typing import Any, Dict, Optional10 11MCP_URL = os.environ.get("MCP_URL", "http://127.0.0.1:8701/complete")12 13 14def mcp_complete(prompt: str, timeout: float = 30.0, context: Optional[Dict[str, Any]] = None) -> str:15    """Send a completion request to the local MCP server and return text.16 17    Raises requests.RequestException on transport errors or ValueError if MCP returns error.18    """19    payload = {20        "prompt": prompt,21        "context": context or {},22        "timeout": timeout23    }24    headers = {"Content-Type": "application/json"}25    resp = requests.post(MCP_URL, data=json.dumps(payload), headers=headers, timeout=timeout)26    resp.raise_for_status()27    body = resp.json()28    # Expecting {'text': '...'} from MCP server29    if not isinstance(body, dict) or 'text' not in body:30        raise ValueError(f"Unexpected MCP response: {body}")31    return body['text']32 33 34def mcp_call_tool(prompt: str, timeout: float = 60.0, context: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:35    """Call the MCP /complete endpoint and return the full JSON response.36 37    This is useful when asking MCP to run a registered tool; the tool result will38    appear under the `tool_result` key in the returned dict.39    """40    payload = {41        "prompt": prompt,42        "context": context or {},43        "timeout": timeout,44    }45    headers = {"Content-Type": "application/json"}46    resp = requests.post(MCP_URL, data=json.dumps(payload), headers=headers, timeout=timeout)47    resp.raise_for_status()48    return resp.json()49