v0.3.1

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "token0"
-version = "0.3.0"
+version = "0.3.1"
 description = "Open-source API proxy that makes vision LLM calls 5-10x cheaper"
 readme = "README.md"
 license = "Apache-2.0"
@@ -21,9 +21,13 @@ dependencies = [
     "anthropic>=0.40.0",
     "openai>=1.50.0",
     "google-genai>=1.0.0",
+    "pypdf>=4.0.0",
 ]
 
 [project.optional-dependencies]
+langchain = [
+    "langchain-core>=0.2.0",
+]
 full = [
     "asyncpg>=0.30.0",
     "redis>=5.0.0",
@@ -35,6 +39,7 @@ dev = [
     "pytest-cov>=6.0.0",
     "ruff>=0.8.0",
     "mypy>=1.13.0",
+    "langchain-core>=0.2.0",
 ]
 
 [project.scripts]

tests/test_estimate.py

@@ -0,0 +1,149 @@
+"""Tests for the /v1/estimate endpoint."""
+
+import asyncio
+import base64
+import io
+
+from PIL import Image
+
+
+def _make_image_data_uri(width: int = 800, height: int = 600, fmt: str = "JPEG") -> str:
+    img = Image.new("RGB", (width, height), color=(100, 150, 200))
+    buf = io.BytesIO()
+    img.save(buf, format=fmt)
+    b64 = base64.b64encode(buf.getvalue()).decode()
+    mime = "image/jpeg" if fmt == "JPEG" else "image/png"
+    return f"data:{mime};base64,{b64}"
+
+
+class TestEstimateEndpoint:
+    def test_estimate_single_image(self):
+        from token0.api.v1.estimate import EstimateRequest, estimate
+        from token0.models.request import ContentPart, ImageUrl, Message
+
+        req = EstimateRequest(
+            model="gpt-4o",
+            messages=[
+                Message(
+                    role="user",
+                    content=[
+                        ContentPart(type="text", text="What's in this image?"),
+                        ContentPart(
+                            type="image_url",
+                            image_url=ImageUrl(url=_make_image_data_uri(800, 600)),
+                        ),
+                    ],
+                )
+            ],
+        )
+        result = asyncio.run(estimate(req))
+
+        assert result.model == "gpt-4o"
+        assert result.provider == "openai"
+        assert len(result.images) == 1
+        assert result.images[0].original_tokens > 0
+        assert result.total_original_tokens > 0
+
+    def test_estimate_returns_savings(self):
+        from token0.api.v1.estimate import EstimateRequest, estimate
+        from token0.models.request import ContentPart, ImageUrl, Message
+
+        # Large image — should trigger resize, yielding savings
+        req = EstimateRequest(
+            model="gpt-4o",
+            messages=[
+                Message(
+                    role="user",
+                    content=[
+                        ContentPart(
+                            type="image_url",
+                            image_url=ImageUrl(url=_make_image_data_uri(3000, 2000)),
+                        )
+                    ],
+                )
+            ],
+        )
+        result = asyncio.run(estimate(req))
+        assert result.total_original_tokens >= result.total_optimized_tokens
+
+    def test_estimate_text_only_no_images(self):
+        from token0.api.v1.estimate import EstimateRequest, estimate
+        from token0.models.request import Message
+
+        req = EstimateRequest(
+            model="gpt-4o",
+            messages=[Message(role="user", content="Just a text message")],
+        )
+        result = asyncio.run(estimate(req))
+        assert result.images == []
+        assert result.total_original_tokens == 0
+
+    def test_estimate_remote_url_skipped_with_note(self):
+        from token0.api.v1.estimate import EstimateRequest, estimate
+        from token0.models.request import ContentPart, ImageUrl, Message
+
+        req = EstimateRequest(
+            model="gpt-4o",
+            messages=[
+                Message(
+                    role="user",
+                    content=[
+                        ContentPart(
+                            type="image_url",
+                            image_url=ImageUrl(url="https://example.com/image.jpg"),
+                        )
+                    ],
+                )
+            ],
+        )
+        result = asyncio.run(estimate(req))
+        assert result.images == []
+        assert result.note is not None
+        assert "remote" in result.note.lower()
+
+    def test_estimate_multiple_images(self):
+        from token0.api.v1.estimate import EstimateRequest, estimate
+        from token0.models.request import ContentPart, ImageUrl, Message
+
+        req = EstimateRequest(
+            model="claude-sonnet-4-6",
+            messages=[
+                Message(
+                    role="user",
+                    content=[
+                        ContentPart(
+                            type="image_url",
+                            image_url=ImageUrl(url=_make_image_data_uri(800, 600)),
+                        ),
+                        ContentPart(
+                            type="image_url",
+                            image_url=ImageUrl(url=_make_image_data_uri(400, 300)),
+                        ),
+                    ],
+                )
+            ],
+        )
+        result = asyncio.run(estimate(req))
+        assert len(result.images) == 2
+        assert result.provider == "anthropic"
+
+    def test_estimate_cost_saved_is_non_negative(self):
+        from token0.api.v1.estimate import EstimateRequest, estimate
+        from token0.models.request import ContentPart, ImageUrl, Message
+
+        req = EstimateRequest(
+            model="gpt-4o",
+            messages=[
+                Message(
+                    role="user",
+                    content=[
+                        ContentPart(
+                            type="image_url",
+                            image_url=ImageUrl(url=_make_image_data_uri(100, 100)),
+                        )
+                    ],
+                )
+            ],
+        )
+        result = asyncio.run(estimate(req))
+        assert result.total_cost_saved_usd >= 0

tests/test_langchain_callback.py

@@ -0,0 +1,120 @@
+"""Tests for the LangChain callback handler."""
+
+import base64
+import io
+
+import pytest
+from PIL import Image
+
+pytest.importorskip("langchain_core", reason="langchain-core not installed")
+
+
+def _make_image_data_uri(width: int = 800, height: int = 600) -> str:
+    img = Image.new("RGB", (width, height), color=(100, 150, 200))
+    buf = io.BytesIO()
+    img.save(buf, format="JPEG")
+    b64 = base64.b64encode(buf.getvalue()).decode()
+    return f"data:image/jpeg;base64,{b64}"
+
+
+class TestToken0Callback:
+    def test_import(self):
+        from token0.langchain_callback import Token0Callback
+
+        cb = Token0Callback()
+        assert cb is not None
+
+    def test_init_defaults(self):
+        from token0.langchain_callback import Token0Callback
+
+        cb = Token0Callback()
+        assert cb.enable_cascade is False
+        assert cb.detail_override is None
+
+    def test_init_custom(self):
+        from token0.langchain_callback import Token0Callback
+
+        cb = Token0Callback(enable_cascade=True, detail_override="low")
+        assert cb.enable_cascade is True
+        assert cb.detail_override == "low"
+
+    def test_text_only_message_unchanged(self):
+        """Text-only messages should pass through without modification."""
+        from token0.langchain_callback import Token0Callback
+
+        try:
+            from langchain_core.messages import HumanMessage
+        except ImportError:
+            pytest.skip("langchain-core not installed")
+
+        cb = Token0Callback()
+        msg = HumanMessage(content="Hello, what is 2+2?")
+        original_content = msg.content
+
+        cb.on_chat_model_start(
+            serialized={"kwargs": {"model_name": "gpt-4o"}},
+            messages=[[msg]],
+        )
+
+        assert msg.content == original_content
+
+    def test_image_message_content_is_list(self):
+        """After optimization, image message content remains a list."""
+        from token0.langchain_callback import Token0Callback
+
+        try:
+            from langchain_core.messages import HumanMessage
+        except ImportError:
+            pytest.skip("langchain-core not installed")
+
+        cb = Token0Callback()
+        content = [
+            {"type": "text", "text": "Describe this image"},
+            {
+                "type": "image_url",
+                "image_url": {"url": _make_image_data_uri(800, 600)},
+            },
+        ]
+        msg = HumanMessage(content=content)
+
+        cb.on_chat_model_start(
+            serialized={"kwargs": {"model_name": "gpt-4o"}},
+            messages=[[msg]],
+        )
+
+        assert isinstance(msg.content, list)
+
+    def test_empty_serialized_does_not_crash(self):
+        """Missing model name in serialized should not crash."""
+        from token0.langchain_callback import Token0Callback
+
+        try:
+            from langchain_core.messages import HumanMessage
+        except ImportError:
+            pytest.skip("langchain-core not installed")
+
+        cb = Token0Callback()
+        msg = HumanMessage(content="hello")
+
+        # Should not raise
+        cb.on_chat_model_start(serialized={}, messages=[[msg]])
+
+    def test_extract_model_name(self):
+        from token0.langchain_callback import _extract_model_name
+
+        assert _extract_model_name({"kwargs": {"model_name": "gpt-4o"}}) == "gpt-4o"
+        model_id = "claude-sonnet-4-6"
+        assert _extract_model_name({"kwargs": {"model": model_id}}) == model_id
+        assert _extract_model_name({}) == ""
+
+    def test_role_for_messages(self):
+        from token0.langchain_callback import _role_for
+
+        try:
+            from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
+        except ImportError:
+            pytest.skip("langchain-core not installed")
+
+        assert _role_for(HumanMessage(content="hi")) == "user"
+        assert _role_for(AIMessage(content="hi")) == "assistant"
+        assert _role_for(SystemMessage(content="hi")) == "system"