Azure AI Agent Service + Azure AI Search MCP Server

"""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()