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

kodit/infrastructure/embedding/embedding_providers/openai_embedding_provider.py

@@ -1,11 +1,12 @@
-"""OpenAI embedding provider implementation."""
+"""OpenAI embedding provider implementation using httpx."""
 
 import asyncio
 from collections.abc import AsyncGenerator
+from typing import Any
 
+import httpx
 import structlog
 import tiktoken
-from openai import AsyncOpenAI
 from tiktoken import Encoding
 
 from kodit.domain.services.embedding_service import EmbeddingProvider
@@ -22,29 +23,53 @@ OPENAI_NUM_PARALLEL_TASKS = 10  # Semaphore limit for concurrent OpenAI requests
 
 
 class OpenAIEmbeddingProvider(EmbeddingProvider):
-    """OpenAI embedding provider that uses OpenAI's embedding API."""
+    """OpenAI embedding provider that uses OpenAI's embedding API via httpx."""
 
-    def __init__(
+    def __init__(  # noqa: PLR0913
         self,
-        openai_client: AsyncOpenAI,
+        api_key: str | None = None,
+        base_url: str = "https://api.openai.com",
         model_name: str = "text-embedding-3-small",
         num_parallel_tasks: int = OPENAI_NUM_PARALLEL_TASKS,
+        socket_path: str | None = None,
+        timeout: float = 30.0,
     ) -> None:
         """Initialize the OpenAI embedding provider.
 
         Args:
-            openai_client: The OpenAI client instance
-            model_name: The model name to use for embeddings
+            api_key: The OpenAI API key.
+            base_url: The base URL for the OpenAI API.
+            model_name: The model name to use for embeddings.
+            num_parallel_tasks: Maximum number of concurrent requests.
+            socket_path: Optional Unix socket path for local communication.
+            timeout: Request timeout in seconds.
 
         """
-        self.openai_client = openai_client
         self.model_name = model_name
         self.num_parallel_tasks = num_parallel_tasks
         self.log = structlog.get_logger(__name__)
+        self.api_key = api_key
+        self.base_url = base_url
+        self.socket_path = socket_path
+        self.timeout = timeout
 
         # Lazily initialised token encoding
         self._encoding: Encoding | None = None
 
+        # Create httpx client with optional Unix socket support
+        if socket_path:
+            transport = httpx.AsyncHTTPTransport(uds=socket_path)
+            self.http_client = httpx.AsyncClient(
+                transport=transport,
+                base_url="http://localhost",  # Base URL for Unix socket
+                timeout=timeout,
+            )
+        else:
+            self.http_client = httpx.AsyncClient(
+                base_url=base_url,
+                timeout=timeout,
+            )
+
     # ---------------------------------------------------------------------
     # Helper utilities
     # ---------------------------------------------------------------------
@@ -76,6 +101,37 @@ class OpenAIEmbeddingProvider(EmbeddingProvider):
             batch_size=BATCH_SIZE,
         )
 
+    async def _call_embeddings_api(
+        self, texts: list[str]
+    ) -> dict[str, Any]:
+        """Call the embeddings API using httpx.
+
+        Args:
+            texts: The texts to embed.
+
+        Returns:
+            The API response as a dictionary.
+
+        """
+        headers = {
+            "Content-Type": "application/json",
+        }
+        if self.api_key:
+            headers["Authorization"] = f"Bearer {self.api_key}"
+
+        data = {
+            "model": self.model_name,
+            "input": texts,
+        }
+
+        response = await self.http_client.post(
+            "/v1/embeddings",
+            json=data,
+            headers=headers,
+        )
+        response.raise_for_status()
+        return response.json()
+
     async def embed(
         self, data: list[EmbeddingRequest]
     ) -> AsyncGenerator[list[EmbeddingResponse], None]:
@@ -99,17 +155,17 @@ class OpenAIEmbeddingProvider(EmbeddingProvider):
         ) -> list[EmbeddingResponse]:
             async with sem:
                 try:
-                    response = await self.openai_client.embeddings.create(
-                        model=self.model_name,
-                        input=[item.text for item in batch],
+                    response = await self._call_embeddings_api(
+                        [item.text for item in batch]
                     )
+                    embeddings_data = response.get("data", [])
 
                     return [
                         EmbeddingResponse(
                             snippet_id=item.snippet_id,
-                            embedding=embedding.embedding,
+                            embedding=emb_data.get("embedding", []),
                         )
-                        for item, embedding in zip(batch, response.data, strict=True)
+                        for item, emb_data in zip(batch, embeddings_data, strict=True)
                     ]
                 except Exception as e:
                     self.log.exception("Error embedding batch", error=str(e))
@@ -119,3 +175,9 @@ class OpenAIEmbeddingProvider(EmbeddingProvider):
         tasks = [_process_batch(batch) for batch in batched_data]
         for task in asyncio.as_completed(tasks):
             yield await task
+
+    async def close(self) -> None:
+        """Close the HTTP client."""
+        if hasattr(self, "http_client"):
+            await self.http_client.aclose()
+

kodit/infrastructure/enrichment/enrichment_factory.py

@@ -45,17 +45,14 @@ def enrichment_domain_service_factory(
     enrichment_provider: EnrichmentProvider | None = None
     if endpoint and endpoint.type == "openai":
         log_event("kodit.enrichment", {"provider": "openai"})
-        from openai import AsyncOpenAI
-
+        # Use new httpx-based provider with socket support
         enrichment_provider = OpenAIEnrichmentProvider(
-            openai_client=AsyncOpenAI(
-                api_key=endpoint.api_key or "default",
-                base_url=endpoint.base_url or "https://api.openai.com/v1",
-                timeout=60,
-                max_retries=2,
-            ),
+            api_key=endpoint.api_key,
+            base_url=endpoint.base_url or "https://api.openai.com/v1",
             model_name=endpoint.model or "gpt-4o-mini",
             num_parallel_tasks=endpoint.num_parallel_tasks or OPENAI_NUM_PARALLEL_TASKS,
+            socket_path=endpoint.socket_path,
+            timeout=endpoint.timeout or 30.0,
         )
     else:
         log_event("kodit.enrichment", {"provider": "local"})

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

@@ -1,13 +1,15 @@
-"""OpenAI enrichment provider implementation."""
+"""OpenAI enrichment provider implementation using httpx."""
 
 import asyncio
 from collections.abc import AsyncGenerator
 from typing import Any
 
+import httpx
 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.
@@ -18,26 +20,82 @@ Please provide a concise explanation of the code.
 OPENAI_NUM_PARALLEL_TASKS = 40
 
 
+
 class OpenAIEnrichmentProvider(EnrichmentProvider):
-    """OpenAI enrichment provider implementation."""
+    """OpenAI enrichment provider implementation using httpx."""
 
-    def __init__(
+    def __init__(  # noqa: PLR0913
         self,
-        openai_client: Any,
+        api_key: str | None = None,
+        base_url: str = "https://api.openai.com",
         model_name: str = "gpt-4o-mini",
         num_parallel_tasks: int = OPENAI_NUM_PARALLEL_TASKS,
+        socket_path: str | None = None,
+        timeout: float = 30.0,
     ) -> None:
         """Initialize the OpenAI enrichment provider.
 
         Args:
-            openai_client: The OpenAI client instance.
+            api_key: The OpenAI API key.
+            base_url: The base URL for the OpenAI API.
             model_name: The model name to use for enrichment.
+            num_parallel_tasks: Maximum number of concurrent requests.
+            socket_path: Optional Unix socket path for local communication.
+            timeout: Request timeout in seconds.
 
         """
         self.log = structlog.get_logger(__name__)
-        self.openai_client = openai_client
         self.model_name = model_name
         self.num_parallel_tasks = num_parallel_tasks
+        self.api_key = api_key
+        self.base_url = base_url
+        self.socket_path = socket_path
+        self.timeout = timeout
+
+        # Create httpx client with optional Unix socket support
+        if socket_path:
+            transport = httpx.AsyncHTTPTransport(uds=socket_path)
+            self.http_client = httpx.AsyncClient(
+                transport=transport,
+                base_url="http://localhost",  # Base URL for Unix socket
+                timeout=timeout,
+            )
+        else:
+            self.http_client = httpx.AsyncClient(
+                base_url=base_url,
+                timeout=timeout,
+            )
+
+    async def _call_chat_completion(
+        self, messages: list[dict[str, str]]
+    ) -> dict[str, Any]:
+        """Call the chat completion API using httpx.
+
+        Args:
+            messages: The messages to send to the API.
+
+        Returns:
+            The API response as a dictionary.
+
+        """
+        headers = {
+            "Content-Type": "application/json",
+        }
+        if self.api_key:
+            headers["Authorization"] = f"Bearer {self.api_key}"
+
+        data = {
+            "model": self.model_name,
+            "messages": messages,
+        }
+
+        response = await self.http_client.post(
+            "/v1/chat/completions",
+            json=data,
+            headers=headers,
+        )
+        response.raise_for_status()
+        return response.json()
 
     async def enrich(
         self, requests: list[EnrichmentRequest]
@@ -66,19 +124,24 @@ class OpenAIEnrichmentProvider(EnrichmentProvider):
                         text="",
                     )
                 try:
-                    response = await self.openai_client.chat.completions.create(
-                        model=self.model_name,
-                        messages=[
-                            {
-                                "role": "system",
-                                "content": ENRICHMENT_SYSTEM_PROMPT,
-                            },
-                            {"role": "user", "content": request.text},
-                        ],
+                    messages = [
+                        {
+                            "role": "system",
+                            "content": ENRICHMENT_SYSTEM_PROMPT,
+                        },
+                        {"role": "user", "content": request.text},
+                    ]
+                    response = await self._call_chat_completion(messages)
+                    content = (
+                        response.get("choices", [{}])[0]
+                        .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=response.choices[0].message.content or "",
+                        text=cleaned_content,
                     )
                 except Exception as e:
                     self.log.exception("Error enriching request", error=str(e))
@@ -93,3 +156,8 @@ class OpenAIEnrichmentProvider(EnrichmentProvider):
         # Process all requests and yield results as they complete
         for task in asyncio.as_completed(tasks):
             yield await task
+
+    async def close(self) -> None:
+        """Close the HTTP client."""
+        if hasattr(self, "http_client"):
+            await self.http_client.aclose()

kodit/infrastructure/enrichment/utils.py

@@ -0,0 +1,30 @@
+"""Utility functions for enrichment processing."""
+
+import re
+
+
+def clean_thinking_tags(content: str) -> str:
+    """Remove <think>...</think> tags from content.
+
+    This utility handles thinking tags that may be produced by various AI models,
+    including both local and remote models. It safely removes thinking content
+    while preserving the actual response.
+
+    Args:
+        content: The content that may contain thinking tags.
+
+    Returns:
+        The content with thinking tags removed and cleaned up.
+
+    """
+    if not content:
+        return content
+
+    # Remove thinking tags using regex with DOTALL flag to match across newlines
+    cleaned = re.sub(
+        r"<think>.*?</think>", "", content, flags=re.DOTALL | re.IGNORECASE
+    )
+
+    # Clean up extra whitespace that may be left behind
+    cleaned = re.sub(r"\n\s*\n\s*\n", "\n\n", cleaned)
+    return cleaned.strip()  # Remove leading/trailing whitespace

kodit/infrastructure/mappers/task_mapper.py

@@ -0,0 +1,81 @@
+"""Task mapper for the task queue."""
+
+from typing import ClassVar
+
+from kodit.domain.entities import Task
+from kodit.domain.value_objects import TaskType
+from kodit.infrastructure.sqlalchemy import entities as db_entities
+
+
+class TaskTypeMapper:
+    """Maps between domain QueuedTaskType and SQLAlchemy TaskType."""
+
+    # Map TaskType enum to QueuedTaskType
+    TASK_TYPE_MAPPING: ClassVar[dict[db_entities.TaskType, TaskType]] = {
+        db_entities.TaskType.INDEX_UPDATE: TaskType.INDEX_UPDATE,
+    }
+
+    @staticmethod
+    def to_domain_type(task_type: db_entities.TaskType) -> TaskType:
+        """Convert SQLAlchemy TaskType to domain QueuedTaskType."""
+        if task_type not in TaskTypeMapper.TASK_TYPE_MAPPING:
+            raise ValueError(f"Unknown task type: {task_type}")
+        return TaskTypeMapper.TASK_TYPE_MAPPING[task_type]
+
+    @staticmethod
+    def from_domain_type(task_type: TaskType) -> db_entities.TaskType:
+        """Convert domain QueuedTaskType to SQLAlchemy TaskType."""
+        if task_type not in TaskTypeMapper.TASK_TYPE_MAPPING.values():
+            raise ValueError(f"Unknown task type: {task_type}")
+
+        # Find value in TASK_TYPE_MAPPING
+        return next(
+            (
+                db_task_type
+                for db_task_type, domain_task_type in TaskTypeMapper.TASK_TYPE_MAPPING.items()  # noqa: E501
+                if domain_task_type == task_type
+            )
+        )
+
+
+class TaskMapper:
+    """Maps between domain QueuedTask and SQLAlchemy Task entities.
+
+    This mapper handles the conversion between the existing domain and
+    persistence layers without creating any new entities.
+    """
+
+    @staticmethod
+    def to_domain_task(record: db_entities.Task) -> Task:
+        """Convert SQLAlchemy Task record to domain QueuedTask.
+
+        Since QueuedTask doesn't have status fields, we store processing
+        state in the payload.
+        """
+        # Get the task type
+        task_type = TaskTypeMapper.to_domain_type(record.type)
+
+        # The dedup_key becomes the id in the domain entity
+        return Task(
+            id=record.dedup_key,  # Use dedup_key as the unique identifier
+            type=task_type,
+            priority=record.priority,
+            payload=record.payload or {},
+            created_at=record.created_at,
+            updated_at=record.updated_at,
+        )
+
+    @staticmethod
+    def from_domain_task(task: Task) -> db_entities.Task:
+        """Convert domain QueuedTask to SQLAlchemy Task record."""
+        if task.type not in TaskTypeMapper.TASK_TYPE_MAPPING.values():
+            raise ValueError(f"Unknown task type: {task.type}")
+
+        # Find value in TASK_TYPE_MAPPING
+        task_type = TaskTypeMapper.from_domain_type(task.type)
+        return db_entities.Task(
+            dedup_key=task.id,
+            type=task_type,
+            payload=task.payload,
+            priority=task.priority,
+        )

kodit/infrastructure/sqlalchemy/entities.py

@@ -201,3 +201,38 @@ class Snippet(Base, CommonMixin):
         self.index_id = index_id
         self.content = content
         self.summary = summary
+
+
+class TaskType(Enum):
+    """Task type."""
+
+    INDEX_UPDATE = 1
+
+
+class Task(Base, CommonMixin):
+    """Queued tasks."""
+
+    __tablename__ = "tasks"
+
+    # dedup_key is used to deduplicate items in the queue
+    dedup_key: Mapped[str] = mapped_column(String(255), index=True)
+    # type represents what the task is meant to achieve
+    type: Mapped[TaskType] = mapped_column(SQLAlchemyEnum(TaskType), index=True)
+    # payload contains the task-specific payload data
+    payload: Mapped[dict] = mapped_column(JSON)
+    # priority is used to determine the order of the items in the queue
+    priority: Mapped[int] = mapped_column(Integer)
+
+    def __init__(
+        self,
+        dedup_key: str,
+        type: TaskType,  # noqa: A002
+        payload: dict,
+        priority: int,
+    ) -> None:
+        """Initialize the queue item."""
+        super().__init__()
+        self.dedup_key = dedup_key
+        self.type = type
+        self.payload = payload
+        self.priority = priority

kodit/infrastructure/sqlalchemy/index_repository.py

@@ -597,12 +597,12 @@ class SqlAlchemyIndexRepository(IndexRepository):
         )
         await self._session.execute(stmt)
 
+        # Delete the index
+        stmt = delete(db_entities.Index).where(db_entities.Index.id == index.id)
+        await self._session.execute(stmt)
+
         # Delete the source
         stmt = delete(db_entities.Source).where(
             db_entities.Source.id == index.source.id
         )
         await self._session.execute(stmt)
-
-        # Delete the index
-        stmt = delete(db_entities.Index).where(db_entities.Index.id == index.id)
-        await self._session.execute(stmt)

kodit/infrastructure/sqlalchemy/task_repository.py

@@ -0,0 +1,81 @@
+"""Task repository for the task queue."""
+
+import structlog
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from kodit.domain.entities import Task
+from kodit.domain.protocols import TaskRepository
+from kodit.domain.value_objects import TaskType
+from kodit.infrastructure.mappers.task_mapper import TaskMapper, TaskTypeMapper
+from kodit.infrastructure.sqlalchemy import entities as db_entities
+
+
+class SqlAlchemyTaskRepository(TaskRepository):
+    """Repository for task persistence using the existing Task entity."""
+
+    def __init__(self, session: AsyncSession) -> None:
+        """Initialize the repository."""
+        self.session = session
+        self.log = structlog.get_logger(__name__)
+
+    async def add(
+        self,
+        task: Task,
+    ) -> None:
+        """Create a new task in the database."""
+        self.session.add(TaskMapper.from_domain_task(task))
+
+    async def get(self, task_id: str) -> Task | None:
+        """Get a task by ID."""
+        stmt = select(db_entities.Task).where(db_entities.Task.dedup_key == task_id)
+        result = await self.session.execute(stmt)
+        db_task = result.scalar_one_or_none()
+        if not db_task:
+            return None
+        return TaskMapper.to_domain_task(db_task)
+
+    async def take(self) -> Task | None:
+        """Take a task for processing and remove it from the database."""
+        stmt = (
+            select(db_entities.Task)
+            .order_by(db_entities.Task.priority.desc(), db_entities.Task.created_at)
+            .limit(1)
+        )
+        result = await self.session.execute(stmt)
+        db_task = result.scalar_one_or_none()
+        if not db_task:
+            return None
+        await self.session.delete(db_task)
+        return TaskMapper.to_domain_task(db_task)
+
+    async def update(self, task: Task) -> None:
+        """Update a task in the database."""
+        stmt = select(db_entities.Task).where(db_entities.Task.dedup_key == task.id)
+        result = await self.session.execute(stmt)
+        db_task = result.scalar_one_or_none()
+
+        if not db_task:
+            raise ValueError(f"Task not found: {task.id}")
+
+        db_task.priority = task.priority
+        db_task.payload = task.payload
+
+    async def list(self, task_type: TaskType | None = None) -> list[Task]:
+        """List tasks with optional status filter."""
+        stmt = select(db_entities.Task)
+
+        if task_type:
+            stmt = stmt.where(
+                db_entities.Task.type == TaskTypeMapper.from_domain_type(task_type)
+            )
+
+        stmt = stmt.order_by(
+            db_entities.Task.priority.desc(), db_entities.Task.created_at
+        )
+
+        result = await self.session.execute(stmt)
+        records = result.scalars().all()
+
+        # Convert to domain entities
+        return [TaskMapper.to_domain_task(record) for record in records]

kodit/middleware.py

@@ -53,6 +53,7 @@ async def logging_middleware(request: Request, call_next: Callable) -> Response:
                 "client_host": client_host,
                 "client_port": client_port,
             },
+            headers=dict(request.headers),
             network={"client": {"ip": client_host, "port": client_port}},
             duration=process_time,
         )

kodit/migrations/versions/9cf0e87de578_add_queue.py

@@ -0,0 +1,47 @@
+# ruff: noqa
+"""add queue
+
+Revision ID: 9cf0e87de578
+Revises: 4073b33f9436
+Create Date: 2025-08-06 17:38:21.055235
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = '9cf0e87de578'
+down_revision: Union[str, None] = '4073b33f9436'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('tasks',
+    sa.Column('dedup_key', sa.String(length=255), nullable=False),
+    sa.Column('type', sa.Enum('INDEX_UPDATE', name='tasktype'), nullable=False),
+    sa.Column('payload', sa.JSON(), nullable=False),
+    sa.Column('priority', sa.Integer(), nullable=False),
+    sa.Column('id', sa.Integer(), autoincrement=True, nullable=False),
+    sa.Column('created_at', sa.DateTime(timezone=True), nullable=False),
+    sa.Column('updated_at', sa.DateTime(timezone=True), nullable=False),
+    sa.PrimaryKeyConstraint('id')
+    )
+    op.create_index(op.f('ix_tasks_dedup_key'), 'tasks', ['dedup_key'], unique=False)
+    op.create_index(op.f('ix_tasks_type'), 'tasks', ['type'], unique=False)
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index(op.f('ix_tasks_type'), table_name='tasks')
+    op.drop_index(op.f('ix_tasks_dedup_key'), table_name='tasks')
+    op.drop_table('tasks')
+    # ### end Alembic commands ###