Chris Coutinho
208365cd3d
Adds OpenAI provider to the unified provider architecture (ADR-015), supporting: - OpenAI API (api.openai.com) - GitHub Models API (models.github.ai/inference) - OpenAI-compatible endpoints (Fireworks, Together, etc.) Features: - Embedding support with text-embedding-3-small/large models - Text generation via chat completions API - Automatic retry with exponential backoff for rate limits - Provider auto-detection in registry (priority after Bedrock) Environment variables: - OPENAI_API_KEY: API key (required) - OPENAI_BASE_URL: Base URL override (optional) - OPENAI_EMBEDDING_MODEL: Embedding model (default: text-embedding-3-small) - OPENAI_GENERATION_MODEL: Generation model (default: gpt-4o-mini) Also adds: - Integration tests for RAG pipeline with MCP sampling - MCP client sampling support for integration tests - Ground truth Q&A pairs for Nextcloud User Manual 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
95 lines
3.1 KiB
Python
95 lines
3.1 KiB
Python
"""MCP sampling support for integration tests.
|
|
|
|
This module provides utilities to enable real LLM-based sampling in integration tests
|
|
using OpenAI or GitHub Models API.
|
|
"""
|
|
|
|
import logging
|
|
from typing import Any
|
|
|
|
from mcp import types
|
|
from mcp.client.session import ClientSession, RequestContext
|
|
|
|
from nextcloud_mcp_server.providers.openai import OpenAIProvider
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def create_openai_sampling_callback(provider: OpenAIProvider):
|
|
"""Factory to create a sampling callback using OpenAI provider.
|
|
|
|
The callback conforms to MCP's SamplingFnT protocol and can be passed
|
|
to ClientSession for handling sampling requests from the server.
|
|
|
|
Args:
|
|
provider: OpenAIProvider instance configured with a generation model
|
|
|
|
Returns:
|
|
Async callback function for MCP sampling
|
|
|
|
Example:
|
|
```python
|
|
provider = OpenAIProvider(
|
|
api_key=os.getenv("OPENAI_API_KEY"),
|
|
base_url=os.getenv("OPENAI_BASE_URL"),
|
|
generation_model="gpt-4o-mini",
|
|
)
|
|
callback = create_openai_sampling_callback(provider)
|
|
|
|
async for session in create_mcp_client_session(
|
|
url="http://localhost:8000/mcp",
|
|
sampling_callback=callback,
|
|
):
|
|
# Session now supports sampling
|
|
pass
|
|
```
|
|
"""
|
|
|
|
async def sampling_callback(
|
|
context: RequestContext[ClientSession, Any],
|
|
params: types.CreateMessageRequestParams,
|
|
) -> types.CreateMessageResult | types.ErrorData:
|
|
"""Handle sampling requests using OpenAI provider."""
|
|
logger.debug(f"Sampling callback invoked with {len(params.messages)} messages")
|
|
|
|
# Extract messages and build prompt
|
|
messages_text = []
|
|
for msg in params.messages:
|
|
if hasattr(msg.content, "text"):
|
|
role_prefix = "User" if msg.role == "user" else "Assistant"
|
|
messages_text.append(f"{role_prefix}: {msg.content.text}")
|
|
|
|
prompt = "\n\n".join(messages_text)
|
|
|
|
# Add system prompt if provided
|
|
if params.systemPrompt:
|
|
prompt = f"System: {params.systemPrompt}\n\n{prompt}"
|
|
|
|
logger.debug(f"Generating response for prompt ({len(prompt)} chars)")
|
|
|
|
try:
|
|
# Generate response using OpenAI provider
|
|
# Note: temperature is hardcoded in the provider at 0.7
|
|
response = await provider.generate(
|
|
prompt=prompt,
|
|
max_tokens=params.maxTokens,
|
|
)
|
|
|
|
model_name = provider.generation_model or "unknown"
|
|
logger.info(f"Sampling completed: {len(response)} chars from {model_name}")
|
|
|
|
return types.CreateMessageResult(
|
|
role="assistant",
|
|
content=types.TextContent(type="text", text=response),
|
|
model=model_name,
|
|
stopReason="endTurn",
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"OpenAI generation failed: {e}")
|
|
return types.ErrorData(
|
|
code=types.INTERNAL_ERROR,
|
|
message=f"OpenAI generation failed: {e!s}",
|
|
)
|
|
|
|
return sampling_callback
|