Make tiktoken and PyJWT optional (saas group only)
This commit is contained in:
+22
-8
@@ -24,11 +24,24 @@ from starlette.middleware.base import BaseHTTPMiddleware
|
|||||||
# Import the proper create_app function that includes all middleware
|
# Import the proper create_app function that includes all middleware
|
||||||
from mcp_server_main import create_app
|
from mcp_server_main import create_app
|
||||||
|
|
||||||
# Import Stripe webhook router
|
# Conditional auth-related imports (only if auth enabled)
|
||||||
from stripe_webhook import router as stripe_router
|
_auth_check = os.getenv("ENABLE_AUTH", "false").lower() == "true"
|
||||||
|
|
||||||
# Import simplified MCP Auth HTTP adapter
|
if _auth_check:
|
||||||
from mcp_auth_http_simple import router as mcp_auth_router
|
# Import MCP Auth HTTP adapter (OAuth endpoints)
|
||||||
|
try:
|
||||||
|
from mcp_auth_http_simple import router as mcp_auth_router
|
||||||
|
except ImportError:
|
||||||
|
mcp_auth_router = None
|
||||||
|
|
||||||
|
# Import Stripe webhook router
|
||||||
|
try:
|
||||||
|
from stripe_webhook import router as stripe_router
|
||||||
|
except ImportError:
|
||||||
|
stripe_router = None
|
||||||
|
else:
|
||||||
|
mcp_auth_router = None
|
||||||
|
stripe_router = None
|
||||||
|
|
||||||
# OAuth configuration from environment variables
|
# OAuth configuration from environment variables
|
||||||
CLERK_ISSUER = os.getenv("CLERK_ISSUER", "https://clerk.yargimcp.com")
|
CLERK_ISSUER = os.getenv("CLERK_ISSUER", "https://clerk.yargimcp.com")
|
||||||
@@ -121,11 +134,12 @@ app = FastAPI(
|
|||||||
redirect_slashes=False # Disable to prevent 307 redirects on /mcp endpoint
|
redirect_slashes=False # Disable to prevent 307 redirects on /mcp endpoint
|
||||||
)
|
)
|
||||||
|
|
||||||
# Add Stripe webhook router to FastAPI
|
# Add auth-related routers to FastAPI (only if available)
|
||||||
app.include_router(stripe_router, prefix="/api/stripe")
|
if stripe_router:
|
||||||
|
app.include_router(stripe_router, prefix="/api/stripe")
|
||||||
|
|
||||||
# Add MCP Auth HTTP adapter to FastAPI (handles OAuth endpoints)
|
if mcp_auth_router:
|
||||||
app.include_router(mcp_auth_router)
|
app.include_router(mcp_auth_router)
|
||||||
|
|
||||||
# Custom 401 exception handler for MCP spec compliance
|
# Custom 401 exception handler for MCP spec compliance
|
||||||
@app.exception_handler(401)
|
@app.exception_handler(401)
|
||||||
|
|||||||
+19
-4
@@ -10,8 +10,15 @@ from collections import defaultdict
|
|||||||
from pydantic import HttpUrl, Field
|
from pydantic import HttpUrl, Field
|
||||||
from typing import Optional, Dict, List, Literal, Any, Union
|
from typing import Optional, Dict, List, Literal, Any, Union
|
||||||
import urllib.parse
|
import urllib.parse
|
||||||
import tiktoken
|
|
||||||
from fastmcp.server.middleware import Middleware, MiddlewareContext
|
from fastmcp.server.middleware import Middleware, MiddlewareContext
|
||||||
|
|
||||||
|
# Optional tiktoken import for token counting
|
||||||
|
try:
|
||||||
|
import tiktoken
|
||||||
|
TIKTOKEN_AVAILABLE = True
|
||||||
|
except ImportError:
|
||||||
|
TIKTOKEN_AVAILABLE = False
|
||||||
|
tiktoken = None
|
||||||
from fastmcp.server.dependencies import get_access_token, AccessToken
|
from fastmcp.server.dependencies import get_access_token, AccessToken
|
||||||
from fastmcp import Context
|
from fastmcp import Context
|
||||||
|
|
||||||
@@ -43,6 +50,9 @@ class TokenCountingMiddleware(Middleware):
|
|||||||
Args:
|
Args:
|
||||||
model: Tiktoken model name (cl100k_base for GPT-4/Claude compatibility)
|
model: Tiktoken model name (cl100k_base for GPT-4/Claude compatibility)
|
||||||
"""
|
"""
|
||||||
|
if not TIKTOKEN_AVAILABLE:
|
||||||
|
raise ImportError("tiktoken is required for token counting. Install with: pip install tiktoken")
|
||||||
|
|
||||||
self.encoder = tiktoken.get_encoding(model)
|
self.encoder = tiktoken.get_encoding(model)
|
||||||
self.model = model
|
self.model = model
|
||||||
self.token_stats = defaultdict(lambda: {"input": 0, "output": 0, "calls": 0})
|
self.token_stats = defaultdict(lambda: {"input": 0, "output": 0, "calls": 0})
|
||||||
@@ -228,9 +238,14 @@ def create_app(auth=None):
|
|||||||
else:
|
else:
|
||||||
logger.info("MCP server created with standard capabilities...")
|
logger.info("MCP server created with standard capabilities...")
|
||||||
|
|
||||||
token_counter = TokenCountingMiddleware()
|
# Add token counting middleware only if tiktoken is available
|
||||||
app.add_middleware(token_counter)
|
if TIKTOKEN_AVAILABLE:
|
||||||
logger.info("Token counting middleware added to MCP server")
|
try:
|
||||||
|
token_counter = TokenCountingMiddleware()
|
||||||
|
app.add_middleware(token_counter)
|
||||||
|
logger.info("Token counting middleware added to MCP server")
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(f"Failed to add token counting middleware: {e}")
|
||||||
|
|
||||||
return app
|
return app
|
||||||
|
|
||||||
|
|||||||
+1
-2
@@ -29,8 +29,6 @@ dependencies = [
|
|||||||
"fastmcp>=2.10.5",
|
"fastmcp>=2.10.5",
|
||||||
"pypdf>=5.5.0",
|
"pypdf>=5.5.0",
|
||||||
"fastapi>=0.115.14",
|
"fastapi>=0.115.14",
|
||||||
"PyJWT>=2.8.0",
|
|
||||||
"tiktoken>=0.5.0",
|
|
||||||
]
|
]
|
||||||
|
|
||||||
[project.optional-dependencies]
|
[project.optional-dependencies]
|
||||||
@@ -51,6 +49,7 @@ saas = [
|
|||||||
"stripe>=9.1.0",
|
"stripe>=9.1.0",
|
||||||
"upstash-redis>=1.1.0",
|
"upstash-redis>=1.1.0",
|
||||||
"tiktoken>=0.5.0",
|
"tiktoken>=0.5.0",
|
||||||
|
"PyJWT>=2.8.0",
|
||||||
]
|
]
|
||||||
|
|
||||||
[project.scripts]
|
[project.scripts]
|
||||||
|
|||||||
+1
-3
@@ -8,6 +8,4 @@ playwright
|
|||||||
pypdf
|
pypdf
|
||||||
fastapi>=0.115.14
|
fastapi>=0.115.14
|
||||||
uvicorn[standard]>=0.30.0
|
uvicorn[standard]>=0.30.0
|
||||||
starlette>=0.37.0
|
starlette>=0.37.0
|
||||||
tiktoken>=0.5.0
|
|
||||||
PyJWT>=2.8.0
|
|
||||||
Reference in New Issue
Block a user