| import os |
| import requests |
| import json |
| from typing import Dict, Any, List, Optional |
| from langchain.tools import BaseTool |
| from pydantic import Field |
| import logging |
|
|
| logger = logging.getLogger(__name__) |
|
|
| class MCPClient: |
| """Client for making authenticated REST API calls to the MCP server.""" |
| |
| def __init__(self, mcp_url: str, api_key: str): |
| self.mcp_url = mcp_url |
| self.headers = { |
| "x-api-key": api_key, |
| "Content-Type": "application/json" |
| } |
| |
| def post(self, endpoint: str, data: Dict[str, Any]) -> Dict[str, Any]: |
| """Send a POST request to a given MCP endpoint.""" |
| try: |
| url = f"{self.mcp_url}/{endpoint}" |
| response = requests.post(url, headers=self.headers, data=json.dumps(data)) |
| response.raise_for_status() |
| return response.json() |
| except requests.exceptions.HTTPError as http_err: |
| logger.error(f"HTTP error occurred: {http_err} - {response.text}") |
| return {"status": "error", "message": f"HTTP error: {response.status_code} {response.reason}"} |
| except requests.exceptions.RequestException as req_err: |
| logger.error(f"Request error occurred: {req_err}") |
| return {"status": "error", "message": f"Request failed: {req_err}"} |
| except json.JSONDecodeError: |
| logger.error("Failed to decode JSON response.") |
| return {"status": "error", "message": "Invalid JSON response from server."} |
|
|
|
|
| class SchemaSearchTool(BaseTool): |
| """LangChain tool for searching database schemas.""" |
| |
| name: str = "schema_search" |
| description: str = """ |
| Search for relevant database schemas based on a natural language query. |
| Use this when you need to find which tables/columns are relevant to a user's question. |
| Input should be a descriptive query like 'patient information' or 'drug trials'. |
| """ |
| mcp_client: MCPClient |
| |
| def _run(self, query: str) -> str: |
| """Execute schema search.""" |
| response = self.mcp_client.post("discovery/get_relevant_schemas", {"query": query}) |
| |
| if response.get("status") == "success": |
| schemas = response.get("schemas", []) |
| if schemas: |
| schema_text = "Found relevant schemas:\\n" |
| for schema in schemas: |
| schema_text += f"- {schema.get('database', 'Unknown')}.{schema.get('table', 'Unknown')}.{schema.get('name', 'Unknown')} ({schema.get('type', ['Unknown'])[0]})\\n" |
| return schema_text |
| else: |
| return "No relevant schemas found." |
| else: |
| return f"Error searching schemas: {response.get('message', 'Unknown error')}" |
| |
| async def _arun(self, query: str) -> str: |
| """Async version - just calls sync version.""" |
| return self._run(query) |
|
|
|
|
| class JoinPathFinderTool(BaseTool): |
| """LangChain tool for finding join paths between tables.""" |
| |
| name: str = "find_join_path" |
| description: str = """ |
| Find how to join two tables together using foreign key relationships. |
| Use this when you need to query across multiple tables. |
| Input should be two table names separated by a comma, like 'patients,studies'. |
| """ |
| mcp_client: MCPClient |
|
|
| def _run(self, table_names: str) -> str: |
| """Find join path.""" |
| try: |
| tables = [t.strip() for t in table_names.split(',')] |
| if len(tables) != 2: |
| return "Please provide exactly two table names separated by a comma." |
| |
| response = self.mcp_client.post( |
| "graph/find_join_path", |
| {"table1": tables[0], "table2": tables[1]} |
| ) |
| |
| if response.get("status") == "success": |
| path = response.get("path", "No path found") |
| return f"Join path: {path}" |
| else: |
| return f"Error finding join path: {response.get('message', 'Unknown error')}" |
| except Exception as e: |
| return f"Failed to find join path: {str(e)}" |
|
|
| async def _arun(self, table_names: str) -> str: |
| """Async version - just calls sync version.""" |
| return self._run(table_names) |
|
|
|
|
| class QueryExecutorTool(BaseTool): |
| """LangChain tool for executing SQL queries.""" |
| |
| name: str = "execute_query" |
| description: str = """ |
| Execute a SQL query against the databases and return results. |
| Use this after you have a valid SQL query. |
| Input should be a valid SQL query string. |
| """ |
| mcp_client: MCPClient |
|
|
| def _run(self, sql: str) -> str: |
| """Execute query.""" |
| try: |
| response = self.mcp_client.post( |
| "intelligence/execute_query", |
| {"sql": sql} |
| ) |
| |
| if response.get("status") == "success": |
| results = response.get("results", []) |
| |
| if results: |
| |
| result_text = f"Query returned {len(results)} rows:\\n" |
| headers = list(results[0].keys()) |
| result_text += " | ".join(headers) + "\n" |
| result_text += "-" * (len(" | ".join(headers))) + "\n" |
| |
| for row in results[:10]: |
| values = [str(row.get(h, "")) for h in headers] |
| result_text += " | ".join(values) + "\n" |
| |
| if len(results) > 10: |
| result_text += f"... and {len(results) - 10} more rows\n" |
| |
| return result_text |
| else: |
| return "Query executed successfully but returned no results." |
| else: |
| return f"Error executing query: {response.get('message', 'Unknown error')}" |
| except Exception as e: |
| return f"Failed to execute query: {str(e)}" |
|
|
| async def _arun(self, sql: str) -> str: |
| """Async version - just calls sync version.""" |
| return self._run(sql) |
|
|
