Building an AI Gateway from Scratch — From LLM Gateway to Agentic Gateway

import time
import uuid
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Callable

@dataclass
class ProviderConfig:
    name: str
    call_fn: Callable[[str, dict], str]   # actual API call wrapper
    cost_per_1k_tokens: float
    max_rpm: int = 60

class LLMGateway:
    """A minimal gateway: routing, retries, fallback, and usage tracking"""

    def __init__(self):
        self.providers: Dict[str, ProviderConfig] = {}
        self.usage_log: List[dict] = []
        self._request_counts: Dict[str, List[float]] = {}

    def register_provider(self, config: ProviderConfig):
        self.providers[config.name] = config
        self._request_counts[config.name] = []

    def _rate_limited(self, provider_name: str) -> bool:
        """Simple sliding-window rate limit check"""
        now = time.time()
        window = self._request_counts[provider_name]
        window[:] = [t for t in window if now - t < 60]
        return len(window) >= self.providers[provider_name].max_rpm

    def route(self, prompt: str, preferred: str, fallback_order: List[str], **kwargs) -> dict:
        """Try the preferred provider, fall back in order on failure or rate limit"""
        request_id = str(uuid.uuid4())
        order = [preferred] + [p for p in fallback_order if p != preferred]

        for provider_name in order:
            if provider_name not in self.providers:
                continue
            if self._rate_limited(provider_name):
                continue

            provider = self.providers[provider_name]
            try:
                self._request_counts[provider_name].append(time.time())
                response = provider.call_fn(prompt, kwargs)
                self._track_usage(request_id, provider_name, prompt, response)
                return {"request_id": request_id, "provider": provider_name, "response": response}
            except Exception as e:
                self._track_usage(request_id, provider_name, prompt, None, error=str(e))
                continue

        raise RuntimeError(f"All providers exhausted for request {request_id}")

    def _track_usage(self, request_id, provider_name, prompt, response, error=None):
        approx_tokens = len(prompt.split()) * 1.3
        cost = (approx_tokens / 1000) * self.providers[provider_name].cost_per_1k_tokens
        self.usage_log.append({
            "request_id": request_id,
            "provider": provider_name,
            "approx_tokens": approx_tokens,
            "approx_cost": cost,
            "error": error,
            "timestamp": time.time()
        })

# Usage
gateway = LLMGateway()

gateway.register_provider(ProviderConfig(
    name="anthropic",
    call_fn=lambda prompt, kw: f"[claude response to: {prompt}]",
    cost_per_1k_tokens=0.003,
    max_rpm=50
))
gateway.register_provider(ProviderConfig(
    name="openai",
    call_fn=lambda prompt, kw: f"[gpt response to: {prompt}]",
    cost_per_1k_tokens=0.005,
    max_rpm=60
))

result = gateway.route("Summarize this ticket", preferred="anthropic", fallback_order=["openai"])
print(result)

import numpy as np

class SemanticCache:
    def __init__(self, embed_fn, similarity_threshold: float = 0.92):
        self.embed_fn = embed_fn
        self.threshold = similarity_threshold
        self.entries: List[dict] = []  # {"embedding": vec, "prompt": str, "response": str}

    def _cosine_sim(self, a, b):
        return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

    def get(self, prompt: str) -> Optional[str]:
        query_vec = self.embed_fn(prompt)
        for entry in self.entries:
            if self._cosine_sim(query_vec, entry["embedding"]) >= self.threshold:
                return entry["response"]
        return None

    def set(self, prompt: str, response: str):
        self.entries.append({
            "embedding": self.embed_fn(prompt),
            "prompt": prompt,
            "response": response
        })

import re

class GuardrailHook:
    PII_PATTERNS = {
        "email": r"[\w\.-]+@[\w\.-]+\.\w+",
        "phone": r"\b\d{3}[-.]?\d{3}[-.]?\d{4}\b",
    }

    def redact(self, text: str) -> str:
        for label, pattern in self.PII_PATTERNS.items():
            text = re.sub(pattern, f"[REDACTED_{label.upper()}]", text)
        return text

    def pre_request(self, prompt: str) -> str:
        return self.redact(prompt)

    def post_response(self, response: str) -> str:
        return self.redact(response)

import json
import logging

logger = logging.getLogger("ai_gateway")

def log_request(request_id: str, provider: str, latency_ms: float, tokens: int, cost: float, status: str):
    logger.info(json.dumps({
        "request_id": request_id,
        "provider": provider,
        "latency_ms": latency_ms,
        "tokens": tokens,
        "cost": cost,
        "status": status
    }))

from litellm import completion

response = completion(
    model="claude-3-5-sonnet-20241022",
    messages=[{"role": "user", "content": "Summarize this ticket"}],
    fallbacks=["gpt-4o", "gemini/gemini-1.5-pro"]
)

print(response["choices"][0]["message"]["content"])

from portkey_ai import Portkey

client = Portkey(
    api_key="your-portkey-api-key",
    virtual_key="your-provider-virtual-key",
    config={
        "cache": {"mode": "semantic"},
        "retry": {"attempts": 3},
    }
)

response = client.chat.completions.create(
    model="gpt-4o",
    messages=[{"role": "user", "content": "Summarize this ticket"}]
)

from dataclasses import dataclass
from typing import Callable, Dict, List
import time
import json

@dataclass
class ToolPolicy:
    tool_name: str
    allowed_agents: List[str]
    requires_approval: bool = False

class MCPGatewayProxy:
    """Sits between an agent and its MCP servers, enforcing policy on every tool call"""

    def __init__(self):
        self.policies: Dict[str, ToolPolicy] = {}
        self.audit_log: List[dict] = []
        self.mcp_servers: Dict[str, Callable] = {}  # tool_name -> actual MCP call fn

    def register_tool(self, tool_name: str, mcp_call_fn: Callable, policy: ToolPolicy):
        self.mcp_servers[tool_name] = mcp_call_fn
        self.policies[tool_name] = policy

    def invoke_tool(self, agent_id: str, tool_name: str, arguments: dict) -> dict:
        policy = self.policies.get(tool_name)
        entry = {
            "timestamp": time.time(),
            "agent_id": agent_id,
            "tool_name": tool_name,
            "arguments": arguments,
        }

        if not policy:
            entry["status"] = "blocked_no_policy"
            self.audit_log.append(entry)
            raise PermissionError(f"No policy registered for tool '{tool_name}'")

        if agent_id not in policy.allowed_agents:
            entry["status"] = "blocked_unauthorized_agent"
            self.audit_log.append(entry)
            raise PermissionError(f"Agent '{agent_id}' is not authorized for tool '{tool_name}'")

        if policy.requires_approval:
            entry["status"] = "pending_approval"
            self.audit_log.append(entry)
            raise PermissionError(f"Tool '{tool_name}' requires human approval before execution")

        # Passed policy — relay to the actual MCP server
        try:
            result = self.mcp_servers[tool_name](arguments)
            entry["status"] = "executed"
            entry["result_summary"] = str(result)[:200]
            self.audit_log.append(entry)
            return {"status": "success", "result": result}
        except Exception as e:
            entry["status"] = "execution_error"
            entry["error"] = str(e)
            self.audit_log.append(entry)
            raise

# Usage
proxy = MCPGatewayProxy()

proxy.register_tool(
    tool_name="send_refund",
    mcp_call_fn=lambda args: f"Refund of {args['amount']} issued to {args['customer_id']}",
    policy=ToolPolicy(tool_name="send_refund", allowed_agents=["billing_agent"], requires_approval=True)
)

proxy.register_tool(
    tool_name="lookup_order",
    mcp_call_fn=lambda args: f"Order {args['order_id']}: shipped",
    policy=ToolPolicy(tool_name="lookup_order", allowed_agents=["billing_agent", "support_agent"])
)

# This succeeds — support_agent is allowed, no approval needed
print(proxy.invoke_tool("support_agent", "lookup_order", {"order_id": "A123"}))

# This raises — refunds require approval regardless of agent
try:
    proxy.invoke_tool("billing_agent", "send_refund", {"amount": 50, "customer_id": "C1"})
except PermissionError as e:
    print(f"Blocked: {e}")

                          AGENT REQUEST
                               │
                               ▼
                   ┌───────────────────────┐
                   │   1. Guardrail (pre)  │  redact PII in the prompt
                   └───────────┬───────────┘
                               ▼
                   ┌───────────────────────┐
                   │   2. Semantic Cache    │  hit? → skip straight to response
                   └───────────┬───────────┘
                        miss   │
                               ▼
                   ┌───────────────────────┐
                   │  3. LLM Router         │  pick provider, retry/fallback
                   │     (Part 1 Gateway)   │
                   └───────────┬───────────┘
                               ▼
                       model returns either
                       plain text OR a tool_use request
                               │
                 ┌─────────────┴─────────────┐
                 ▼                           ▼
        plain text response          tool_use requested
                 │                           │
                 │                           ▼
                 │              ┌───────────────────────┐
                 │              │ 4. MCP Gateway Proxy   │
                 │              │   - check policy       │
                 │              │   - log to audit trail │
                 │              │   - execute or block   │
                 │              └───────────┬───────────┘
                 │                           │
                 │                  tool result fed back
                 │                  into the LLM Router (step 3)
                 │                           │
                 └─────────────┬─────────────┘
                               ▼
                   ┌───────────────────────┐
                   │  5. Guardrail (post)   │  redact PII in the response
                   └───────────┬───────────┘
                               ▼
                   ┌───────────────────────┐
                   │  6. Cache + Log        │  store response, log usage/cost
                   └───────────┬───────────┘
                               ▼
                          FINAL RESPONSE

class AgenticGateway:
    """Composes LLMGateway + SemanticCache + GuardrailHook + MCPGatewayProxy
    into a single request path for agent traffic."""

    def __init__(self, llm_gateway: LLMGateway, cache: SemanticCache,
                 guardrail: GuardrailHook, mcp_proxy: MCPGatewayProxy):
        self.llm_gateway = llm_gateway
        self.cache = cache
        self.guardrail = guardrail
        self.mcp_proxy = mcp_proxy

    def handle(self, agent_id: str, prompt: str, preferred_provider: str,
               fallback_order: List[str], max_tool_hops: int = 3) -> dict:

        # 1. Guardrail — redact before anything leaves the building
        clean_prompt = self.guardrail.pre_request(prompt)

        # 2. Cache check
        cached = self.cache.get(clean_prompt)
        if cached:
            return {"source": "cache", "response": cached}

        # 3. LLM call with routing + fallback (Part 1)
        result = self.llm_gateway.route(clean_prompt, preferred_provider, fallback_order)
        model_output = result["response"]

        # 3a. Tool-call loop — bounded so a misbehaving agent can't loop forever
        hops = 0
        while self._is_tool_call(model_output) and hops < max_tool_hops:
            tool_name, arguments = self._parse_tool_call(model_output)

            # 4. MCP Gateway Proxy — policy check, execute, audit (Part 4)
            try:
                tool_result = self.mcp_proxy.invoke_tool(agent_id, tool_name, arguments)
            except PermissionError as e:
                tool_result = {"status": "blocked", "reason": str(e)}

            # Feed the tool result back through the LLM router as new context
            follow_up_prompt = f"{clean_prompt}\n\n[Tool result for {tool_name}]: {tool_result}"
            result = self.llm_gateway.route(follow_up_prompt, preferred_provider, fallback_order)
            model_output = result["response"]
            hops += 1

        # 5. Guardrail on the way out
        final_response = self.guardrail.post_response(model_output)

        # 6. Cache + log
        self.cache.set(clean_prompt, final_response)
        log_request(
            request_id=result["request_id"],
            provider=result["provider"],
            latency_ms=0,  # measure around llm_gateway.route in production
            tokens=len(final_response.split()),
            cost=0,        # pulled from llm_gateway.usage_log in production
            status="success"
        )

        return {"source": "live", "response": final_response, "tool_hops": hops}

    def _is_tool_call(self, model_output) -> bool:
        return isinstance(model_output, dict) and model_output.get("type") == "tool_use"

    def _parse_tool_call(self, model_output):
        return model_output["tool_name"], model_output["arguments"]

gateway = AgenticGateway(
    llm_gateway=gateway,          # LLMGateway instance from Part 1
    cache=SemanticCache(embed_fn=my_embed_fn),
    guardrail=GuardrailHook(),
    mcp_proxy=proxy                # MCPGatewayProxy instance from Part 4
)

response = gateway.handle(
    agent_id="support_agent",
    prompt="Look up order A123 and tell the customer its status",
    preferred_provider="anthropic",
    fallback_order=["openai"]
)
print(response)