kodit 0.5.3__tar.gz → 0.5.4__tar.gz

{kodit-0.5.3 → kodit-0.5.4}/Makefile

@@ -21,7 +21,7 @@ type:
 lint:
 	uv run ruff check --fix --unsafe-fixes
 
-test: lint type docs-check
+test: lint type
 	uv run pytest -s --cov=src --cov-report=xml tests/kodit
 
 no-database-changes-check:

{kodit-0.5.3 → kodit-0.5.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kodit
-Version: 0.5.3
+Version: 0.5.4
 Summary: Code indexing for better AI code generation
 Project-URL: Homepage, https://docs.helixml.tech/kodit/
 Project-URL: Documentation, https://docs.helixml.tech/kodit/

{kodit-0.5.3 → kodit-0.5.4}/docs/reference/api/index.md

@@ -12,7 +12,7 @@ look at the [hosted version](https://kodit.helix.ml/docs).
 This is the REST API for the Kodit server. Please refer to the
 [Kodit documentation](https://docs.helix.ml/kodit/) for more information.
     
-Current version: 0.5.2
+Current version: 0.5.3
 
 ## Authentication
 

{kodit-0.5.3 → kodit-0.5.4}/docs/reference/api/openapi.json

@@ -3,7 +3,7 @@
   "info": {
     "title": "kodit API",
     "description": "\nThis is the REST API for the Kodit server. Please refer to the\n[Kodit documentation](https://docs.helix.ml/kodit/) for more information.\n    ",
-    "version": "0.5.2"
+    "version": "0.5.3"
   },
   "paths": {
     "/healthz": {

{kodit-0.5.3 → kodit-0.5.4}/src/kodit/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.5.3'
-__version_tuple__ = version_tuple = (0, 5, 3)
+__version__ = version = '0.5.4'
+__version_tuple__ = version_tuple = (0, 5, 4)
 
 __commit_id__ = commit_id = None

{kodit-0.5.3 → kodit-0.5.4}/src/kodit/infrastructure/providers/litellm_provider.py

@@ -5,6 +5,7 @@ import functools
 from collections.abc import Callable, Coroutine
 from typing import Any
 
+import httpx
 import litellm
 import structlog
 from litellm import acompletion, aembedding
@@ -94,6 +95,20 @@ class LiteLLMProvider:
     def __init__(self, endpoint: Endpoint) -> None:
         """Initialize the LiteLLM provider."""
         self.endpoint = endpoint
+        self._setup_litellm_client()
+
+    def _setup_litellm_client(self) -> None:
+        """Set up LiteLLM with custom HTTPX client for Unix socket support."""
+        if self.endpoint.socket_path:
+            # Create HTTPX client with Unix socket transport
+            transport = httpx.AsyncHTTPTransport(uds=self.endpoint.socket_path)
+            unix_client = httpx.AsyncClient(
+                transport=transport,
+                base_url="http://localhost",  # Base URL for Unix socket
+                timeout=self.endpoint.timeout,
+            )
+            # Set as LiteLLM's async client session
+            litellm.aclient_session = unix_client
 
     def _populate_base_kwargs(self) -> dict[str, Any]:
         """Populate base kwargs common to all API calls."""
@@ -129,4 +144,11 @@ class LiteLLMProvider:
         return response.model_dump()
 
     async def close(self) -> None:
-        """Close the provider - litellm handles its own connection cleanup."""
+        """Close the provider and cleanup HTTPX client if using Unix sockets."""
+        if (
+            self.endpoint.socket_path
+            and hasattr(litellm, "aclient_session")
+            and litellm.aclient_session
+        ):
+            await litellm.aclient_session.aclose()
+            litellm.aclient_session = None

kodit-0.5.4/tests/kodit/infrastructure/providers/__init__.py

@@ -0,0 +1 @@
+"""Tests for infrastructure providers."""

kodit-0.5.4/tests/kodit/infrastructure/providers/litellm_provider_test.py

@@ -0,0 +1,273 @@
+"""Tests for the LiteLLM provider."""
+
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import httpx
+import litellm
+import pytest
+
+from kodit.config import Endpoint
+from kodit.infrastructure.providers.litellm_provider import LiteLLMProvider
+
+
+class TestLiteLLMProvider:
+    """Test the LiteLLM provider."""
+
+    def test_init_without_socket_path(self) -> None:
+        """Test initialization without Unix socket path."""
+        endpoint = Endpoint(model="gpt-4o-mini", api_key="test-key")
+        provider = LiteLLMProvider(endpoint)
+
+        assert provider.endpoint.model == "gpt-4o-mini"
+        assert provider.endpoint.api_key == "test-key"
+        assert provider.endpoint.socket_path is None
+
+    @patch("kodit.infrastructure.providers.litellm_provider.httpx.AsyncClient")
+    @patch("kodit.infrastructure.providers.litellm_provider.httpx.AsyncHTTPTransport")
+    def test_init_with_socket_path(
+        self, mock_transport: MagicMock, mock_client: MagicMock
+    ) -> None:
+        """Test initialization with Unix socket path configures httpx client."""
+        endpoint = Endpoint(
+            model="gpt-4o-mini",
+            socket_path="/var/run/llm.sock",
+            timeout=60.0,
+        )
+
+        # Create mock instances
+        mock_transport_instance = MagicMock()
+        mock_transport.return_value = mock_transport_instance
+        mock_client_instance = MagicMock()
+        mock_client.return_value = mock_client_instance
+
+        LiteLLMProvider(endpoint)
+
+        # Verify transport was created with Unix socket
+        mock_transport.assert_called_once_with(uds="/var/run/llm.sock")
+
+        # Verify httpx client was created with correct parameters
+        mock_client.assert_called_once_with(
+            transport=mock_transport_instance,
+            base_url="http://localhost",
+            timeout=60.0,
+        )
+
+        # Verify litellm's aclient_session was set
+        assert litellm.aclient_session == mock_client_instance
+
+    def test_socket_path_not_set_when_none(self) -> None:
+        """Test that socket setup is skipped when socket_path is None."""
+        endpoint = Endpoint(model="gpt-4o-mini")
+
+        # Store original session
+        original_session = getattr(litellm, "aclient_session", None)
+
+        try:
+            # Set to None to ensure clean state
+            litellm.aclient_session = None
+
+            provider = LiteLLMProvider(endpoint)
+
+            # Verify socket_path is None
+            assert provider.endpoint.socket_path is None
+
+            # Verify litellm.aclient_session was not set (still None)
+            assert litellm.aclient_session is None
+        finally:
+            # Restore original session
+            litellm.aclient_session = original_session
+
+    def test_populate_base_kwargs_with_all_params(self) -> None:
+        """Test _populate_base_kwargs includes all endpoint parameters."""
+        endpoint = Endpoint(
+            model="gpt-4o-mini",
+            api_key="test-key",
+            base_url="https://custom.openai.com",
+            timeout=120.0,
+            extra_params={"temperature": 0.7, "max_tokens": 1000},
+        )
+        provider = LiteLLMProvider(endpoint)
+
+        kwargs = provider._populate_base_kwargs()  # noqa: SLF001
+
+        assert kwargs["model"] == "gpt-4o-mini"
+        assert kwargs["timeout"] == 120.0
+        assert kwargs["api_key"] == "test-key"
+        assert kwargs["api_base"] == "https://custom.openai.com"
+        assert kwargs["temperature"] == 0.7
+        assert kwargs["max_tokens"] == 1000
+
+    def test_populate_base_kwargs_minimal(self) -> None:
+        """Test _populate_base_kwargs with minimal configuration."""
+        endpoint = Endpoint(model="gpt-4o-mini")
+        provider = LiteLLMProvider(endpoint)
+
+        kwargs = provider._populate_base_kwargs()  # noqa: SLF001
+
+        assert kwargs["model"] == "gpt-4o-mini"
+        assert kwargs["timeout"] == 60
+        assert "api_key" not in kwargs
+        assert "api_base" not in kwargs
+
+    @pytest.mark.asyncio
+    async def test_chat_completion_calls_acompletion(self) -> None:
+        """Test chat_completion calls litellm's acompletion with correct params."""
+        endpoint = Endpoint(model="gpt-4o-mini", api_key="test-key")
+        provider = LiteLLMProvider(endpoint)
+
+        messages = [
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": "Hello!"},
+        ]
+
+        # Mock the acompletion function
+        mock_response = MagicMock()
+        mock_response.model_dump.return_value = {
+            "choices": [{"message": {"content": "Hi there!"}}]
+        }
+
+        with patch(
+            "kodit.infrastructure.providers.litellm_provider.acompletion",
+            new_callable=AsyncMock,
+            return_value=mock_response,
+        ) as mock_acompletion:
+            result = await provider.chat_completion(messages)
+
+            # Verify acompletion was called with correct parameters
+            mock_acompletion.assert_called_once()
+            call_kwargs = mock_acompletion.call_args[1]
+            assert call_kwargs["model"] == "gpt-4o-mini"
+            assert call_kwargs["messages"] == messages
+            assert call_kwargs["api_key"] == "test-key"
+            assert call_kwargs["max_retries"] == 0
+
+            # Verify result
+            assert result == {"choices": [{"message": {"content": "Hi there!"}}]}
+
+    @pytest.mark.asyncio
+    async def test_embedding_calls_aembedding(self) -> None:
+        """Test embedding calls litellm's aembedding with correct params."""
+        endpoint = Endpoint(model="text-embedding-3-small", api_key="test-key")
+        provider = LiteLLMProvider(endpoint)
+
+        texts = ["Hello world", "Test text"]
+
+        # Mock the aembedding function
+        mock_response = MagicMock()
+        mock_response.model_dump.return_value = {
+            "data": [
+                {"embedding": [0.1, 0.2, 0.3]},
+                {"embedding": [0.4, 0.5, 0.6]},
+            ]
+        }
+
+        with patch(
+            "kodit.infrastructure.providers.litellm_provider.aembedding",
+            new_callable=AsyncMock,
+            return_value=mock_response,
+        ) as mock_aembedding:
+            result = await provider.embedding(texts)
+
+            # Verify aembedding was called with correct parameters
+            mock_aembedding.assert_called_once()
+            call_kwargs = mock_aembedding.call_args[1]
+            assert call_kwargs["model"] == "text-embedding-3-small"
+            assert call_kwargs["input"] == texts
+            assert call_kwargs["api_key"] == "test-key"
+            assert call_kwargs["max_retries"] == 0
+
+            # Verify result
+            assert result == {
+                "data": [
+                    {"embedding": [0.1, 0.2, 0.3]},
+                    {"embedding": [0.4, 0.5, 0.6]},
+                ]
+            }
+
+    @pytest.mark.asyncio
+    async def test_close_without_socket_path(self) -> None:
+        """Test close method when not using Unix sockets."""
+        endpoint = Endpoint(model="gpt-4o-mini")
+        provider = LiteLLMProvider(endpoint)
+
+        # Should complete without error
+        await provider.close()
+
+    @pytest.mark.asyncio
+    async def test_close_with_socket_path(self) -> None:
+        """Test close method cleans up httpx client when using Unix sockets."""
+        endpoint = Endpoint(
+            model="gpt-4o-mini",
+            socket_path="/var/run/llm.sock",
+        )
+
+        # Mock httpx client
+        mock_client = AsyncMock(spec=httpx.AsyncClient)
+        mock_client.aclose = AsyncMock()
+
+        with patch(
+            "kodit.infrastructure.providers.litellm_provider.httpx.AsyncClient",
+            return_value=mock_client,
+        ), patch(
+            "kodit.infrastructure.providers.litellm_provider.httpx.AsyncHTTPTransport"
+        ):
+            provider = LiteLLMProvider(endpoint)
+
+            # Verify client was set
+            assert litellm.aclient_session == mock_client
+
+            # Close the provider
+            await provider.close()
+
+            # Verify client was closed
+            mock_client.aclose.assert_called_once()
+
+            # Verify litellm.aclient_session was set to None
+            assert litellm.aclient_session is None
+
+    @pytest.mark.asyncio
+    @patch("kodit.infrastructure.providers.litellm_provider.httpx.AsyncClient")
+    @patch("kodit.infrastructure.providers.litellm_provider.httpx.AsyncHTTPTransport")
+    async def test_socket_path_end_to_end(
+        self, mock_transport: MagicMock, mock_client_class: MagicMock
+    ) -> None:
+        """Test Unix socket setup and cleanup end-to-end."""
+        # Create mock client instance
+        mock_client = AsyncMock()
+        mock_client.aclose = AsyncMock()
+        mock_client_class.return_value = mock_client
+
+        # Create mock transport
+        mock_transport_instance = MagicMock()
+        mock_transport.return_value = mock_transport_instance
+
+        endpoint = Endpoint(
+            model="gpt-4o-mini",
+            socket_path="/tmp/test.sock",
+            timeout=30.0,
+        )
+
+        # Initialize provider (should setup socket)
+        provider = LiteLLMProvider(endpoint)
+
+        # Verify transport was created
+        mock_transport.assert_called_once_with(uds="/tmp/test.sock")
+
+        # Verify client was created
+        mock_client_class.assert_called_once_with(
+            transport=mock_transport_instance,
+            base_url="http://localhost",
+            timeout=30.0,
+        )
+
+        # Verify litellm session was set
+        assert litellm.aclient_session == mock_client
+
+        # Close provider (should cleanup socket)
+        await provider.close()
+
+        # Verify client was closed
+        mock_client.aclose.assert_called_once()
+
+        # Verify session was cleared
+        assert litellm.aclient_session is None