Tidy up code and use a provider model with base

Tidy up code and use a provider model with base

Author: lukehinds

src/mockllm/providers/__init__.py

@@ -0,0 +1 @@
+# Empty init file to make this a package

src/mockllm/providers/anthropic.py

@@ -0,0 +1,68 @@
+from typing import Any, AsyncGenerator, Dict, Union
+
+from fastapi import HTTPException
+from fastapi.responses import StreamingResponse
+
+from ..config import ResponseConfig
+from ..models import (
+    AnthropicChatRequest,
+    AnthropicChatResponse,
+    AnthropicStreamDelta,
+    AnthropicStreamResponse,
+)
+from ..utils import count_tokens
+from .base import LLMProvider
+
+
+class AnthropicProvider(LLMProvider):
+    def __init__(self, response_config: ResponseConfig):
+        self.response_config = response_config
+
+    async def generate_stream_response(
+        self, content: str, model: str
+    ) -> AsyncGenerator[str, None]:
+        async for chunk in self.response_config.get_streaming_response_with_lag(
+            content
+        ):
+            stream_response = AnthropicStreamResponse(
+                delta=AnthropicStreamDelta(delta={"text": chunk})
+            )
+            yield f"data: {stream_response.model_dump_json()}\n\n"
+
+        yield "data: [DONE]\n\n"
+
+    async def handle_chat_completion(
+        self, request: AnthropicChatRequest
+    ) -> Union[Dict[str, Any], StreamingResponse]:
+        last_message = next(
+            (msg for msg in reversed(request.messages) if msg.role == "user"), None
+        )
+
+        if not last_message:
+            raise HTTPException(
+                status_code=400, detail="No user message found in request"
+            )
+
+        if request.stream:
+            return StreamingResponse(
+                self.generate_stream_response(last_message.content, request.model),
+                media_type="text/event-stream",
+            )
+
+        response_content = await self.response_config.get_response_with_lag(
+            last_message.content
+        )
+
+        prompt_tokens = count_tokens(str(request.messages), request.model)
+        completion_tokens = count_tokens(response_content, request.model)
+        total_tokens = prompt_tokens + completion_tokens
+
+        return AnthropicChatResponse(
+            model=request.model,
+            content=[{"type": "text", "text": response_content}],
+            usage={
+                "input_tokens": prompt_tokens,
+                "output_tokens": completion_tokens,
+                "total_tokens": total_tokens,
+            },
+        ).model_dump()

src/mockllm/providers/base.py

@@ -0,0 +1,19 @@
+from abc import ABC, abstractmethod
+from typing import Any, AsyncGenerator, Dict, Union
+
+from fastapi.responses import StreamingResponse
+
+
+class LLMProvider(ABC):
+    @abstractmethod
+    async def handle_chat_completion(
+        self, request: Any
+    ) -> Union[Dict[str, Any], StreamingResponse]:
+        pass
+
+    @abstractmethod
+    async def generate_stream_response(
+        self, content: str, model: str
+    ) -> AsyncGenerator[str, None]:
+        """Generate streaming response"""
+        yield ""  # pragma: no cover

src/mockllm/providers/openai.py

@@ -0,0 +1,89 @@
+from typing import Any, AsyncGenerator, Dict, Union
+
+from fastapi import HTTPException
+from fastapi.responses import StreamingResponse
+
+from ..config import ResponseConfig
+from ..models import (
+    OpenAIChatRequest,
+    OpenAIChatResponse,
+    OpenAIDeltaMessage,
+    OpenAIStreamChoice,
+    OpenAIStreamResponse,
+)
+from ..utils import count_tokens
+from .base import LLMProvider
+
+
+class OpenAIProvider(LLMProvider):
+    def __init__(self, response_config: ResponseConfig):
+        self.response_config = response_config
+
+    async def generate_stream_response(
+        self, content: str, model: str
+    ) -> AsyncGenerator[str, None]:
+        first_chunk = OpenAIStreamResponse(
+            model=model,
+            choices=[OpenAIStreamChoice(delta=OpenAIDeltaMessage(role="assistant"))],
+        )
+        yield f"data: {first_chunk.model_dump_json()}\n\n"
+
+        async for chunk in self.response_config.get_streaming_response_with_lag(
+            content
+        ):
+            chunk_response = OpenAIStreamResponse(
+                model=model,
+                choices=[OpenAIStreamChoice(delta=OpenAIDeltaMessage(content=chunk))],
+            )
+            yield f"data: {chunk_response.model_dump_json()}\n\n"
+
+        final_chunk = OpenAIStreamResponse(
+            model=model,
+            choices=[
+                OpenAIStreamChoice(delta=OpenAIDeltaMessage(), finish_reason="stop")
+            ],
+        )
+        yield f"data: {final_chunk.model_dump_json()}\n\n"
+        yield "data: [DONE]\n\n"
+
+    async def handle_chat_completion(
+        self, request: OpenAIChatRequest
+    ) -> Union[Dict[str, Any], StreamingResponse]:
+        last_message = next(
+            (msg for msg in reversed(request.messages) if msg.role == "user"), None
+        )
+
+        if not last_message:
+            raise HTTPException(
+                status_code=400, detail="No user message found in request"
+            )
+
+        if request.stream:
+            return StreamingResponse(
+                self.generate_stream_response(last_message.content, request.model),
+                media_type="text/event-stream",
+            )
+
+        response_content = await self.response_config.get_response_with_lag(
+            last_message.content
+        )
+
+        prompt_tokens = count_tokens(str(request.messages), request.model)
+        completion_tokens = count_tokens(response_content, request.model)
+        total_tokens = prompt_tokens + completion_tokens
+
+        return OpenAIChatResponse(
+            model=request.model,
+            choices=[
+                {
+                    "index": 0,
+                    "message": {"role": "assistant", "content": response_content},
+                    "finish_reason": "stop",
+                }
+            ],
+            usage={
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": completion_tokens,
+                "total_tokens": total_tokens,
+            },
+        ).model_dump()

src/mockllm/server.py

@@ -1,5 +1,5 @@
 import logging
-from typing import AsyncGenerator, Union
+from typing import Any, AsyncGenerator, Dict, Union
 
 import tiktoken
 from fastapi import FastAPI, HTTPException
@@ -9,15 +9,15 @@
 from .config import ResponseConfig
 from .models import (
     AnthropicChatRequest,
-    AnthropicChatResponse,
     AnthropicStreamDelta,
     AnthropicStreamResponse,
     OpenAIChatRequest,
-    OpenAIChatResponse,
     OpenAIDeltaMessage,
     OpenAIStreamChoice,
     OpenAIStreamResponse,
 )
+from .providers.anthropic import AnthropicProvider
+from .providers.openai import OpenAIProvider
 
 log_handler = logging.StreamHandler()
 log_handler.setFormatter(jsonlogger.JsonFormatter())
@@ -27,6 +27,8 @@
 app = FastAPI(title="Mock LLM Server")
 
 response_config = ResponseConfig()
+openai_provider = OpenAIProvider(response_config)
+anthropic_provider = AnthropicProvider(response_config)
 
 
 def count_tokens(text: str, model: str) -> int:
@@ -80,9 +82,8 @@ async def anthropic_stream_response(
 @app.post("/v1/chat/completions", response_model=None)
 async def openai_chat_completion(
     request: OpenAIChatRequest,
-) -> Union[OpenAIChatResponse, StreamingResponse]:
-    """Handle chat completion requests, supporting
-    both regular and streaming responses."""
+) -> Union[Dict[str, Any], StreamingResponse]:
+    """Handle OpenAI chat completion requests"""
     try:
         logger.info(
             "Received chat completion request",
@@ -92,47 +93,7 @@ async def openai_chat_completion(
                 "stream": request.stream,
             },
         )
-
-        last_message = next(
-            (msg for msg in reversed(request.messages) if msg.role == "user"), None
-        )
-
-        if not last_message:
-            raise HTTPException(
-                status_code=400, detail="No user message found in request"
-            )
-
-        if request.stream:
-            return StreamingResponse(
-                openai_stream_response(last_message.content, request.model),
-                media_type="text/event-stream",
-            )
-
-        response_content = await response_config.get_response_with_lag(
-            last_message.content
-        )
-
-        # Calculate mock token counts
-        prompt_tokens = count_tokens(str(request.messages), request.model)
-        completion_tokens = count_tokens(response_content, request.model)
-        total_tokens = prompt_tokens + completion_tokens
-
-        return OpenAIChatResponse(
-            model=request.model,
-            choices=[
-                {
-                    "index": 0,
-                    "message": {"role": "assistant", "content": response_content},
-                    "finish_reason": "stop",
-                }
-            ],
-            usage={
-                "prompt_tokens": prompt_tokens,
-                "completion_tokens": completion_tokens,
-                "total_tokens": total_tokens,
-            },
-        )
-
+        return await openai_provider.handle_chat_completion(request)
     except Exception as e:
         logger.error(f"Error processing request: {str(e)}")
         raise HTTPException(
@@ -143,9 +104,8 @@ async def openai_chat_completion(
 @app.post("/v1/messages", response_model=None)
 async def anthropic_chat_completion(
     request: AnthropicChatRequest,
-) -> Union[AnthropicChatResponse, StreamingResponse]:
-    """Handle Anthropic chat completion requests,
-    supporting both regular and streaming responses."""
+) -> Union[Dict[str, Any], StreamingResponse]:
+    """Handle Anthropic chat completion requests"""
     try:
         logger.info(
             "Received Anthropic chat completion request",
@@ -155,41 +115,7 @@ async def anthropic_chat_completion(
                 "stream": request.stream,
             },
         )
-
-        last_message = next(
-            (msg for msg in reversed(request.messages) if msg.role == "user"), None
-        )
-
-        if not last_message:
-            raise HTTPException(
-                status_code=400, detail="No user message found in request"
-            )
-
-        if request.stream:
-            return StreamingResponse(
-                anthropic_stream_response(last_message.content, request.model),
-                media_type="text/event-stream",
-            )
-
-        response_content = await response_config.get_response_with_lag(
-            last_message.content
-        )
-
-        # Calculate mock token counts
-        prompt_tokens = count_tokens(str(request.messages), request.model)
-        completion_tokens = count_tokens(response_content, request.model)
-        total_tokens = prompt_tokens + completion_tokens
-
-        return AnthropicChatResponse(
-            model=request.model,
-            content=[{"type": "text", "text": response_content}],
-            usage={
-                "input_tokens": prompt_tokens,
-                "output_tokens": completion_tokens,
-                "total_tokens": total_tokens,
-            },
-        )
-
+        return await anthropic_provider.handle_chat_completion(request)
     except Exception as e:
         logger.error(f"Error processing request: {str(e)}")
         raise HTTPException(

src/mockllm/utils.py

@@ -0,0 +1,11 @@
+import tiktoken
+
+
+def count_tokens(text: str, model: str) -> int:
+    """Get realistic token count for text using tiktoken"""
+    try:
+        encoding = tiktoken.encoding_for_model(model)
+        return len(encoding.encode(text))
+    except Exception:
+        # Fallback to rough estimation if model not supported
+        return len(text.split())