Make tiktoken and PyJWT optional (saas group only)

This commit is contained in:
saidsurucu
2025-10-06 15:45:16 +03:00
parent 25723f070f
commit 3961a23d3a
5 changed files with 2528 additions and 17 deletions
+22 -8
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
Generated
+2485
View File
File diff suppressed because it is too large Load Diff