Glama
Azure AI Search MCP Server

"""Azure AI Agent Service MCP Server for Claude Desktop using Azure AI Search and Bing Web Grounding Tools."""

import os
import sys
import asyncio
from dotenv import load_dotenv
from mcp.server.fastmcp import FastMCP

# Import Azure AI Agent Service modules
from azure.ai.projects import AIProjectClient
from azure.ai.projects.models import AzureAISearchTool, BingGroundingTool, MessageRole
from azure.identity import DefaultAzureCredential

# Add startup message
print("Starting Azure AI Agent Service MCP Server...", file=sys.stderr)

# Load environment variables
load_dotenv()
print("Environment variables loaded", file=sys.stderr)

# Create MCP server
mcp = FastMCP(
    "azure-ai-agent", 
    description="MCP server for Azure AI Agent Service integration with AzureAISearch and Bing Web Grounding tools",
    dependencies=[
        "azure-identity",
        "python-dotenv",
        "azure-ai-projects"
    ]
)
print("MCP server instance created", file=sys.stderr)

class AzureAIAgentClient:
    """Client for Azure AI Agent Service with Azure AI Search and Bing Web Grounding tools."""
    
    def __init__(self):
        """Initialize Azure AI Agent Service client with credentials from environment variables."""
        print("Initializing Azure AI Agent client...", file=sys.stderr)
        
        # Load environment variables
        self.project_connection_string = os.getenv("PROJECT_CONNECTION_STRING")
        self.model_deployment_name = os.getenv("MODEL_DEPLOYMENT_NAME")
        self.search_connection_name = os.getenv("AI_SEARCH_CONNECTION_NAME")
        self.bing_connection_name = os.getenv("BING_CONNECTION_NAME")
        self.index_name = os.getenv("AI_SEARCH_INDEX_NAME")
        
        # Validate environment variables
        required_vars = {
            "PROJECT_CONNECTION_STRING": self.project_connection_string,
            "MODEL_DEPLOYMENT_NAME": self.model_deployment_name,
            "AI_SEARCH_CONNECTION_NAME": self.search_connection_name,
            "BING_CONNECTION_NAME": self.bing_connection_name,
            "AI_SEARCH_INDEX_NAME": self.index_name
        }
        
        missing = [k for k, v in required_vars.items() if not v]
        if missing:
            error_msg = f"Missing environment variables: {', '.join(missing)}"
            print(f"Error: {error_msg}", file=sys.stderr)
            raise ValueError(error_msg)
        
        # Initialize AIProjectClient
        try:
            self.client = AIProjectClient.from_connection_string(
                credential=DefaultAzureCredential(),
                conn_str=self.project_connection_string
            )
            print("AIProjectClient initialized successfully", file=sys.stderr)
        except Exception as e:
            print(f"Error initializing AIProjectClient: {str(e)}", file=sys.stderr)
            raise
        
        print(f"Azure AI Agent client initialized for AI Search connection: {self.search_connection_name}, Bing connection: {self.bing_connection_name}", file=sys.stderr)
    
    def search_index(self, query, top=5):
        """
        Perform a search using Azure AI Search Tool (default: best/hybrid mode).
        
        Args:
            query: The search query text
            top: Maximum number of results to return
            
        Returns:
            Formatted search results
        """
        print(f"Performing AI Search for: {query}", file=sys.stderr)
        
        try:
            # Get Azure AI Search connection
            search_connection = self.client.connections.get(connection_name=self.search_connection_name)
            if not search_connection:
                raise ValueError(f"Connection '{self.search_connection_name}' not found")
            
            # Create search tool
            search_tool = AzureAISearchTool(
                index_connection_id=search_connection.id,
                index_name=self.index_name
            )
            
            # Create agent with the search tool
            agent = self.client.agents.create_agent(
                model=self.model_deployment_name,
                name="search-agent",
                instructions=f"You are an Azure AI Search expert. Use the Azure AI Search Tool to find the most relevant information for: '{query}'. Return only the top {top} most relevant results. For each result, provide a title, content excerpt, and relevance score if available. Format your response as Markdown with each result clearly separated.",
                tools=search_tool.definitions,
                tool_resources=search_tool.resources,
                headers={"x-ms-enable-preview": "true"}
            )
            
            # Create thread for communication
            thread = self.client.agents.create_thread()
            
            # Create message to thread
            self.client.agents.create_message(
                thread_id=thread.id,
                role=MessageRole.USER,
                content=query
            )
            
            # Process the run
            run = self.client.agents.create_and_process_run(
                thread_id=thread.id,
                agent_id=agent.id
            )
            
            if run.status == "failed":
                print(f"Run failed: {run.last_error}", file=sys.stderr)
                return f"Search failed: {run.last_error}"
            
            # Get the agent's response
            response_message = self.client.agents.list_messages(thread_id=thread.id).get_last_message_by_role(
                MessageRole.AGENT
            )
            
            result = ""
            if response_message:
                for text_message in response_message.text_messages:
                    result += text_message.text.value + "\n"
                
                # Include any citations
                for annotation in response_message.url_citation_annotations:
                    result += f"\nCitation: [{annotation.url_citation.title}]({annotation.url_citation.url})\n"
            
            # Clean up resources
            self.client.agents.delete_agent(agent.id)
            
            return result
        
        except Exception as e:
            print(f"Error during search: {str(e)}", file=sys.stderr)
            raise
    
    def web_search(self, query):
        """
        Perform a web search using Bing Web Grounding Tool.
        
        Args:
            query: The search query text
            
        Returns:
            Formatted search results from the web
        """
        print(f"Performing Bing Web search for: {query}", file=sys.stderr)
        
        try:
            # Get Bing connection
            bing_connection = self.client.connections.get(connection_name=self.bing_connection_name)
            if not bing_connection:
                raise ValueError(f"Connection '{self.bing_connection_name}' not found")
            
            # Initialize Bing Web Grounding Tool
            bing_tool = BingGroundingTool(connection_id=bing_connection.id)
            
            # Create agent with the Bing tool
            agent = self.client.agents.create_agent(
                model=self.model_deployment_name,
                name="web-search-agent",
                instructions=f"You are a helpful web search assistant. Use the Bing Web Grounding Tool to find the most current and accurate information for: '{query}'. Provide a comprehensive answer with citations to sources. Format your response as Markdown.",
                tools=bing_tool.definitions,
                headers={"x-ms-enable-preview": "true"}
            )
            
            # Create thread for communication
            thread = self.client.agents.create_thread()
            
            # Create message to thread
            self.client.agents.create_message(
                thread_id=thread.id,
                role=MessageRole.USER,
                content=query
            )
            
            # Process the run
            run = self.client.agents.create_and_process_run(
                thread_id=thread.id,
                agent_id=agent.id
            )
            
            if run.status == "failed":
                print(f"Run failed: {run.last_error}", file=sys.stderr)
                return f"Web search failed: {run.last_error}"
            
            # Get the agent's response
            response_message = self.client.agents.list_messages(thread_id=thread.id).get_last_message_by_role(
                MessageRole.AGENT
            )
            
            result = ""
            if response_message:
                for text_message in response_message.text_messages:
                    result += text_message.text.value + "\n"
                
                # Include any citations
                for annotation in response_message.url_citation_annotations:
                    result += f"\nCitation: [{annotation.url_citation.title}]({annotation.url_citation.url})\n"
            
            # Clean up resources
            self.client.agents.delete_agent(agent.id)
            
            return result
        
        except Exception as e:
            print(f"Error during web search: {str(e)}", file=sys.stderr)
            raise

# Initialize Azure AI Agent client
try:
    print("Starting initialization of agent client...", file=sys.stderr)
    agent_client = AzureAIAgentClient()
    print("Agent client initialized successfully", file=sys.stderr)
except Exception as e:
    print(f"Error initializing agent client: {str(e)}", file=sys.stderr)
    # Don't exit - we'll handle errors in the tool functions
    agent_client = None

@mcp.tool()
def search_index(query: str, top: int = 5) -> str:
    """
    Search your Azure AI Search index using the optimal retrieval method.
    
    Args:
        query: The search query text
        top: Maximum number of results to return (default: 5)
    
    Returns:
        Formatted search results from your indexed documents
    """
    print(f"Tool called: search_index({query}, {top})", file=sys.stderr)
    if agent_client is None:
        return "Error: Azure AI Agent client is not initialized. Check server logs for details."
    
    try:
        results = agent_client.search_index(query, top)
        return f"## Azure AI Search Results\n\n{results}"
    except Exception as e:
        error_msg = f"Error performing index search: {str(e)}"
        print(error_msg, file=sys.stderr)
        return error_msg

@mcp.tool()
def web_search(query: str) -> str:
    """
    Search the web using Bing Web Grounding to find the most current information.
    
    Args:
        query: The search query text
    
    Returns:
        Formatted search results from the web with citations
    """
    print(f"Tool called: web_search({query})", file=sys.stderr)
    if agent_client is None:
        return "Error: Azure AI Agent client is not initialized. Check server logs for details."
    
    try:
        results = agent_client.web_search(query)
        return f"## Bing Web Search Results\n\n{results}"
    except Exception as e:
        error_msg = f"Error performing web 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()