Updated token counting and removed non streaming support

README.md

@@ -221,17 +221,9 @@ Extend `BaseProvider` in `providers/` to add support for other APIs:
 from providers.base import BaseProvider, ProviderConfig
 
 class MyProvider(BaseProvider):
-    async def complete(self, request):
-        # Make API call, return raw JSON
-        pass
-
     async def stream_response(self, request, input_tokens=0):
         # Yield Anthropic SSE format events
         pass
-
-    def convert_response(self, response_json, original_request):
-        # Convert to Anthropic response format
-        pass
 ```
 
 ### Adding Your Own Messaging App

api/models/anthropic.py

@@ -99,7 +99,7 @@ class MessagesRequest(BaseModel):
     messages: List[Message]
     system: Optional[Union[str, List[SystemContent]]] = None
     stop_sequences: Optional[List[str]] = None
-    stream: Optional[bool] = False
+    stream: Optional[bool] = True
     temperature: Optional[float] = None
     top_p: Optional[float] = None
     top_k: Optional[int] = None

api/request_utils.py

@@ -32,33 +32,78 @@ def get_token_count(
             total_tokens += len(ENCODER.encode(system))
         elif isinstance(system, list):
             for block in system:
-                if hasattr(block, "text"):
-                    total_tokens += len(ENCODER.encode(block.text))
+                text = (
+                    getattr(block, "text", None)
+                    if hasattr(block, "text")
+                    else (block.get("text", "") if isinstance(block, dict) else "")
+                )
+                if text:
+                    total_tokens += len(ENCODER.encode(text))
+        total_tokens += 4  # System block formatting overhead
 
     for msg in messages:
         if isinstance(msg.content, str):
             total_tokens += len(ENCODER.encode(msg.content))
         elif isinstance(msg.content, list):
             for block in msg.content:
-                b_type = getattr(block, "type", None)
+                b_type = getattr(block, "type", None) or (
+                    block.get("type") if isinstance(block, dict) else None
+                )
 
                 if b_type == "text":
-                    total_tokens += len(ENCODER.encode(getattr(block, "text", "")))
+                    text = getattr(block, "text", "") or (
+                        block.get("text", "") if isinstance(block, dict) else ""
+                    )
+                    total_tokens += len(ENCODER.encode(text))
                 elif b_type == "thinking":
-                    total_tokens += len(ENCODER.encode(getattr(block, "thinking", "")))
+                    thinking = getattr(block, "thinking", "") or (
+                        block.get("thinking", "") if isinstance(block, dict) else ""
+                    )
+                    total_tokens += len(ENCODER.encode(thinking))
                 elif b_type == "tool_use":
-                    name = getattr(block, "name", "")
-                    inp = getattr(block, "input", {})
+                    name = getattr(block, "name", "") or (
+                        block.get("name", "") if isinstance(block, dict) else ""
+                    )
+                    inp = getattr(block, "input", {}) or (
+                        block.get("input", {}) if isinstance(block, dict) else {}
+                    )
+                    block_id = getattr(block, "id", "") or (
+                        block.get("id", "") if isinstance(block, dict) else ""
+                    )
                     total_tokens += len(ENCODER.encode(name))
                     total_tokens += len(ENCODER.encode(json.dumps(inp)))
-                    total_tokens += 10
+                    total_tokens += len(ENCODER.encode(str(block_id)))
+                    total_tokens += 15
+                elif b_type == "image":
+                    source = getattr(block, "source", None) or (
+                        block.get("source", {}) if isinstance(block, dict) else {}
+                    )
+                    if isinstance(source, dict):
+                        data = source.get("data") or source.get("base64") or ""
+                        if data:
+                            total_tokens += max(85, len(data) // 3000)
+                        else:
+                            total_tokens += 765
+                    else:
+                        total_tokens += 765
                 elif b_type == "tool_result":
-                    content = getattr(block, "content", "")
+                    content = getattr(block, "content", "") or (
+                        block.get("content", "") if isinstance(block, dict) else ""
+                    )
+                    tool_use_id = getattr(block, "tool_use_id", "") or (
+                        block.get("tool_use_id", "") if isinstance(block, dict) else ""
+                    )
                     if isinstance(content, str):
                         total_tokens += len(ENCODER.encode(content))
                     else:
                         total_tokens += len(ENCODER.encode(json.dumps(content)))
-                    total_tokens += 5
+                    total_tokens += len(ENCODER.encode(str(tool_use_id)))
+                    total_tokens += 8
+                else:
+                    try:
+                        total_tokens += len(ENCODER.encode(json.dumps(block)))
+                    except (TypeError, ValueError):
+                        total_tokens += len(ENCODER.encode(str(block)))
 
     if tools:
         for tool in tools:
@@ -67,7 +112,7 @@ def get_token_count(
             )
             total_tokens += len(ENCODER.encode(tool_str))
 
-    total_tokens += len(messages) * 3
+    total_tokens += len(messages) * 4
     if tools:
         total_tokens += len(tools) * 5
 

api/routes.py

@@ -33,7 +33,7 @@ async def create_message(
     provider: BaseProvider = Depends(get_provider),
     settings: Settings = Depends(get_settings),
 ):
-    """Create a message (streaming or non-streaming)."""
+    """Create a message (always streaming)."""
 
     try:
         optimized = try_optimizations(request_data, settings)
@@ -43,22 +43,18 @@ async def create_message(
         request_id = f"req_{uuid.uuid4().hex[:12]}"
         log_request_compact(logger, request_id, request_data)
 
-        if request_data.stream:
-            input_tokens = get_token_count(
-                request_data.messages, request_data.system, request_data.tools
-            )
-            return StreamingResponse(
-                provider.stream_response(request_data, input_tokens=input_tokens),
-                media_type="text/event-stream",
-                headers={
-                    "X-Accel-Buffering": "no",
-                    "Cache-Control": "no-cache",
-                    "Connection": "keep-alive",
-                },
-            )
-        else:
-            response_json = await provider.complete(request_data)
-            return provider.convert_response(response_json, request_data)
+        input_tokens = get_token_count(
+            request_data.messages, request_data.system, request_data.tools
+        )
+        return StreamingResponse(
+            provider.stream_response(request_data, input_tokens=input_tokens),
+            media_type="text/event-stream",
+            headers={
+                "X-Accel-Buffering": "no",
+                "Cache-Control": "no-cache",
+                "Connection": "keep-alive",
+            },
+        )
 
     except ProviderError:
         raise

providers/base.py

@@ -28,11 +28,6 @@ class BaseProvider(ABC):
     def __init__(self, config: ProviderConfig):
         self.config = config
 
-    @abstractmethod
-    async def complete(self, request: Any) -> dict:
-        """Make a non-streaming completion request. Returns raw JSON response."""
-        pass
-
     @abstractmethod
     async def stream_response(
         self, request: Any, input_tokens: int = 0
@@ -40,8 +35,3 @@ class BaseProvider(ABC):
         """Stream response in Anthropic SSE format."""
         if False:
             yield ""
-
-    @abstractmethod
-    def convert_response(self, response_json: dict, original_request: Any) -> Any:
-        """Convert provider response to Anthropic format."""
-        pass

providers/nvidia_nim/client.py

@@ -10,7 +10,6 @@ from openai import AsyncOpenAI
 from providers.base import BaseProvider, ProviderConfig
 from providers.rate_limit import GlobalRateLimiter
 from .request import build_request_body
-from .response import convert_response
 from .errors import map_error
 from .utils import (
     SSEBuilder,
@@ -44,9 +43,9 @@ class NvidiaNimProvider(BaseProvider):
             timeout=300.0,
         )
 
-    def _build_request_body(self, request: Any, stream: bool = False) -> dict:
+    def _build_request_body(self, request: Any) -> dict:
         """Internal helper for tests and shared building."""
-        return build_request_body(request, self._nim_settings, stream=stream)
+        return build_request_body(request, self._nim_settings)
 
     async def stream_response(
         self, request: Any, input_tokens: int = 0
@@ -55,7 +54,7 @@ class NvidiaNimProvider(BaseProvider):
         message_id = f"msg_{uuid.uuid4()}"
         sse = SSEBuilder(message_id, request.model, input_tokens)
 
-        body = self._build_request_body(request, stream=True)
+        body = self._build_request_body(request)
         logger.info(
             f"NIM_STREAM: model={body.get('model')} msgs={len(body.get('messages', []))} tools={len(body.get('tools', []))}"
         )
@@ -219,31 +218,17 @@ class NvidiaNimProvider(BaseProvider):
             if usage_info and hasattr(usage_info, "completion_tokens")
             else sse.estimate_output_tokens()
         )
+        if usage_info and hasattr(usage_info, "prompt_tokens"):
+            provider_input = usage_info.prompt_tokens
+            if isinstance(provider_input, int):
+                diff = provider_input - input_tokens
+                logger.debug(
+                    f"TOKEN_ESTIMATE: our={input_tokens} provider={provider_input} diff={diff:+d}"
+                )
         yield sse.message_delta(map_stop_reason(finish_reason), output_tokens)
         yield sse.message_stop()
         yield sse.done()
 
-    async def complete(self, request: Any) -> dict:
-        """Make a non-streaming completion request."""
-        body = self._build_request_body(request, stream=False)
-        logger.info(
-            f"NIM_COMPLETE: model={body.get('model')} msgs={len(body.get('messages', []))} tools={len(body.get('tools', []))}"
-        )
-
-        try:
-            response = await self._global_rate_limiter.execute_with_retry(
-                self._client.chat.completions.create, **body
-            )
-            # Response converter expects a dict
-            return response.model_dump()
-        except Exception as e:
-            logger.error(f"NIM_ERROR: {type(e).__name__}: {e}")
-            raise map_error(e)
-
-    def convert_response(self, response_json: dict, original_request: Any) -> Any:
-        """Convert provider response to Anthropic format."""
-        return convert_response(response_json, original_request)
-
     def _process_tool_call(self, tc: dict, sse: Any):
         """Process a single tool call delta and yield SSE events.
 

providers/nvidia_nim/request.py

@@ -23,9 +23,7 @@ def _set_extra(
     extra_body[key] = value
 
 
-def build_request_body(
-    request_data: Any, nim: NimSettings, stream: bool = False
-) -> dict:
+def build_request_body(request_data: Any, nim: NimSettings) -> dict:
     """Build OpenAI-format request body from Anthropic request."""
     messages = AnthropicToOpenAIConverter.convert_messages(request_data.messages)
 

providers/nvidia_nim/utils/sse_builder.py

@@ -281,9 +281,17 @@ class SSEBuilder:
                 name = self.blocks.tool_names.get(idx, "")
                 tool_tokens += len(ENCODER.encode(name))
                 tool_tokens += len(ENCODER.encode(content))
-                tool_tokens += 10  # Control tokens overhead
+                tool_tokens += 15  # Control tokens overhead per tool
 
-            return text_tokens + reasoning_tokens + tool_tokens
+            # Per-block overhead (~4 tokens per content block)
+            block_count = (
+                (1 if self._accumulated_reasoning else 0)
+                + (1 if self._accumulated_text else 0)
+                + len(self.blocks.tool_indices)
+            )
+            block_overhead = block_count * 4
+
+            return text_tokens + reasoning_tokens + tool_tokens + block_overhead
 
         text_tokens = len(self._accumulated_text) // 4
         reasoning_tokens = len(self._accumulated_reasoning) // 4

tests/test_api.py

@@ -7,9 +7,19 @@ from providers.exceptions import APIError
 
 # Mock provider
 mock_provider = MagicMock(spec=NvidiaNimProvider)
-mock_provider.complete = AsyncMock()
-mock_provider.stream_response = AsyncMock()
-mock_provider.convert_response = MagicMock()
+
+# Track stream_response calls for test_model_mapping
+_stream_response_calls = []
+
+
+async def _mock_stream_response(*args, **kwargs):
+    """Minimal async generator for streaming tests."""
+    _stream_response_calls.append((args, kwargs))
+    yield "event: message_start\ndata: {}\n\n"
+    yield "[DONE]\n\n"
+
+
+mock_provider.stream_response = _mock_stream_response
 
 
 def override_get_provider():
@@ -32,39 +42,33 @@ def test_health():
     assert response.json()["status"] == "healthy"
 
 
-def test_create_message_non_stream():
-    mock_provider.complete.return_value = {"id": "123", "choices": []}
-    mock_provider.convert_response.return_value = {
-        "id": "msg_123",
-        "type": "message",
-        "role": "assistant",
-        "model": "test-model",
-        "content": [{"type": "text", "text": "Hello"}],
-        "usage": {"input_tokens": 10, "output_tokens": 5},
-    }
-
+def test_create_message_stream():
+    """Create message returns streaming response."""
     payload = {
         "model": "claude-3-sonnet",
         "messages": [{"role": "user", "content": "Hi"}],
         "max_tokens": 100,
-        "stream": False,
+        "stream": True,
     }
-
     response = client.post("/v1/messages", json=payload)
     assert response.status_code == 200
-    assert response.json()["content"][0]["text"] == "Hello"
-    mock_provider.complete.assert_called_once()
+    assert "text/event-stream" in response.headers.get("content-type", "")
+    content = b"".join(response.iter_bytes())
+    assert b"message_start" in content or b"event:" in content
 
 
 def test_model_mapping():
     # Test Haiku mapping
+    _stream_response_calls.clear()
     payload_haiku = {
         "model": "claude-3-haiku-20240307",
         "messages": [{"role": "user", "content": "Hi"}],
         "max_tokens": 100,
+        "stream": True,
     }
     client.post("/v1/messages", json=payload_haiku)
-    args, _ = mock_provider.complete.call_args
+    assert len(_stream_response_calls) == 1
+    args = _stream_response_calls[0][0]
     assert args[0].model != "claude-3-haiku-20240307"
     assert args[0].original_model == "claude-3-haiku-20240307"
 
@@ -76,65 +80,77 @@ def test_error_fallbacks():
         OverloadedError,
     )
 
+    base_payload = {"model": "test", "messages": [], "max_tokens": 10, "stream": True}
+
+    def _raise_auth(*args, **kwargs):
+        raise AuthenticationError("Invalid Key")
+
+    def _raise_rate_limit(*args, **kwargs):
+        raise RateLimitError("Too Many Requests")
+
+    def _raise_overloaded(*args, **kwargs):
+        raise OverloadedError("Server Overloaded")
+
     # 1. Authentication Error (401)
-    mock_provider.complete.side_effect = AuthenticationError("Invalid Key")
-    response = client.post(
-        "/v1/messages", json={"model": "test", "messages": [], "max_tokens": 10}
-    )
+    mock_provider.stream_response = _raise_auth
+    response = client.post("/v1/messages", json=base_payload)
     assert response.status_code == 401
     assert response.json()["error"]["type"] == "authentication_error"
 
     # 2. Rate Limit (429)
-    mock_provider.complete.side_effect = RateLimitError("Too Many Requests")
-    response = client.post(
-        "/v1/messages", json={"model": "test", "messages": [], "max_tokens": 10}
-    )
+    mock_provider.stream_response = _raise_rate_limit
+    response = client.post("/v1/messages", json=base_payload)
     assert response.status_code == 429
     assert response.json()["error"]["type"] == "rate_limit_error"
 
     # 3. Overloaded (529)
-    mock_provider.complete.side_effect = OverloadedError("Server Overloaded")
-    response = client.post(
-        "/v1/messages", json={"model": "test", "messages": [], "max_tokens": 10}
-    )
+    mock_provider.stream_response = _raise_overloaded
+    response = client.post("/v1/messages", json=base_payload)
     assert response.status_code == 529
     assert response.json()["error"]["type"] == "overloaded_error"
 
-    # Reset side_effect for subsequent tests
-    mock_provider.complete.side_effect = None
+    # Reset for subsequent tests
+    mock_provider.stream_response = _mock_stream_response
 
 
 def test_generic_exception_returns_500():
     """Non-ProviderError exceptions are caught and returned as HTTPException(500)."""
-    mock_provider.complete.side_effect = RuntimeError("unexpected crash")
+
+    def _raise_runtime(*args, **kwargs):
+        raise RuntimeError("unexpected crash")
+
+    mock_provider.stream_response = _raise_runtime
     response = client.post(
         "/v1/messages",
         json={
             "model": "test",
             "messages": [{"role": "user", "content": "Hi"}],
             "max_tokens": 10,
-            "stream": False,
+            "stream": True,
         },
     )
     assert response.status_code == 500
-    mock_provider.complete.side_effect = None
+    mock_provider.stream_response = _mock_stream_response
 
 
 def test_generic_exception_with_status_code():
     """Exception with status_code attribute uses that status."""
-    exc = APIError("bad gateway", status_code=502)
-    mock_provider.complete.side_effect = exc
+
+    def _raise_api_error(*args, **kwargs):
+        raise APIError("bad gateway", status_code=502)
+
+    mock_provider.stream_response = _raise_api_error
     response = client.post(
         "/v1/messages",
         json={
             "model": "test",
             "messages": [{"role": "user", "content": "Hi"}],
             "max_tokens": 10,
-            "stream": False,
+            "stream": True,
         },
     )
     assert response.status_code == 502
-    mock_provider.complete.side_effect = None
+    mock_provider.stream_response = _mock_stream_response
 
 
 def test_count_tokens_endpoint():

tests/test_nvidia_nim.py

@@ -65,7 +65,7 @@ async def test_init(provider_config):
 async def test_build_request_body(nim_provider):
     """Test request body construction."""
     req = MockRequest()
-    body = nim_provider._build_request_body(req, stream=True)
+    body = nim_provider._build_request_body(req)
 
     assert body["model"] == "test-model"
     assert body["temperature"] == 0.5
@@ -164,50 +164,6 @@ async def test_stream_response_thinking_reasoning_content(nim_provider):
         assert found_thinking
 
 
-@pytest.mark.asyncio
-async def test_complete_success(nim_provider):
-    """Test successful completion."""
-    req = MockRequest()
-
-    mock_response = MagicMock()
-    mock_response.model_dump.return_value = {
-        "id": "test_id",
-        "choices": [
-            {
-                "message": {"role": "assistant", "content": "Hello world"},
-                "finish_reason": "stop",
-            }
-        ],
-        "usage": {"prompt_tokens": 10, "completion_tokens": 5},
-    }
-
-    with patch.object(
-        nim_provider._client.chat.completions, "create", new_callable=AsyncMock
-    ) as mock_create:
-        mock_create.return_value = mock_response
-
-        result = await nim_provider.complete(req)
-        assert result["id"] == "test_id"
-        assert result["choices"][0]["message"]["content"] == "Hello world"
-
-
-@pytest.mark.asyncio
-async def test_complete_error_handling(nim_provider):
-    """Test error handling on completion."""
-    req = MockRequest()
-
-    import openai
-
-    with patch.object(
-        nim_provider._client.chat.completions,
-        "create",
-        side_effect=openai.APIError("API Error", request=MagicMock(), body=None),
-    ):
-        with pytest.raises(APIError) as exc:
-            await nim_provider.complete(req)
-        assert "API Error" in str(exc.value)
-
-
 @pytest.mark.asyncio
 async def test_tool_call_stream(nim_provider):
     """Test streaming tool calls."""

tests/test_request_utils.py

@@ -453,6 +453,100 @@ class TestGetTokenCount:
         # Double message should have more tokens (including overhead)
         assert count_double > count_single
 
+    def test_per_message_overhead_four_tokens(self):
+        """Per-message overhead is 4 tokens (was 3)."""
+        msg = MagicMock()
+        msg.content = "x"  # Minimal content
+        count = get_token_count([msg])
+        # 1 msg * 4 overhead + content tokens
+        assert count >= 5
+
+    def test_system_overhead_added(self):
+        """System prompt adds ~4 tokens overhead."""
+        msg = MagicMock()
+        msg.content = "Hi"
+        count_no_sys = get_token_count([msg])
+        count_with_sys = get_token_count([msg], system="You are helpful")
+        assert count_with_sys >= count_no_sys + 4
+
+    def test_system_as_list_of_dicts(self):
+        """System blocks as dicts (not objects) are counted."""
+        msg = MagicMock()
+        msg.content = "Hi"
+        count_no_sys = get_token_count([msg])
+        system_dicts = [{"type": "text", "text": "System prompt from dict"}]
+        count_with_dict_sys = get_token_count([msg], system=system_dicts)
+        assert count_with_dict_sys > count_no_sys
+
+    def test_tool_use_includes_id(self):
+        """Tool use blocks count id field."""
+        tool_block = MagicMock()
+        tool_block.type = "tool_use"
+        tool_block.name = "search"
+        tool_block.input = {"q": "test"}
+        tool_block.id = "call_abc123"
+        msg = MagicMock()
+        msg.content = [tool_block]
+        count = get_token_count([msg])
+        assert count > 0
+
+    def test_tool_result_includes_tool_use_id(self):
+        """Tool result blocks count tool_use_id field."""
+        result_block = MagicMock()
+        result_block.type = "tool_result"
+        result_block.content = "ok"
+        result_block.tool_use_id = "call_xyz"
+        msg = MagicMock()
+        msg.content = [result_block]
+        count = get_token_count([msg])
+        assert count > 0
+
+    def test_unrecognized_block_type_fallback(self):
+        """Unrecognized block types are tokenized via json.dumps fallback."""
+        unknown_block = {"type": "custom", "spec": "data"}
+        msg = MagicMock()
+        msg.content = [unknown_block]
+        count = get_token_count([msg])
+        assert count > 0
+
+    def test_message_with_image_block(self):
+        """Test token count includes image blocks."""
+        image_block = MagicMock()
+        image_block.type = "image"
+        image_block.source = {
+            "type": "base64",
+            "media_type": "image/png",
+            "data": "x" * 3000,
+        }
+        msg = MagicMock()
+        msg.content = [image_block]
+        count = get_token_count([msg])
+        assert count >= 85
+
+    def test_image_block_with_dict_source(self):
+        """Image block with dict-style source is counted."""
+        image_block = {"type": "image", "source": {"data": "a" * 10000}}
+        msg = MagicMock()
+        msg.content = [image_block]
+        count = get_token_count([msg])
+        assert count >= 85
+
+    def test_known_payload_estimate_range(self):
+        """Known payload produces estimate within expected range (validation harness)."""
+        import tiktoken
+
+        enc = tiktoken.get_encoding("cl100k_base")
+        system_text = "You are a helpful assistant."
+        user_text = "Hello, how are you?"
+        sys_tokens = len(enc.encode(system_text))
+        user_tokens = len(enc.encode(user_text))
+        # Min: content tokens + system overhead (4) + per-msg overhead (4)
+        expected_min = sys_tokens + user_tokens + 4 + 4
+        msg = MagicMock()
+        msg.content = user_text
+        count = get_token_count([msg], system=system_text)
+        assert count >= expected_min, f"count={count} < expected_min={expected_min}"
+
 
 # --- Parametrized Edge Case Tests ---