MCP Research Server (arXiv)

import asyncio
import json
import re
import ollama
from typing import Optional
from mcp import ClientSession
from mcp.client.stdio import stdio_client, StdioServerParameters

OLLAMA_MODEL = "mistral"


class MCPChatbot:
    def __init__(self):
        self.session: Optional[ClientSession] = None

    async def connect(self):
        server_params = StdioServerParameters(
            command="uv",
            args=["run", "research_server.py"],
        )

        async with stdio_client(server_params) as (read, write):
            async with ClientSession(read, write) as session:
                self.session = session
                await session.initialize()
                tools = await session.list_tools()
                print("Connected. Available tools:", [t.name for t in tools.tools])
                await self.chat_loop()

    async def chat_loop(self):
        while True:
            query = input("\nQuery (or 'quit'): ").strip()
            if query.lower() == "quit":
                break
            await self.handle_query(query)

    async def handle_query(self, query: str):
        # 🔒 CLIENT decides: is this a paper search?
        if self.is_search_request(query):
            topic, count = await self.extract_search_args(query)
            await self.run_search(topic, count)
        else:
            print("💬 Assistant: I can help you search for academic papers on arXiv. Try asking me to find papers on a topic.")

    def is_search_request(self, query: str) -> bool:
        keywords = ["search", "find", "papers", "research", "fetch"]
        return any(k in query.lower() for k in keywords)

    async def extract_search_args(self, query: str):
        # LLM used ONLY for argument extraction
        prompt = f"""
Extract search parameters from the user query.

Return ONLY valid JSON:
{{
  "topic": "<string>",
  "max_results": <integer>
}}

User query:
{query}
"""
        response = ollama.generate(
            model=OLLAMA_MODEL,
            prompt=prompt,
            stream=False,
        )

        try:
            data = json.loads(response["response"])
            topic = data.get("topic", "physics")
            max_results = int(data.get("max_results", 5))
        except Exception:
            # Safe fallback
            topic = "physics"
            max_results = 5

        return topic, max_results

    async def run_search(self, topic: str, max_results: int):
        print(f"Searching arXiv for '{topic}' ({max_results} papers)...")

        result = await self.session.call_tool(
            "search_papers",
            {"topic": topic, "max_results": max_results}
        )

        paper_ids = [
            c.text for c in result.content if hasattr(c, "text")
        ]

        if not paper_ids:
            print("Assistant: No papers found.")
            return

        print(" Papers found:")
        for pid in paper_ids:
            print(f" - {pid}")


async def main():
    bot = MCPChatbot()
    await bot.connect()


if __name__ == "__main__":
    asyncio.run(main())