uer-mcp 1.0.0

package/python/pyproject.toml

@@ -0,0 +1,127 @@
+[project]
+name = "uer"
+version = "0.1.0"
+description = "UER - Universal Expert Registry: Manipulation detection for AI systems"
+readme = "README.md"
+requires-python = ">=3.11"
+license = { text = "MIT" }
+authors = [
+    { name = "Margus Martsepp" },
+    { name = "Marco Lackovic" },
+    { name = "Martin Martsepp" },
+    { name = "Zane Estere" },
+    { name = "Anirudh" },
+]
+
+dependencies = [
+    "litellm>=1.40.0",
+    "datasets>=2.18.0",
+    "huggingface-hub>=0.20.0",
+    "mcp>=1.0.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "black>=24.1.1",
+    "ruff>=0.1.14",
+    "mypy>=1.8.0",
+    "pytest>=8.0.0",
+    "pytest-asyncio>=0.23.0",
+    "pytest-cov>=4.1.0",
+    "pre-commit>=3.6.0",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/uer"]
+
+[tool.black]
+line-length = 100
+target-version = ['py311']
+skip-string-normalization = false
+
+[tool.ruff]
+line-length = 100
+target-version = "py311"
+
+[tool.ruff.lint]
+select = [
+    "E",   # pycodestyle errors
+    "F",   # pyflakes
+    "I",   # isort
+    "N",   # pep8-naming
+    "W",   # pycodestyle warnings
+    "UP",  # pyupgrade
+    "B",   # flake8-bugbear
+    "C4",  # flake8-comprehensions
+    "SIM", # flake8-simplify
+]
+ignore = []
+
+[tool.ruff.lint.per-file-ignores]
+"__init__.py" = ["F401"]  # Allow unused imports in __init__.py
+"tests/**/*.py" = ["S101"]  # Allow assert statements in tests
+"context/scripts/*.py" = ["E501"]  # Allow long lines in test scripts
+"seed_datasets.py" = ["E501"]  # Allow long lines in setup script
+"src/uer/models/*.py" = ["E501", "N805"]  # Allow long lines in descriptions, Pydantic validators use cls
+"src/uer/server.py" = ["E501"]  # Allow long lines in tool descriptions
+
+[tool.mypy]
+python_version = "3.11"
+strict = true
+warn_return_any = true
+warn_unused_configs = true
+disallow_untyped_defs = true
+disallow_any_generics = true
+check_untyped_defs = true
+no_implicit_optional = true
+warn_redundant_casts = true
+warn_unused_ignores = true
+
+[[tool.mypy.overrides]]
+module = [
+    "datasets.*",
+    "huggingface_hub.*",
+    "litellm.*",
+    "mcp.*",
+    "pydantic.*",
+]
+ignore_missing_imports = true
+
+[[tool.mypy.overrides]]
+module = [
+    "seed_datasets",
+    "context.scripts.*",
+]
+disallow_untyped_defs = false
+check_untyped_defs = false
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+asyncio_mode = "auto"
+addopts = [
+    "--cov=.",
+    "--cov-report=term-missing",
+    "--cov-report=html",
+    "-v",
+]
+
+[tool.coverage.run]
+omit = [
+    "tests/*",
+    "context/*",
+    "*/site-packages/*",
+]
+
+[tool.coverage.report]
+exclude_lines = [
+    "pragma: no cover",
+    "def __repr__",
+    "raise AssertionError",
+    "raise NotImplementedError",
+    "if __name__ == .__main__.:",
+    "if TYPE_CHECKING:",
+]

package/python/src/uer/__init__.py

@@ -0,0 +1,17 @@
+"""
+UER - Universal Expert Registry.
+
+An MCP server providing universal LLM access via LiteLLM.
+"""
+
+__version__ = "0.1.0"
+
+# Export main components for programmatic use
+from uer.llm.gateway import LLMGateway
+from uer.models.llm import LLMCallRequest, LLMCallResponse
+
+__all__ = [
+    "LLMGateway",
+    "LLMCallRequest",
+    "LLMCallResponse",
+]

package/python/src/uer/__main__.py

@@ -0,0 +1,12 @@
+"""
+UER (Universal Expert Registry) MCP Server Entry Point.
+
+Run as: python -m uer.server
+"""
+
+import asyncio
+
+from uer.server import main
+
+if __name__ == "__main__":
+    asyncio.run(main())

package/python/src/uer/llm/gateway.py

@@ -0,0 +1,87 @@
+"""LiteLLM Gateway - Unified interface to 100+ LLM providers."""
+
+import os
+from typing import Any
+
+import litellm
+from litellm import acompletion
+
+
+class LLMGateway:
+    """Unified gateway to LLM providers via LiteLLM."""
+
+    def __init__(self) -> None:
+        """Initialize gateway and detect available providers."""
+        self.available_providers = self._detect_providers()
+
+        # Configure LiteLLM
+        litellm.set_verbose = False
+        litellm.drop_params = True  # Drop unsupported params instead of erroring
+
+    def _detect_providers(self) -> list[str]:
+        """Detect which LLM providers have API keys configured."""
+        providers = []
+
+        if os.getenv("GEMINI_API_KEY"):
+            providers.append("gemini")
+        if os.getenv("ANTHROPIC_API_KEY"):
+            providers.append("anthropic")
+        if os.getenv("OPENAI_API_KEY"):
+            providers.append("openai")
+        if os.getenv("AWS_ACCESS_KEY_ID"):
+            providers.append("bedrock")
+        if os.getenv("AZURE_API_KEY"):
+            providers.append("azure")
+
+        return providers
+
+    async def call(
+        self, model: str, messages: list[dict[str, str]], **kwargs: Any
+    ) -> dict[str, Any]:
+        """
+        Call any LLM via LiteLLM.
+
+        Args:
+            model: LiteLLM identifier (e.g., "gemini/gemini-3-flash-preview")
+            messages: List of chat messages with role and content
+            **kwargs: Additional parameters (temperature, max_tokens, response_format,
+                     thinking_level, thinking_budget, etc.)
+
+        Returns:
+            OpenAI-compatible response dict
+
+        Raises:
+            ValueError: Invalid model format
+            RuntimeError: API key missing or call failed
+        """
+        # Validate model format
+        if "/" not in model:
+            raise ValueError(
+                f"Invalid model format: '{model}'. "
+                f"Expected 'provider/model' (e.g., 'gemini/gemini-3-flash-preview')"
+            )
+
+        provider = model.split("/")[0]
+
+        # Check if provider is available
+        if provider not in self.available_providers:
+            raise RuntimeError(
+                f"Provider '{provider}' not available. "
+                f"Please set the appropriate API key environment variable. "
+                f"Available providers: {', '.join(self.available_providers) or 'none'}"
+            )
+
+        try:
+            # Call LiteLLM with all parameters
+            # Note: response_format, thinking_level, thinking_budget are passed through **kwargs
+            response = await acompletion(model=model, messages=messages, **kwargs)
+
+            # Convert to dict (LiteLLM returns ModelResponse object)
+            return response.model_dump()
+
+        except Exception as e:
+            raise RuntimeError(f"LLM call failed for model '{model}': {str(e)}") from e
+
+    def get_available_providers(self) -> list[str]:
+        """Get list of available providers."""
+        return self.available_providers

package/python/src/uer/models/llm.py

@@ -0,0 +1,109 @@
+"""Pydantic models for LLM call requests and responses."""
+
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field, validator
+
+
+class LLMCallRequest(BaseModel):
+    """Request for llm_call tool with validation."""
+
+    model: str = Field(
+        ..., description="LiteLLM model identifier (e.g., 'gemini/gemini-3-flash-preview')"
+    )
+    messages: list[dict[str, str]] = Field(
+        ..., min_length=1, description="List of chat messages with 'role' and 'content' keys"
+    )
+    temperature: float | None = Field(
+        default=0.7, ge=0.0, le=2.0, description="Sampling temperature"
+    )
+    max_tokens: int | None = Field(default=4096, ge=1, description="Maximum tokens to generate")
+
+    # Structured Output Support
+    response_format: dict[str, Any] | None = Field(
+        default=None,
+        description=(
+            "Force structured JSON output matching a schema. "
+            "Format: {'type': 'json_schema', 'json_schema': {'name': '...', 'schema': {...}, 'strict': True}}. "
+            "The calling LLM can generate schemas dynamically based on the task."
+        ),
+    )
+
+    # Chain of Thought Support (Gemini-specific)
+    thinking_level: Literal["minimal", "low", "medium", "high"] | None = Field(
+        default=None,
+        description=(
+            "Gemini 3 reasoning level. Use 'high' for complex tasks requiring deep reasoning, "
+            "'low' for simple tasks. Defaults to 'high' if not specified."
+        ),
+    )
+    thinking_budget: int | None = Field(
+        default=None,
+        ge=-1,
+        le=32768,
+        description=(
+            "Gemini 2.5 thinking tokens budget (128-32768, or -1 for dynamic). "
+            "Use -1 for automatic complexity-based allocation."
+        ),
+    )
+
+    # Tool Use Support
+    tools: list[dict[str, Any]] | None = Field(
+        default=None,
+        description=(
+            "List of tools the model can use in OpenAI format. "
+            "For Claude models, use built-in tools like 'web_search_20250305' for web search. "
+            "For Gemini, tools include code_execution, google_search_retrieval, etc. "
+            "Example: [{'type': 'web_search_20250305'}] for Claude web search."
+        ),
+    )
+
+    @validator("model")
+    def validate_model_format(cls, v: str) -> str:
+        """Ensure model is in 'provider/model' format."""
+        if "/" not in v:
+            raise ValueError(f"Model must be in format 'provider/model', got '{v}'")
+        return v
+
+    @validator("messages")
+    def validate_messages(cls, v: list[dict[str, str]]) -> list[dict[str, str]]:
+        """Ensure messages have required fields."""
+        for msg in v:
+            if "role" not in msg:
+                raise ValueError("Each message must have 'role' field")
+            if "content" not in msg:
+                raise ValueError("Each message must have 'content' field")
+        return v
+
+    @validator("response_format")
+    def validate_response_format(cls, v: dict[str, Any] | None) -> dict[str, Any] | None:
+        """Validate response_format structure."""
+        if v is not None:
+            if "type" not in v:
+                raise ValueError("response_format must have 'type' field")
+            if v["type"] == "json_schema" and "json_schema" not in v:
+                raise ValueError("json_schema response_format must have 'json_schema' field")
+        return v
+
+
+class LLMCallResponse(BaseModel):
+    """Response from llm_call tool."""
+
+    model: str
+    choices: list[dict[str, Any]]
+    usage: dict[str, int] | None = None
+
+    class Config:
+        """Pydantic config."""
+
+        json_schema_extra = {
+            "examples": [
+                {
+                    "model": "gemini/gemini-3-flash-preview",
+                    "choices": [
+                        {"message": {"role": "assistant", "content": "Hello! How can I help you?"}}
+                    ],
+                    "usage": {"total_tokens": 25},
+                }
+            ]
+        }

package/python/src/uer/server.py

@@ -0,0 +1,186 @@
+"""UER MCP Server - Universal Expert Registry."""
+
+import asyncio
+import json
+import logging
+from collections.abc import Sequence
+from typing import Any
+
+from mcp.server import Server
+from mcp.server.stdio import stdio_server
+from mcp.types import TextContent, Tool
+
+from uer.llm.gateway import LLMGateway
+from uer.models.llm import LLMCallRequest
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+)
+logger = logging.getLogger("uer.server")
+
+# Initialize MCP server
+app = Server("uer")
+gateway = LLMGateway()
+
+
+@app.list_tools()
+async def list_tools() -> list[Tool]:
+    """List available MCP tools."""
+    available = gateway.get_available_providers()
+
+    return [
+        Tool(
+            name="llm_call",
+            description=(
+                "Call any LLM via LiteLLM unified interface. "
+                f"Available providers: {', '.join(available) or 'none'}. "
+                "Supports Anthropic (Claude), OpenAI (GPT), Google (Gemini), and 100+ more. "
+                "Features: Structured output (response_format), Chain of Thought (thinking_level/thinking_budget), "
+                "Tool use (tools) for web search, code execution, grounding, etc."
+            ),
+            inputSchema={
+                "type": "object",
+                "properties": {
+                    "model": {
+                        "type": "string",
+                        "description": "LiteLLM model identifier (e.g., 'gemini/gemini-3-flash-preview')",
+                    },
+                    "messages": {
+                        "type": "array",
+                        "description": "List of chat messages",
+                        "items": {
+                            "type": "object",
+                            "properties": {
+                                "role": {
+                                    "type": "string",
+                                    "enum": ["system", "user", "assistant"],
+                                },
+                                "content": {"type": "string"},
+                            },
+                            "required": ["role", "content"],
+                        },
+                        "minItems": 1,
+                    },
+                    "temperature": {
+                        "type": "number",
+                        "default": 0.7,
+                        "minimum": 0,
+                        "maximum": 2,
+                    },
+                    "max_tokens": {
+                        "type": "integer",
+                        "default": 4096,
+                        "minimum": 1,
+                    },
+                    "response_format": {
+                        "type": "object",
+                        "description": (
+                            "Force structured JSON output. The calling LLM can generate schemas dynamically. "
+                            "Format: {'type': 'json_schema', 'json_schema': {'name': '...', 'schema': {...}, 'strict': True}}"
+                        ),
+                    },
+                    "thinking_level": {
+                        "type": "string",
+                        "enum": ["minimal", "low", "medium", "high"],
+                        "description": "Gemini 3 reasoning level: 'high' for complex reasoning tasks, 'low' for simple ones",
+                    },
+                    "thinking_budget": {
+                        "type": "integer",
+                        "minimum": -1,
+                        "maximum": 32768,
+                        "description": "Gemini 2.5 thinking tokens (128-32768, or -1 for dynamic allocation)",
+                    },
+                    "tools": {
+                        "type": "array",
+                        "description": (
+                            "List of tools the model can use. "
+                            "Claude: [{'type': 'web_search_20250305'}] for web search, [{'type': 'bash_20250305'}] for code execution. "
+                            "Gemini: [{'type': 'code_execution'}] or [{'type': 'google_search_retrieval'}]. "
+                            "OpenAI: Standard function calling format."
+                        ),
+                        "items": {"type": "object"},
+                    },
+                },
+                "required": ["model", "messages"],
+            },
+        )
+    ]
+
+
+@app.call_tool()
+async def call_tool(name: str, arguments: Any) -> Sequence[TextContent]:
+    """Handle tool invocation."""
+    if name != "llm_call":
+        raise ValueError(f"Unknown tool: {name}")
+
+    try:
+        # Validate input using Pydantic
+        request = LLMCallRequest(**arguments)
+
+        logger.info(
+            f"Calling LLM: model={request.model}, messages={len(request.messages)}, "
+            f"structured_output={request.response_format is not None}, "
+            f"thinking={request.thinking_level or request.thinking_budget}, "
+            f"tools={len(request.tools) if request.tools else 0}"
+        )
+
+        # Call LLM via gateway with all parameters
+        response = await gateway.call(
+            model=request.model,
+            messages=request.messages,
+            temperature=request.temperature,
+            max_tokens=request.max_tokens,
+            response_format=request.response_format,
+            thinking_level=request.thinking_level,
+            thinking_budget=request.thinking_budget,
+            tools=request.tools,
+        )
+
+        # Log result
+        usage = response.get("usage", {})
+        logger.info(
+            f"LLM call successful: tokens={usage.get('total_tokens', 'unknown')}, "
+            f"model={response.get('model', 'unknown')}"
+        )
+
+        # Return response as text content
+        return [TextContent(type="text", text=json.dumps(response, indent=2))]
+
+    except ValueError as e:
+        # Validation error from Pydantic
+        logger.error(f"Invalid input: {str(e)}")
+        return [
+            TextContent(type="text", text=json.dumps({"error": "Invalid input", "message": str(e)}))
+        ]
+
+    except RuntimeError as e:
+        # Error from LLMGateway
+        logger.error(f"LLM call failed: {str(e)}")
+        return [
+            TextContent(
+                type="text", text=json.dumps({"error": "LLM call failed", "message": str(e)})
+            )
+        ]
+
+    except Exception as e:
+        # Unexpected error
+        logger.exception(f"Unexpected error: {str(e)}")
+        return [
+            TextContent(
+                type="text", text=json.dumps({"error": "Internal error", "message": str(e)})
+            )
+        ]
+
+
+async def main() -> None:
+    """Run the MCP server via stdio transport."""
+    logger.info("Starting UER MCP server...")
+    logger.info(f"Available providers: {gateway.get_available_providers()}")
+
+    async with stdio_server() as (read_stream, write_stream):
+        await app.run(read_stream, write_stream, app.create_initialization_options())
+
+
+if __name__ == "__main__":
+    asyncio.run(main())