Upload mcp_client.py

Client/mcp_client.py

@@ -0,0 +1,167 @@
+# fastmcp_client.py
+import asyncio
+import threading
+from fastmcp import Client
+import re
+import logging
+logging.basicConfig(level=logging.INFO, format="[%(asctime)s] %(levelname)-7s %(message)s")
+logger = logging.getLogger("pdf-rag-mcp-client")
+class MCPClientExcelRAG:
+    """
+    SSE-based MCP client for the Excel RAG tool `ask_excel`.
+    """
+
+    def __init__(self, url: str = "http://127.0.0.1:8001/sse"):
+    #def __init__(self, url: str = "http://127.0.0.1:7860/gradio_api/mcp/"):    
+        self.url = url
+        self._loop = None
+        self._thread = None
+        self._client: Client | None = None
+        self._ready = threading.Event()
+        self._available_tools = []
+
+    # ---------- Background runner ----------
+    def _runner(self):
+        self._loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(self._loop)
+
+        async def start_client():
+            self._client = Client(self.url)
+            await self._client.__aenter__()
+            try:
+                tools = await self._client.list_tools()
+                self._available_tools = [t.name for t in tools]
+                print(f"Available MCP tools: {self._available_tools}")
+            except Exception as e:
+                print(f"Failed to list tools: {e}")
+            self._ready.set()
+            await asyncio.Event().wait()  # keep loop alive
+
+        self._loop.run_until_complete(start_client())
+
+    def start(self, timeout: float = 10):
+        if self._thread and self._thread.is_alive():
+            return
+        self._thread = threading.Thread(target=self._runner, daemon=True)
+        self._thread.start()
+        if not self._ready.wait(timeout):
+            raise RuntimeError("SSE MCP client failed to start")
+
+    def stop(self):
+        if not self._loop or not self._thread.is_alive():
+            return
+        for task in asyncio.all_tasks(self._loop):
+            if not task.done():
+                self._loop.call_soon_threadsafe(task.cancel)
+                break
+        self._loop.call_soon_threadsafe(self._loop.stop)
+        self._thread.join(timeout=5)
+
+    # ---------- RAG tool call ----------
+    
+    
+    def ask_excel(self, question: str, top_k: int = 5, sheet: int | str = 0, temperature: float = 0.1) -> str:
+        if "ask_excel" not in self._available_tools:
+            return "❌ Tool 'ask_excel' not available on the server."
+        logger.info(f"[CLIENT] call ask_excel q='{question[:80]}...' top_k={top_k} sheet={sheet} temp={temperature}")
+        fut = asyncio.run_coroutine_threadsafe(
+            self._client.call_tool("ask_excel", {
+                "question": question,
+                "top_k": top_k,
+                "sheet": sheet,
+                "temperature": temperature,
+            }),
+            self._loop,
+        )
+        result = fut.result(timeout=60)
+        logger.info(f"[CLIENT] got result type={type(result)} has_content={hasattr(result, 'content')}")
+
+        # ✅ Extract text from content blocks if available
+        if hasattr(result, "content") and result.content:
+            texts = []
+            for block in result.content:
+                # Handle FastMCP TextContent objects or dict blocks
+                if hasattr(block, "text"):
+                    texts.append(block.text)
+                elif isinstance(block, dict) and block.get("type") == "text":
+                    texts.append(block.get("text", ""))
+            answer = "\n".join(texts).strip() if texts else str(result.content)
+        elif hasattr(result, "structured_content") and result.structured_content:
+            sc = result.structured_content
+            if isinstance(sc, dict) and "result" in sc:
+                answer = sc["result"]
+            else:
+                answer = str(sc)
+        elif hasattr(result, "data") and result.data:
+            answer = str(result.data)
+        else:
+            answer = str(result)
+
+        # ✅ Optional: Remove [Row N] citations if you don't want them
+        answer = re.sub(r"\s*\[Row\s+\d+\]\s*", "", answer).strip()
+        logger.info(f"[CLIENT] answer_len={len(answer)} preview: {answer[:200]}")
+
+        return answer
+    
+
+    def ask_pdf(self, question: str, top_k: int = 5, temperature: float = 0.1) -> str:
+        if "ask_pdf" not in self._available_tools:
+            return "❌ Tool 'ask_pdf' not available on the server."
+        logger.info(f"[CLIENT] call ask_pdf q='{question[:80]}...' top_k={top_k} temp={temperature}")
+        fut = asyncio.run_coroutine_threadsafe(
+            self._client.call_tool("ask_pdf", {
+                "question": question,
+                "top_k": top_k,
+                "temperature": temperature,
+            }),
+            self._loop,
+        )
+        result = fut.result(timeout=60)
+        # Extract text similar to ask_excel
+        if hasattr(result, "content") and result.content:
+            texts = []
+            for block in result.content:
+                # Handle FastMCP TextContent objects or dict blocks
+                if hasattr(block, "text"):
+                    texts.append(block.text)
+                elif isinstance(block, dict) and block.get("type") == "text":
+                    texts.append(block.get("text", ""))
+            answer = "\n".join(texts).strip() if texts else str(result.content)
+        elif hasattr(result, "structured_content") and result.structured_content:
+            sc = result.structured_content
+            if isinstance(sc, dict) and "result" in sc:
+                answer = sc["result"]
+            else:
+                answer = str(sc)
+        elif hasattr(result, "data") and result.data:
+            answer = str(result.data)
+        else:
+            answer = str(result)
+
+        # ✅ Optional: Remove [Row N] citations if you don't want them
+        answer = re.sub(r"\s*\[Row\s+\d+\]\s*", "", answer).strip()
+        logger.info(f"[CLIENT] answer_len={len(answer)} preview: {answer[:200]}")
+
+        return answer
+
+    
+    def ask_link(self, question: str, temperature: float = 0.1, subquery_context: str = None, top_k: int = 5) -> str:
+        if "ask_link" not in self._available_tools:
+            return "❌ Tool 'ask_link' not available on the server."
+        fut = asyncio.run_coroutine_threadsafe(
+            self._client.call_tool("ask_link", {
+                "query": question,
+                "temperature": temperature,
+                "subquery_context": subquery_context or "",
+                "top_k": top_k,
+            }),
+            self._loop,
+        )
+        result = fut.result(timeout=60)
+        # Extract text
+        if hasattr(result, "content") and result.content:
+            texts = [block.text for block in result.content if hasattr(block, "text")]
+            return "\n".join(texts).strip() if texts else str(result.content)
+        return str(result)
+    
+