firstock_docs_mcp/server.py

from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import Tool, TextContent, Prompt, PromptArgument, GetPromptResult, PromptMessage
from pydantic import BaseModel
import httpx
from bs4 import BeautifulSoup
import json
import logging
import urllib.parse

# Set up logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

class FirstockDocServer:
    def __init__(self):
        logger.debug("Initializing FirstockDocServer")
        self.base_url = "https://firstock.in/api/docs/"
        self.categories = {
            "overview": ["Introduction", "Versioning", "Key Generation", "Change Logs"],
            "downloader": ["Downloaders", "Library & SDK", "Data Types", "Request & Responses", "Exception & Errors"],
            "login": ["login", "logout", "User Details"],
            "orders": ["Place Order", "Order Margin", "Order Book", "Cancel Order", "Modify Order", 
                      "Single Order History", "Trade Book", "Position Book", "Product Conversion", 
                      "Holdings", "Limit", "Basket Margin"],
            "market": ["Brokerage Calculator", "Get Security Info", "Get Quote", "Get Quote Ltp",
                      "Get Multi Quote", "Get Multi Quote Ltp", "Index List", "Get Expiry",
                      "Option Chain", "Search Scrips", "Time Price Regular Interval", 
                      "Time Price Day Interval"]
        }
        logger.debug(f"Initialized with {len(self.categories)} categories")

    def _format_url(self, section: str) -> str:
        """Format the section name into a proper URL"""
        # Convert spaces to hyphens and lowercase
        formatted = section.lower().replace(" ", "-")
        # Remove any special characters
        formatted = ''.join(c for c in formatted if c.isalnum() or c == '-')
        return f"{self.base_url}{formatted}"

    async def fetch_documentation(self, section: str) -> str:
        """Fetch documentation for a specific section"""
        logger.debug(f"Fetching documentation for section: {section}")
        try:
            async with httpx.AsyncClient(follow_redirects=True) as client:
                url = self._format_url(section)
                logger.debug(f"Making request to: {url}")
                response = await client.get(url)
                if response.status_code == 200:
                    soup = BeautifulSoup(response.text, 'html.parser')
                    content = soup.find('main')
                    if content:
                        text = content.get_text()
                        logger.debug(f"Successfully fetched content, length: {len(text)}")
                        return text
                    logger.warning("No main content found in the page")
                    return "Documentation not found"
                logger.error(f"Error fetching documentation: {response.status_code}")
                return f"Error fetching documentation: {response.status_code}"
        except Exception as e:
            logger.error(f"Exception while fetching documentation: {str(e)}")
            return f"Error: {str(e)}"

    async def search_documentation(self, query: str) -> str:
        """Search through documentation"""
        logger.debug(f"Searching documentation for query: {query}")
        results = []
        for category, sections in self.categories.items():
            for section in sections:
                content = await self.fetch_documentation(section)
                if query.lower() in content.lower():
                    results.append(f"Found in {section}: {content[:200]}...")
        logger.debug(f"Search completed, found {len(results)} results")
        return "\n".join(results) if results else "No results found"

async def serve() -> None:
    logger.info("Starting Firstock Documentation MCP server")
    server = Server("firstock-docs")
    doc_server = FirstockDocServer()

    @server.list_tools()
    async def list_tools() -> list[Tool]:
        logger.debug("Listing available tools")
        return [
            Tool(
                name="get_documentation",
                description="Get documentation for a specific section of Firstock API",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "section": {
                            "type": "string",
                            "description": "Section of documentation to fetch (e.g., 'Introduction', 'Login', 'Place Order')"
                        }
                    },
                    "required": ["section"]
                }
            ),
            Tool(
                name="search_documentation",
                description="Search through Firstock API documentation",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "Search query to find in documentation"
                        }
                    },
                    "required": ["query"]
                }
            ),
            Tool(
                name="list_categories",
                description="List all available documentation categories",
                inputSchema={
                    "type": "object",
                    "properties": {}
                }
            )
        ]

    @server.list_prompts()
    async def list_prompts() -> list[Prompt]:
        logger.debug("Listing available prompts")
        return [
            Prompt(
                name="explore_docs",
                description="Explore Firstock API documentation",
                arguments=[
                    PromptArgument(
                        name="topic",
                        description="Topic to explore in documentation",
                        required=True
                    )
                ]
            )
        ]

    @server.call_tool()
    async def call_tool(name: str, arguments: dict) -> list[TextContent]:
        logger.debug(f"Calling tool: {name} with arguments: {arguments}")
        try:
            if name == "get_documentation":
                section = arguments["section"]
                content = await doc_server.fetch_documentation(section)
                return [TextContent(type="text", text=content)]
            
            elif name == "search_documentation":
                query = arguments["query"]
                results = await doc_server.search_documentation(query)
                return [TextContent(type="text", text=results)]
            
            elif name == "list_categories":
                categories = json.dumps(doc_server.categories, indent=2)
                return [TextContent(type="text", text=categories)]
            
            else:
                logger.error(f"Unknown tool: {name}")
                raise ValueError(f"Unknown tool: {name}")
                
        except Exception as e:
            logger.error(f"Error in call_tool: {str(e)}")
            return [TextContent(type="text", text=f"Error: {str(e)}")]

    @server.get_prompt()
    async def get_prompt(name: str, arguments: dict | None) -> GetPromptResult:
        logger.debug(f"Getting prompt: {name} with arguments: {arguments}")
        if name == "explore_docs":
            topic = arguments.get("topic", "")
            content = await doc_server.fetch_documentation(topic)
            return GetPromptResult(
                description=f"Documentation for {topic}",
                messages=[
                    PromptMessage(
                        role="user",
                        content=TextContent(type="text", text=content)
                    )
                ]
            )

    logger.info("Server setup complete, starting to serve")
    options = server.create_initialization_options()
    async with stdio_server() as (read_stream, write_stream):
        await server.run(read_stream, write_stream, options)