From 4eaf8acfda14e22188d683d47488d0b40d2824c0 Mon Sep 17 00:00:00 2001 From: 7000pctAUTO Date: Thu, 5 Feb 2026 12:29:04 +0000 Subject: [PATCH] Add source files: core modules (auth, config, models) --- src/mcp_server_cli/auth.py | 243 +++++++++++++++++++++++++++++++++++++ 1 file changed, 243 insertions(+) create mode 100644 src/mcp_server_cli/auth.py diff --git a/src/mcp_server_cli/auth.py b/src/mcp_server_cli/auth.py new file mode 100644 index 0000000..ece6804 --- /dev/null +++ b/src/mcp_server_cli/auth.py @@ -0,0 +1,243 @@ +"""Authentication and local LLM configuration for MCP Server CLI.""" + +from typing import Any, Dict, List, Optional +from pydantic import BaseModel +import httpx +import json + +from mcp_server_cli.models import LocalLLMConfig + + +class LLMMessage(BaseModel): + """A message in an LLM conversation.""" + + role: str + content: str + + +class LLMChoice(BaseModel): + """A choice in an LLM response.""" + + index: int + message: LLMMessage + finish_reason: Optional[str] = None + + +class LLMResponse(BaseModel): + """Response from an LLM provider.""" + + id: str + object: str + created: int + model: str + choices: List[LLMChoice] + usage: Optional[Dict[str, Any]] = None + + +class ChatCompletionRequest(BaseModel): + """Request for chat completion.""" + + messages: List[Dict[str, str]] + model: str + temperature: Optional[float] = None + max_tokens: Optional[int] = None + stream: Optional[bool] = False + + +class LocalLLMClient: + """Client for interacting with local LLM providers.""" + + def __init__(self, config: LocalLLMConfig): + """Initialize the LLM client. + + Args: + config: Local LLM configuration. + """ + self.config = config + self.base_url = config.base_url.rstrip("/") + self.model = config.model + self.timeout = config.timeout + + async def chat_complete( + self, + messages: List[Dict[str, str]], + temperature: Optional[float] = None, + max_tokens: Optional[int] = None, + stream: bool = False, + ) -> LLMResponse: + """Send a chat completion request to the local LLM. + + Args: + messages: List of conversation messages. + temperature: Sampling temperature. + max_tokens: Maximum tokens to generate. + stream: Whether to stream the response. + + Returns: + LLM response with generated text. + """ + payload = { + "messages": messages, + "model": self.model, + "temperature": temperature or self.config.temperature, + "max_tokens": max_tokens or self.config.max_tokens, + "stream": stream, + } + + async with httpx.AsyncClient(timeout=self.timeout) as client: + response = await client.post( + f"{self.base_url}/v1/chat/completions", + json=payload, + headers={"Content-Type": "application/json"}, + ) + response.raise_for_status() + data = response.json() + + return LLMResponse( + id=data.get("id", "local-llm"), + object=data.get("object", "chat.completion"), + created=data.get("created", 0), + model=data.get("model", self.model), + choices=[ + LLMChoice( + index=choice.get("index", 0), + message=LLMMessage( + role=choice.get("message", {}).get("role", "assistant"), + content=choice.get("message", {}).get("content", ""), + ), + finish_reason=choice.get("finish_reason"), + ) + for choice in data.get("choices", []) + ], + usage=data.get("usage"), + ) + + async def stream_chat_complete( + self, + messages: List[Dict[str, str]], + temperature: Optional[float] = None, + max_tokens: Optional[int] = None, + ): + """Stream a chat completion from the local LLM. + + Args: + messages: List of conversation messages. + temperature: Sampling temperature. + max_tokens: Maximum tokens to generate. + + Yields: + Chunks of generated text. + """ + payload = { + "messages": messages, + "model": self.model, + "temperature": temperature or self.config.temperature, + "max_tokens": max_tokens or self.config.max_tokens, + "stream": True, + } + + async with httpx.AsyncClient(timeout=self.timeout) as client: + async with client.stream( + "POST", + f"{self.base_url}/v1/chat/completions", + json=payload, + headers={"Content-Type": "application/json"}, + ) as response: + async for line in response.aiter_lines(): + if line.startswith("data: "): + data = line[6:] + if data == "[DONE]": + break + try: + chunk = json.loads(data) + delta = chunk.get("choices", [{}])[0].get("delta", {}) + content = delta.get("content", "") + if content: + yield content + except json.JSONDecodeError: + continue + + async def test_connection(self) -> Dict[str, Any]: + """Test the connection to the local LLM. + + Returns: + Dictionary with connection status and details. + """ + try: + async with httpx.AsyncClient(timeout=10) as client: + response = await client.get(f"{self.base_url}/api/tags") + if response.status_code == 200: + return {"status": "connected", "details": response.json()} + except httpx.RequestError: + pass + + try: + async with httpx.AsyncClient(timeout=10) as client: + response = await client.get(f"{self.base_url}/v1/models") + if response.status_code == 200: + return {"status": "connected", "details": response.json()} + except httpx.RequestError: + pass + + return {"status": "failed", "error": "Could not connect to local LLM server"} + + +class LLMProviderRegistry: + """Registry for managing LLM providers.""" + + def __init__(self): + """Initialize the provider registry.""" + self._providers: Dict[str, LocalLLMClient] = {} + + def register(self, name: str, client: LocalLLMClient): + """Register an LLM provider. + + Args: + name: Provider name. + client: LLM client instance. + """ + self._providers[name] = client + + def get(self, name: str) -> Optional[LocalLLMClient]: + """Get an LLM provider by name. + + Args: + name: Provider name. + + Returns: + LLM client or None if not found. + """ + return self._providers.get(name) + + def list_providers(self) -> List[str]: + """List all registered provider names. + + Returns: + List of provider names. + """ + return list(self._providers.keys()) + + def create_default(self, config: LocalLLMConfig) -> LocalLLMClient: + """Create and register the default LLM provider. + + Args: + config: Local LLM configuration. + + Returns: + Created LLM client. + """ + client = LocalLLMClient(config) + self.register("default", client) + return client + + +def create_llm_client(config: LocalLLMConfig) -> LocalLLMClient: + """Create an LLM client from configuration. + + Args: + config: Local LLM configuration. + + Returns: + Configured LLM client. + """ + return LocalLLMClient(config)