kodit 0.3.16__py3-none-any.whl → 0.3.17__py3-none-any.whl

kodit/cli_utils.py

@@ -0,0 +1,74 @@
+"""Utilities for CLI commands with remote/local mode support."""
+
+from collections.abc import Callable
+from functools import wraps
+from typing import TYPE_CHECKING, Any
+
+import click
+
+from kodit.infrastructure.api.client import IndexClient, SearchClient
+
+if TYPE_CHECKING:
+    from kodit.config import AppContext
+
+
+def with_client(f: Callable) -> Callable:
+    """Provide appropriate client based on configuration.
+
+    This decorator automatically detects whether to run in local or remote mode
+    based on the presence of remote.server_url in the configuration. In remote
+    mode, it provides API clients. In local mode, it behaves like the existing
+    with_session decorator.
+    """
+
+    @wraps(f)
+    async def wrapper(*args: Any, **kwargs: Any) -> Any:
+        ctx = click.get_current_context()
+        app_context: AppContext = ctx.obj
+
+        # Auto-detect mode based on remote.server_url presence
+        if not app_context.is_remote:
+            # Local mode - use existing database session approach
+            from kodit.config import with_session
+
+            # Apply the session decorator to the original function
+            session_wrapped = with_session(f)
+            # Remove the async wrapper that with_session adds since we're already async
+            inner_func = getattr(
+                getattr(session_wrapped, "__wrapped__", session_wrapped),
+                "__wrapped__",
+                session_wrapped
+            )
+
+            # Get database session manually
+            db = await app_context.get_db()
+            async with db.session_factory() as session:
+                return await inner_func(session, *args, **kwargs)
+        else:
+            # Remote mode - use API clients
+            clients = {
+                "index_client": IndexClient(
+                    base_url=app_context.remote.server_url or "",
+                    api_key=app_context.remote.api_key,
+                    timeout=app_context.remote.timeout,
+                    max_retries=app_context.remote.max_retries,
+                    verify_ssl=app_context.remote.verify_ssl,
+                ),
+                "search_client": SearchClient(
+                    base_url=app_context.remote.server_url or "",
+                    api_key=app_context.remote.api_key,
+                    timeout=app_context.remote.timeout,
+                    max_retries=app_context.remote.max_retries,
+                    verify_ssl=app_context.remote.verify_ssl,
+                ),
+            }
+
+            try:
+                # Pass clients to the command function
+                return await f(*args, clients=clients, **kwargs)
+            finally:
+                # Clean up client connections
+                for client in clients.values():
+                    await client.close()
+
+    return wrapper

kodit/config.py

@@ -111,6 +111,20 @@ class PeriodicSyncConfig(BaseModel):
     )
 
 
+class RemoteConfig(BaseModel):
+    """Configuration for remote server connection."""
+
+    server_url: str | None = Field(
+        default=None, description="Remote Kodit server URL"
+    )
+    api_key: str | None = Field(default=None, description="API key for authentication")
+    timeout: float = Field(default=30.0, description="Request timeout in seconds")
+    max_retries: int = Field(default=3, description="Maximum retry attempts")
+    verify_ssl: bool = Field(
+        default=True, description="Verify SSL certificates"
+    )
+
+
 class CustomAutoIndexingEnvSource(EnvSettingsSource):
     """Custom environment source for parsing AutoIndexingConfig."""
 
@@ -212,6 +226,9 @@ class AppContext(BaseSettings):
         default_factory=list,
         description="Comma-separated list of valid API keys (e.g. 'key1,key2')",
     )
+    remote: RemoteConfig = Field(
+        default_factory=RemoteConfig, description="Remote server configuration"
+    )
 
     @field_validator("api_keys", mode="before")
     @classmethod
@@ -234,6 +251,11 @@ class AppContext(BaseSettings):
         # Call this to ensure the data dir exists for the default db location
         self.get_data_dir()
 
+    @property
+    def is_remote(self) -> bool:
+        """Check if running in remote mode."""
+        return self.remote.server_url is not None
+
     def get_data_dir(self) -> Path:
         """Get the data directory."""
         self.data_dir.mkdir(parents=True, exist_ok=True)
@@ -248,11 +270,19 @@ class AppContext(BaseSettings):
     async def get_db(self, *, run_migrations: bool = True) -> Database:
         """Get the database."""
         if self._db is None:
-            self._db = Database(self.db_url)
-        if run_migrations:
-            await self._db.run_migrations(self.db_url)
+            self._db = await self.new_db(run_migrations=run_migrations)
         return self._db
 
+    async def new_db(self, *, run_migrations: bool = True) -> Database:
+        """Get a completely fresh connection to a database.
+
+        This is required when running tasks in a thread pool.
+        """
+        db = Database(self.db_url)
+        if run_migrations:
+            await db.run_migrations(self.db_url)
+        return db
+
 
 with_app_context = click.make_pass_decorator(AppContext)
 

kodit/domain/entities.py

@@ -4,16 +4,18 @@ import shutil
 from dataclasses import dataclass
 from datetime import datetime
 from pathlib import Path
-from typing import Protocol
+from typing import Any, Protocol
 from urllib.parse import urlparse, urlunparse
 
 from pydantic import AnyUrl, BaseModel
 
 from kodit.domain.value_objects import (
     FileProcessingStatus,
+    QueuePriority,
     SnippetContent,
     SnippetContentType,
     SourceType,
+    TaskType,
 )
 from kodit.utils.path_utils import path_from_uri
 
@@ -274,3 +276,48 @@ class SnippetWithContext:
     file: File
     authors: list[Author]
     snippet: Snippet
+
+
+class Task(BaseModel):
+    """Represents an item in the queue waiting to be processed.
+
+    If the item exists, that means it is in the queue and waiting to be processed. There
+    is no status associated.
+    """
+
+    id: str  # Is a unique key to deduplicate items in the queue
+    type: TaskType  # Task type
+    priority: int  # Priority (higher number = higher priority)
+    payload: dict[str, Any]  # Task-specific data
+
+    created_at: datetime | None = None  # Is populated by repository
+    updated_at: datetime | None = None  # Is populated by repository
+
+    @staticmethod
+    def create(task_type: TaskType, priority: int, payload: dict[str, Any]) -> "Task":
+        """Create a task."""
+        return Task(
+            id=Task._create_id(task_type, payload),
+            type=task_type,
+            priority=priority,
+            payload=payload,
+        )
+
+    @staticmethod
+    def _create_id(task_type: TaskType, payload: dict[str, Any]) -> str:
+        """Create a unique id for a task."""
+        if task_type == TaskType.INDEX_UPDATE:
+            return str(payload["index_id"])
+
+        raise ValueError(f"Unknown task type: {task_type}")
+
+    @staticmethod
+    def create_index_update_task(
+        index_id: int, priority: QueuePriority = QueuePriority.USER_INITIATED
+    ) -> "Task":
+        """Create an index update task."""
+        return Task.create(
+            task_type=TaskType.INDEX_UPDATE,
+            priority=priority.value,
+            payload={"index_id": index_id},
+        )

kodit/domain/protocols.py

@@ -5,8 +5,35 @@ from typing import Protocol
 
 from pydantic import AnyUrl
 
-from kodit.domain.entities import Index, Snippet, SnippetWithContext, WorkingCopy
-from kodit.domain.value_objects import MultiSearchRequest
+from kodit.domain.entities import Index, Snippet, SnippetWithContext, Task, WorkingCopy
+from kodit.domain.value_objects import MultiSearchRequest, TaskType
+
+
+class TaskRepository(Protocol):
+    """Repository interface for Task entities."""
+
+    async def add(
+        self,
+        task: Task,
+    ) -> None:
+        """Add a task."""
+        ...
+
+    async def get(self, task_id: str) -> Task | None:
+        """Get a task by ID."""
+        ...
+
+    async def take(self) -> Task | None:
+        """Take a task for processing."""
+        ...
+
+    async def update(self, task: Task) -> None:
+        """Update a task."""
+        ...
+
+    async def list(self, task_type: TaskType | None = None) -> list[Task]:
+        """List tasks with optional status filter."""
+        ...
 
 
 class IndexRepository(Protocol):

kodit/domain/value_objects.py

@@ -649,3 +649,16 @@ class FunctionDefinition:
     qualified_name: str
     start_byte: int
     end_byte: int
+
+
+class TaskType(Enum):
+    """Task type."""
+
+    INDEX_UPDATE = 1
+
+
+class QueuePriority(IntEnum):
+    """Queue priority."""
+
+    BACKGROUND = 10
+    USER_INITIATED = 50

kodit/infrastructure/api/client/__init__.py

@@ -0,0 +1,14 @@
+"""API client for remote Kodit server communication."""
+
+from .base import BaseAPIClient
+from .exceptions import AuthenticationError, KoditAPIError
+from .index_client import IndexClient
+from .search_client import SearchClient
+
+__all__ = [
+    "AuthenticationError",
+    "BaseAPIClient",
+    "IndexClient",
+    "KoditAPIError",
+    "SearchClient",
+]

kodit/infrastructure/api/client/base.py

@@ -0,0 +1,100 @@
+"""Base HTTP client for Kodit API communication."""
+
+import asyncio
+from typing import Any
+
+import httpx
+
+from .exceptions import AuthenticationError, KoditAPIError
+
+
+class BaseAPIClient:
+    """Base API client with authentication and retry logic."""
+
+    def __init__(
+        self,
+        base_url: str,
+        api_key: str | None = None,
+        timeout: float = 30.0,
+        max_retries: int = 3,
+        *,
+        verify_ssl: bool = True,
+    ) -> None:
+        """Initialize the API client.
+
+        Args:
+            base_url: Base URL of the Kodit server
+            api_key: API key for authentication
+            timeout: Request timeout in seconds
+            max_retries: Maximum retry attempts
+            verify_ssl: Whether to verify SSL certificates
+
+        """
+        self.base_url = base_url.rstrip("/")
+        self.api_key = api_key
+        self.timeout = timeout
+        self.max_retries = max_retries
+        self.verify_ssl = verify_ssl
+        self._client = self._create_client()
+
+    def _create_client(self) -> httpx.AsyncClient:
+        """Create the HTTP client with proper configuration."""
+        headers = {}
+        if self.api_key:
+            headers["X-API-Key"] = self.api_key
+
+        return httpx.AsyncClient(
+            base_url=self.base_url,
+            headers=headers,
+            timeout=httpx.Timeout(self.timeout),
+            verify=self.verify_ssl,
+            follow_redirects=True,
+        )
+
+    async def _request(
+        self,
+        method: str,
+        path: str,
+        **kwargs: Any,
+    ) -> httpx.Response:
+        """Make HTTP request with retry logic.
+
+        Args:
+            method: HTTP method (GET, POST, etc.)
+            path: API endpoint path
+            **kwargs: Additional arguments passed to httpx
+
+        Returns:
+            HTTP response object
+
+        Raises:
+            AuthenticationError: If authentication fails
+            KoditAPIError: For other API errors
+
+        """
+        url = f"{self.base_url}{path}"
+
+        for attempt in range(self.max_retries):
+            try:
+                response = await self._client.request(method, url, **kwargs)
+                response.raise_for_status()
+            except httpx.HTTPStatusError as e:
+                if e.response.status_code == 401:
+                    raise AuthenticationError("Invalid API key") from e
+                if e.response.status_code >= 500 and attempt < self.max_retries - 1:
+                    await asyncio.sleep(2**attempt)  # Exponential backoff
+                    continue
+                raise KoditAPIError(f"API request failed: {e}") from e
+            except httpx.RequestError as e:
+                if attempt < self.max_retries - 1:
+                    await asyncio.sleep(2**attempt)
+                    continue
+                raise KoditAPIError(f"Connection error: {e}") from e
+            else:
+                return response
+
+        raise KoditAPIError(f"Max retries ({self.max_retries}) exceeded")
+
+    async def close(self) -> None:
+        """Close the HTTP client."""
+        await self._client.aclose()

kodit/infrastructure/api/client/exceptions.py

@@ -0,0 +1,21 @@
+"""Exceptions for Kodit API client operations."""
+
+
+class KoditAPIError(Exception):
+    """Base exception for Kodit API errors."""
+
+
+
+class AuthenticationError(KoditAPIError):
+    """Authentication failed with the API server."""
+
+
+
+class KoditConnectionError(KoditAPIError):
+    """Connection to API server failed."""
+
+
+
+class ServerError(KoditAPIError):
+    """Server returned an error response."""
+

kodit/infrastructure/api/client/generated_endpoints.py

@@ -0,0 +1,27 @@
+"""API endpoint constants generated from OpenAPI specification.
+
+This file is auto-generated. Do not edit manually.
+Run `make generate-api-paths` to regenerate.
+"""
+
+from typing import Final
+
+
+class APIEndpoints:
+    """API endpoint constants extracted from OpenAPI specification."""
+
+    # /api/v1/indexes
+    API_V1_INDEXES: Final[str] = "/api/v1/indexes"
+
+    # /api/v1/indexes/{index_id}
+    API_V1_INDEXES_INDEX_ID: Final[str] = "/api/v1/indexes/{index_id}"
+
+    # /api/v1/search
+    API_V1_SEARCH: Final[str] = "/api/v1/search"
+
+    # /healthz
+    HEALTHZ: Final[str] = "/healthz"
+
+
+# Generated from: openapi.json
+# Total endpoints: 4

kodit/infrastructure/api/client/index_client.py

@@ -0,0 +1,57 @@
+"""Index operations API client for Kodit server."""
+
+from kodit.infrastructure.api.v1.schemas.index import (
+    IndexCreateAttributes,
+    IndexCreateData,
+    IndexCreateRequest,
+    IndexData,
+    IndexListResponse,
+    IndexResponse,
+)
+
+from .base import BaseAPIClient
+from .exceptions import KoditAPIError
+from .generated_endpoints import APIEndpoints
+
+
+class IndexClient(BaseAPIClient):
+    """API client for index operations."""
+
+    async def list_indexes(self) -> list[IndexData]:
+        """List all indexes."""
+        response = await self._request("GET", APIEndpoints.API_V1_INDEXES)
+        data = IndexListResponse.model_validate_json(response.text)
+        return data.data
+
+    async def create_index(self, uri: str) -> IndexData:
+        """Create a new index."""
+        request = IndexCreateRequest(
+            data=IndexCreateData(
+                type="index", attributes=IndexCreateAttributes(uri=uri)
+            )
+        )
+        response = await self._request(
+            "POST", APIEndpoints.API_V1_INDEXES, json=request.model_dump()
+        )
+        result = IndexResponse.model_validate_json(response.text)
+        return result.data
+
+    async def get_index(self, index_id: str) -> IndexData | None:
+        """Get index by ID."""
+        try:
+            response = await self._request(
+                "GET", APIEndpoints.API_V1_INDEXES_INDEX_ID.format(index_id=index_id)
+            )
+            result = IndexResponse.model_validate_json(response.text)
+        except KoditAPIError as e:
+            if "404" in str(e):
+                return None
+            raise
+        else:
+            return result.data
+
+    async def delete_index(self, index_id: str) -> None:
+        """Delete an index."""
+        await self._request(
+            "DELETE", APIEndpoints.API_V1_INDEXES_INDEX_ID.format(index_id=index_id)
+        )

kodit/infrastructure/api/client/search_client.py

@@ -0,0 +1,86 @@
+"""Search operations API client for Kodit server."""
+
+from datetime import datetime
+
+from kodit.infrastructure.api.v1.schemas.search import (
+    SearchAttributes,
+    SearchData,
+    SearchFilters,
+    SearchRequest,
+    SearchResponse,
+    SnippetData,
+)
+
+from .base import BaseAPIClient
+from .generated_endpoints import APIEndpoints
+
+
+class SearchClient(BaseAPIClient):
+    """API client for search operations."""
+
+    async def search(  # noqa: PLR0913
+        self,
+        keywords: list[str] | None = None,
+        code_query: str | None = None,
+        text_query: str | None = None,
+        limit: int = 10,
+        languages: list[str] | None = None,
+        authors: list[str] | None = None,
+        start_date: datetime | None = None,
+        end_date: datetime | None = None,
+        sources: list[str] | None = None,
+        file_patterns: list[str] | None = None,
+    ) -> list[SnippetData]:
+        """Search for code snippets.
+
+        Args:
+            keywords: Keywords to search for
+            code_query: Code search query
+            text_query: Text search query
+            limit: Maximum number of results
+            languages: Programming languages to filter by
+            authors: Authors to filter by
+            start_date: Filter snippets created after this date
+            end_date: Filter snippets created before this date
+            sources: Source repositories to filter by
+            file_patterns: File path patterns to filter by
+
+        Returns:
+            List of matching snippets
+
+        Raises:
+            KoditAPIError: If the request fails
+
+        """
+        filters = None
+        if any([languages, authors, start_date, end_date, sources, file_patterns]):
+            filters = SearchFilters(
+                languages=languages,
+                authors=authors,
+                start_date=start_date,
+                end_date=end_date,
+                sources=sources,
+                file_patterns=file_patterns,
+            )
+
+        request = SearchRequest(
+            data=SearchData(
+                type="search",
+                attributes=SearchAttributes(
+                    keywords=keywords,
+                    code=code_query,
+                    text=text_query,
+                    limit=limit,
+                    filters=filters,
+                ),
+            )
+        )
+
+        response = await self._request(
+            "POST",
+            APIEndpoints.API_V1_SEARCH,
+            json=request.model_dump(exclude_none=True),
+        )
+
+        result = SearchResponse.model_validate_json(response.text)
+        return result.data

kodit/infrastructure/api/v1/dependencies.py

@@ -12,6 +12,7 @@ from kodit.application.factories.code_indexing_factory import (
 from kodit.application.services.code_indexing_application_service import (
     CodeIndexingApplicationService,
 )
+from kodit.application.services.queue_service import QueueService
 from kodit.config import AppContext
 from kodit.domain.services.index_query_service import IndexQueryService
 from kodit.infrastructure.indexing.fusion_service import ReciprocalRankFusionService
@@ -68,3 +69,15 @@ async def get_indexing_app_service(
 IndexingAppServiceDep = Annotated[
     CodeIndexingApplicationService, Depends(get_indexing_app_service)
 ]
+
+
+async def get_queue_service(
+    session: DBSessionDep,
+) -> QueueService:
+    """Get queue service dependency."""
+    return QueueService(
+        session=session,
+    )
+
+
+QueueServiceDep = Annotated[QueueService, Depends(get_queue_service)]

kodit/infrastructure/api/v1/routers/indexes.py

@@ -1,11 +1,14 @@
 """Index management router for the REST API."""
 
-from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException
+from fastapi import APIRouter, Depends, HTTPException
 
+from kodit.domain.entities import Task
+from kodit.domain.value_objects import QueuePriority
 from kodit.infrastructure.api.middleware.auth import api_key_auth
 from kodit.infrastructure.api.v1.dependencies import (
     IndexingAppServiceDep,
     IndexQueryServiceDep,
+    QueueServiceDep,
 )
 from kodit.infrastructure.api.v1.schemas.index import (
     IndexAttributes,
@@ -52,15 +55,17 @@ async def list_indexes(
 @router.post("", status_code=202)
 async def create_index(
     request: IndexCreateRequest,
-    background_tasks: BackgroundTasks,
     app_service: IndexingAppServiceDep,
+    queue_service: QueueServiceDep,
 ) -> IndexResponse:
     """Create a new index and start async indexing."""
     # Create index using the application service
     index = await app_service.create_index_from_uri(request.data.attributes.uri)
 
-    # Start async indexing in background
-    background_tasks.add_task(app_service.run_index, index)
+    # Add the indexing task to the queue
+    await queue_service.enqueue_task(
+        Task.create_index_update_task(index.id, QueuePriority.USER_INITIATED)
+    )
 
     return IndexResponse(
         data=IndexData(

kodit/infrastructure/enrichment/local_enrichment_provider.py

@@ -8,6 +8,7 @@ import tiktoken
 
 from kodit.domain.services.enrichment_service import EnrichmentProvider
 from kodit.domain.value_objects import EnrichmentRequest, EnrichmentResponse
+from kodit.infrastructure.enrichment.utils import clean_thinking_tags
 
 ENRICHMENT_SYSTEM_PROMPT = """
 You are a professional software developer. You will be given a snippet of code.
@@ -109,7 +110,9 @@ class LocalEnrichmentProvider(EnrichmentProvider):
             content = self.tokenizer.decode(output_ids, skip_special_tokens=True).strip(  # type: ignore[attr-defined]
                 "\n"
             )
+            # Remove thinking tags from the response
+            cleaned_content = clean_thinking_tags(content)
             yield EnrichmentResponse(
                 snippet_id=prompt["id"],
-                text=content,
+                text=cleaned_content,
             )

kodit/infrastructure/enrichment/openai_enrichment_provider.py

@@ -9,6 +9,7 @@ import structlog
 
 from kodit.domain.services.enrichment_service import EnrichmentProvider
 from kodit.domain.value_objects import EnrichmentRequest, EnrichmentResponse
+from kodit.infrastructure.enrichment.utils import clean_thinking_tags
 
 ENRICHMENT_SYSTEM_PROMPT = """
 You are a professional software developer. You will be given a snippet of code.
@@ -19,6 +20,7 @@ Please provide a concise explanation of the code.
 OPENAI_NUM_PARALLEL_TASKS = 40
 
 
+
 class OpenAIEnrichmentProvider(EnrichmentProvider):
     """OpenAI enrichment provider implementation using httpx."""
 
@@ -135,9 +137,11 @@ class OpenAIEnrichmentProvider(EnrichmentProvider):
                         .get("message", {})
                         .get("content", "")
                     )
+                    # Remove thinking tags from the response
+                    cleaned_content = clean_thinking_tags(content or "")
                     return EnrichmentResponse(
                         snippet_id=request.snippet_id,
-                        text=content or "",
+                        text=cleaned_content,
                     )
                 except Exception as e:
                     self.log.exception("Error enriching request", error=str(e))