Azure AI Search MCP Server
by farzad528
"""Azure AI Search MCP Server for Claude Desktop."""
import os
import sys
from dotenv import load_dotenv
from azure.core.credentials import AzureKeyCredential
from azure.search.documents import SearchClient
from azure.search.documents.models import VectorizableTextQuery
from mcp.server.fastmcp import FastMCP
# Add startup message
print("Starting Azure AI Search MCP Server...", file=sys.stderr)
# Load environment variables
load_dotenv()
print("Environment variables loaded", file=sys.stderr)
# Create MCP server
mcp = FastMCP(
"azure-search",
description="MCP server for Azure AI Search integration",
dependencies=["azure-search-documents==11.5.2", "azure-identity", "python-dotenv"]
)
print("MCP server instance created", file=sys.stderr)
class AzureSearchClient:
"""Client for Azure AI Search service."""
def __init__(self):
"""Initialize Azure Search client with credentials from environment variables."""
print("Initializing Azure Search client...", file=sys.stderr)
# Load environment variables
self.endpoint = os.getenv("AZURE_SEARCH_SERVICE_ENDPOINT")
self.index_name = os.getenv("AZURE_SEARCH_INDEX_NAME") # Modified to use AZURE_SEARCH_INDEX
api_key = os.getenv("AZURE_SEARCH_API_KEY")
# Validate environment variables
if not all([self.endpoint, self.index_name, api_key]):
missing = []
if not self.endpoint:
missing.append("AZURE_SEARCH_SERVICE_ENDPOINT")
if not self.index_name:
missing.append("AZURE_SEARCH_INDEX")
if not api_key:
missing.append("AZURE_SEARCH_API_KEY")
error_msg = f"Missing environment variables: {', '.join(missing)}"
print(f"Error: {error_msg}", file=sys.stderr)
raise ValueError(error_msg)
# Initialize the search client
print(f"Connecting to Azure AI Search at {self.endpoint}", file=sys.stderr)
self.credential = AzureKeyCredential(api_key)
self.search_client = SearchClient(
endpoint=self.endpoint,
index_name=self.index_name,
credential=self.credential
)
print(f"Azure Search client initialized for index: {self.index_name}", file=sys.stderr)
def keyword_search(self, query, top=5):
"""Perform keyword search on the index."""
print(f"Performing keyword search for: {query}", file=sys.stderr)
results = self.search_client.search(
search_text=query,
top=top,
select=["title", "chunk"]
)
return self._format_results(results)
def vector_search(self, query, top=5, vector_field="text_vector"):
"""Perform vector search on the index."""
print(f"Performing vector search for: {query}", file=sys.stderr)
results = self.search_client.search(
vector_queries=[
VectorizableTextQuery(
text=query,
k_nearest_neighbors=50,
fields=vector_field
)
],
top=top,
select=["title", "chunk"]
)
return self._format_results(results)
def hybrid_search(self, query, top=5, vector_field="text_vector"):
"""Perform hybrid search (keyword + vector) on the index."""
print(f"Performing hybrid search for: {query}", file=sys.stderr)
results = self.search_client.search(
search_text=query,
vector_queries=[
VectorizableTextQuery(
text=query,
k_nearest_neighbors=50,
fields=vector_field
)
],
top=top,
select=["title", "chunk"]
)
return self._format_results(results)
def _format_results(self, results):
"""Format search results for better readability."""
formatted_results = []
for result in results:
item = {
"title": result.get("title", "Unknown"),
"content": result.get("chunk", "")[:1000], # Limit content length
"score": result.get("@search.score", 0)
}
formatted_results.append(item)
print(f"Formatted {len(formatted_results)} search results", file=sys.stderr)
return formatted_results
# Initialize Azure Search client
try:
print("Starting initialization of search client...", file=sys.stderr)
search_client = AzureSearchClient()
print("Search client initialized successfully", file=sys.stderr)
except Exception as e:
print(f"Error initializing search client: {str(e)}", file=sys.stderr)
# Don't exit - we'll handle errors in the tool functions
search_client = None
def _format_results_as_markdown(results, search_type):
"""Format search results as markdown for better readability."""
if not results:
return f"No results found for your query using {search_type}."
markdown = f"## {search_type} Results\n\n"
for i, result in enumerate(results, 1):
markdown += f"### {i}. {result['title']}\n"
markdown += f"Score: {result['score']:.2f}\n\n"
markdown += f"{result['content']}\n\n"
markdown += "---\n\n"
return markdown
@mcp.tool()
def keyword_search(query: str, top: int = 5) -> str:
"""
Perform a keyword-based search on the Azure AI Search index.
Args:
query: The search query text
top: Maximum number of results to return (default: 5)
Returns:
Formatted search results
"""
print(f"Tool called: keyword_search({query}, {top})", file=sys.stderr)
if search_client is None:
return "Error: Azure Search client is not initialized. Check server logs for details."
try:
results = search_client.keyword_search(query, top)
return _format_results_as_markdown(results, "Keyword Search")
except Exception as e:
error_msg = f"Error performing keyword search: {str(e)}"
print(error_msg, file=sys.stderr)
return error_msg
@mcp.tool()
def vector_search(query: str, top: int = 5) -> str:
"""
Perform a vector similarity search on the Azure AI Search index.
Args:
query: The search query text
top: Maximum number of results to return (default: 5)
Returns:
Formatted search results
"""
print(f"Tool called: vector_search({query}, {top})", file=sys.stderr)
if search_client is None:
return "Error: Azure Search client is not initialized. Check server logs for details."
try:
results = search_client.vector_search(query, top)
return _format_results_as_markdown(results, "Vector Search")
except Exception as e:
error_msg = f"Error performing vector search: {str(e)}"
print(error_msg, file=sys.stderr)
return error_msg
@mcp.tool()
def hybrid_search(query: str, top: int = 5) -> str:
"""
Perform a hybrid search (keyword + vector) on the Azure AI Search index.
Args:
query: The search query text
top: Maximum number of results to return (default: 5)
Returns:
Formatted search results
"""
print(f"Tool called: hybrid_search({query}, {top})", file=sys.stderr)
if search_client is None:
return "Error: Azure Search client is not initialized. Check server logs for details."
try:
results = search_client.hybrid_search(query, top)
return _format_results_as_markdown(results, "Hybrid Search")
except Exception as e:
error_msg = f"Error performing hybrid search: {str(e)}"
print(error_msg, file=sys.stderr)
return error_msg
if __name__ == "__main__":
# Run the server with stdio transport (default)
print("Starting MCP server run...", file=sys.stderr)
mcp.run()