xai-review 0.20.0__py3-none-any.whl → 0.22.0__py3-none-any.whl

ai_review/clients/claude/client.py

@@ -26,7 +26,7 @@ class ClaudeHTTPClient(HTTPClient):
 def get_claude_http_client() -> ClaudeHTTPClient:
     logger = get_logger("CLAUDE_HTTP_CLIENT")
     logger_event_hook = LoggerEventHook(logger=logger)
-    retry_transport = RetryTransport(transport=AsyncHTTPTransport())
+    retry_transport = RetryTransport(logger=logger, transport=AsyncHTTPTransport())
 
     client = AsyncClient(
         timeout=settings.llm.http_client.timeout,

ai_review/clients/gemini/client.py

@@ -29,7 +29,7 @@ class GeminiHTTPClient(HTTPClient):
 def get_gemini_http_client() -> GeminiHTTPClient:
     logger = get_logger("GEMINI_HTTP_CLIENT")
     logger_event_hook = LoggerEventHook(logger=logger)
-    retry_transport = RetryTransport(transport=AsyncHTTPTransport())
+    retry_transport = RetryTransport(logger=logger, transport=AsyncHTTPTransport())
 
     client = AsyncClient(
         timeout=settings.llm.http_client.timeout,

ai_review/clients/github/client.py

@@ -15,7 +15,7 @@ class GitHubHTTPClient:
 def get_github_http_client() -> GitHubHTTPClient:
     logger = get_logger("GITHUB_HTTP_CLIENT")
     logger_event_hook = LoggerEventHook(logger=logger)
-    retry_transport = RetryTransport(transport=AsyncHTTPTransport())
+    retry_transport = RetryTransport(logger=logger, transport=AsyncHTTPTransport())
 
     client = AsyncClient(
         timeout=settings.llm.http_client.timeout,

ai_review/clients/github/pr/client.py

@@ -1,15 +1,23 @@
-from httpx import Response
+from httpx import Response, QueryParams
 
 from ai_review.clients.github.pr.schema.comments import (
+    GitHubGetPRCommentsQuerySchema,
     GitHubGetPRCommentsResponseSchema,
     GitHubCreateIssueCommentRequestSchema,
     GitHubCreateIssueCommentResponseSchema,
     GitHubCreateReviewCommentRequestSchema,
     GitHubCreateReviewCommentResponseSchema
 )
-from ai_review.clients.github.pr.schema.files import GitHubGetPRFilesResponseSchema
+from ai_review.clients.github.pr.schema.files import (
+    GitHubGetPRFilesQuerySchema,
+    GitHubGetPRFilesResponseSchema
+)
 from ai_review.clients.github.pr.schema.pull_request import GitHubGetPRResponseSchema
-from ai_review.clients.github.pr.schema.reviews import GitHubGetPRReviewsResponseSchema
+from ai_review.clients.github.pr.schema.reviews import (
+    GitHubGetPRReviewsQuerySchema,
+    GitHubGetPRReviewsResponseSchema
+)
+from ai_review.clients.github.pr.types import GitHubPullRequestsHTTPClientProtocol
 from ai_review.libs.http.client import HTTPClient
 from ai_review.libs.http.handlers import HTTPClientError, handle_http_error
 
@@ -18,22 +26,49 @@ class GitHubPullRequestsHTTPClientError(HTTPClientError):
     pass
 
 
-class GitHubPullRequestsHTTPClient(HTTPClient):
+class GitHubPullRequestsHTTPClient(HTTPClient, GitHubPullRequestsHTTPClientProtocol):
     @handle_http_error(client="GitHubPullRequestsHTTPClient", exception=GitHubPullRequestsHTTPClientError)
     async def get_pull_request_api(self, owner: str, repo: str, pull_number: str) -> Response:
         return await self.get(f"/repos/{owner}/{repo}/pulls/{pull_number}")
 
     @handle_http_error(client="GitHubPullRequestsHTTPClient", exception=GitHubPullRequestsHTTPClientError)
-    async def get_files_api(self, owner: str, repo: str, pull_number: str) -> Response:
-        return await self.get(f"/repos/{owner}/{repo}/pulls/{pull_number}/files")
+    async def get_files_api(
+            self,
+            owner: str,
+            repo: str,
+            pull_number: str,
+            query: GitHubGetPRFilesQuerySchema
+    ) -> Response:
+        return await self.get(
+            f"/repos/{owner}/{repo}/pulls/{pull_number}/files",
+            query=QueryParams(**query.model_dump())
+        )
 
     @handle_http_error(client="GitHubPullRequestsHTTPClient", exception=GitHubPullRequestsHTTPClientError)
-    async def get_issue_comments_api(self, owner: str, repo: str, issue_number: str) -> Response:
-        return await self.get(f"/repos/{owner}/{repo}/issues/{issue_number}/comments")
+    async def get_issue_comments_api(
+            self,
+            owner: str,
+            repo: str,
+            issue_number: str,
+            query: GitHubGetPRCommentsQuerySchema,
+    ) -> Response:
+        return await self.get(
+            f"/repos/{owner}/{repo}/issues/{issue_number}/comments",
+            query=QueryParams(**query.model_dump())
+        )
 
     @handle_http_error(client="GitHubPullRequestsHTTPClient", exception=GitHubPullRequestsHTTPClientError)
-    async def get_review_comments_api(self, owner: str, repo: str, pull_number: str) -> Response:
-        return await self.get(f"/repos/{owner}/{repo}/pulls/{pull_number}/comments")
+    async def get_review_comments_api(
+            self,
+            owner: str,
+            repo: str,
+            pull_number: str,
+            query: GitHubGetPRCommentsQuerySchema,
+    ) -> Response:
+        return await self.get(
+            f"/repos/{owner}/{repo}/pulls/{pull_number}/comments",
+            query=QueryParams(**query.model_dump())
+        )
 
     @handle_http_error(client="GitHubPullRequestsHTTPClient", exception=GitHubPullRequestsHTTPClientError)
     async def create_review_comment_api(
@@ -62,27 +97,40 @@ class GitHubPullRequestsHTTPClient(HTTPClient):
         )
 
     @handle_http_error(client="GitHubPullRequestsHTTPClient", exception=GitHubPullRequestsHTTPClientError)
-    async def get_reviews_api(self, owner: str, repo: str, pull_number: str) -> Response:
-        return await self.get(f"/repos/{owner}/{repo}/pulls/{pull_number}/reviews")
+    async def get_reviews_api(
+            self,
+            owner: str,
+            repo: str,
+            pull_number: str,
+            query: GitHubGetPRReviewsQuerySchema
+    ) -> Response:
+        return await self.get(
+            f"/repos/{owner}/{repo}/pulls/{pull_number}/reviews",
+            query=QueryParams(**query.model_dump())
+        )
 
     async def get_pull_request(self, owner: str, repo: str, pull_number: str) -> GitHubGetPRResponseSchema:
         response = await self.get_pull_request_api(owner, repo, pull_number)
         return GitHubGetPRResponseSchema.model_validate_json(response.text)
 
     async def get_files(self, owner: str, repo: str, pull_number: str) -> GitHubGetPRFilesResponseSchema:
-        response = await self.get_files_api(owner, repo, pull_number)
+        query = GitHubGetPRFilesQuerySchema(per_page=100)
+        response = await self.get_files_api(owner, repo, pull_number, query)
         return GitHubGetPRFilesResponseSchema.model_validate_json(response.text)
 
     async def get_issue_comments(self, owner: str, repo: str, issue_number: str) -> GitHubGetPRCommentsResponseSchema:
-        response = await self.get_issue_comments_api(owner, repo, issue_number)
+        query = GitHubGetPRCommentsQuerySchema(per_page=100)
+        response = await self.get_issue_comments_api(owner, repo, issue_number, query)
         return GitHubGetPRCommentsResponseSchema.model_validate_json(response.text)
 
     async def get_review_comments(self, owner: str, repo: str, pull_number: str) -> GitHubGetPRCommentsResponseSchema:
-        response = await self.get_review_comments_api(owner, repo, pull_number)
+        query = GitHubGetPRCommentsQuerySchema(per_page=100)
+        response = await self.get_review_comments_api(owner, repo, pull_number, query)
         return GitHubGetPRCommentsResponseSchema.model_validate_json(response.text)
 
     async def get_reviews(self, owner: str, repo: str, pull_number: str) -> GitHubGetPRReviewsResponseSchema:
-        response = await self.get_reviews_api(owner, repo, pull_number)
+        query = GitHubGetPRReviewsQuerySchema(per_page=100)
+        response = await self.get_reviews_api(owner, repo, pull_number, query)
         return GitHubGetPRReviewsResponseSchema.model_validate_json(response.text)
 
     async def create_review_comment(

ai_review/clients/github/pr/schema/comments.py

@@ -8,6 +8,10 @@ class GitHubPRCommentSchema(BaseModel):
     line: int | None = None
 
 
+class GitHubGetPRCommentsQuerySchema(BaseModel):
+    per_page: int
+
+
 class GitHubGetPRCommentsResponseSchema(RootModel[list[GitHubPRCommentSchema]]):
     root: list[GitHubPRCommentSchema]
 

ai_review/clients/github/pr/schema/files.py

@@ -8,5 +8,9 @@ class GitHubPRFileSchema(BaseModel):
     filename: str
 
 
+class GitHubGetPRFilesQuerySchema(BaseModel):
+    per_page: int
+
+
 class GitHubGetPRFilesResponseSchema(RootModel[list[GitHubPRFileSchema]]):
     root: list[GitHubPRFileSchema]

ai_review/clients/github/pr/schema/reviews.py

@@ -9,5 +9,9 @@ class GitHubPRReviewSchema(BaseModel):
     state: str
 
 
+class GitHubGetPRReviewsQuerySchema(BaseModel):
+    per_page: int
+
+
 class GitHubGetPRReviewsResponseSchema(RootModel[list[GitHubPRReviewSchema]]):
     root: list[GitHubPRReviewSchema]

ai_review/clients/github/pr/types.py

@@ -0,0 +1,49 @@
+from typing import Protocol
+
+from ai_review.clients.github.pr.schema.comments import (
+    GitHubGetPRCommentsResponseSchema,
+    GitHubCreateIssueCommentResponseSchema,
+    GitHubCreateReviewCommentResponseSchema,
+    GitHubCreateReviewCommentRequestSchema,
+)
+from ai_review.clients.github.pr.schema.files import GitHubGetPRFilesResponseSchema
+from ai_review.clients.github.pr.schema.pull_request import GitHubGetPRResponseSchema
+from ai_review.clients.github.pr.schema.reviews import GitHubGetPRReviewsResponseSchema
+
+
+class GitHubPullRequestsHTTPClientProtocol(Protocol):
+    async def get_pull_request(self, owner: str, repo: str, pull_number: str) -> GitHubGetPRResponseSchema: ...
+
+    async def get_files(self, owner: str, repo: str, pull_number: str) -> GitHubGetPRFilesResponseSchema: ...
+
+    async def get_issue_comments(
+            self,
+            owner: str,
+            repo: str,
+            issue_number: str
+    ) -> GitHubGetPRCommentsResponseSchema: ...
+
+    async def get_review_comments(
+            self,
+            owner: str,
+            repo: str,
+            pull_number: str
+    ) -> GitHubGetPRCommentsResponseSchema: ...
+
+    async def get_reviews(self, owner: str, repo: str, pull_number: str) -> GitHubGetPRReviewsResponseSchema: ...
+
+    async def create_review_comment(
+            self,
+            owner: str,
+            repo: str,
+            pull_number: str,
+            request: GitHubCreateReviewCommentRequestSchema,
+    ) -> GitHubCreateReviewCommentResponseSchema: ...
+
+    async def create_issue_comment(
+            self,
+            owner: str,
+            repo: str,
+            issue_number: str,
+            body: str,
+    ) -> GitHubCreateIssueCommentResponseSchema: ...

ai_review/clients/gitlab/client.py

@@ -15,7 +15,7 @@ class GitLabHTTPClient:
 def get_gitlab_http_client() -> GitLabHTTPClient:
     logger = get_logger("GITLAB_HTTP_CLIENT")
     logger_event_hook = LoggerEventHook(logger=logger)
-    retry_transport = RetryTransport(transport=AsyncHTTPTransport())
+    retry_transport = RetryTransport(logger=logger, transport=AsyncHTTPTransport())
 
     client = AsyncClient(
         timeout=settings.llm.http_client.timeout,

ai_review/clients/gitlab/mr/client.py

@@ -1,16 +1,19 @@
-from httpx import Response
+from httpx import Response, QueryParams
 
 from ai_review.clients.gitlab.mr.schema.changes import GitLabGetMRChangesResponseSchema
 from ai_review.clients.gitlab.mr.schema.discussions import (
+    GitLabGetMRDiscussionsQuerySchema,
     GitLabGetMRDiscussionsResponseSchema,
     GitLabCreateMRDiscussionRequestSchema,
     GitLabCreateMRDiscussionResponseSchema
 )
 from ai_review.clients.gitlab.mr.schema.notes import (
+    GitLabGetMRNotesQuerySchema,
     GitLabGetMRNotesResponseSchema,
     GitLabCreateMRNoteRequestSchema,
     GitLabCreateMRNoteResponseSchema,
 )
+from ai_review.clients.gitlab.mr.types import GitLabMergeRequestsHTTPClientProtocol
 from ai_review.libs.http.client import HTTPClient
 from ai_review.libs.http.handlers import handle_http_error, HTTPClientError
 
@@ -19,7 +22,7 @@ class GitLabMergeRequestsHTTPClientError(HTTPClientError):
     pass
 
 
-class GitLabMergeRequestsHTTPClient(HTTPClient):
+class GitLabMergeRequestsHTTPClient(HTTPClient, GitLabMergeRequestsHTTPClientProtocol):
     @handle_http_error(client="GitLabMergeRequestsHTTPClient", exception=GitLabMergeRequestsHTTPClientError)
     async def get_changes_api(self, project_id: str, merge_request_id: str) -> Response:
         return await self.get(
@@ -27,15 +30,27 @@ class GitLabMergeRequestsHTTPClient(HTTPClient):
         )
 
     @handle_http_error(client="GitLabMergeRequestsHTTPClient", exception=GitLabMergeRequestsHTTPClientError)
-    async def get_notes_api(self, project_id: str, merge_request_id: str) -> Response:
+    async def get_notes_api(
+            self,
+            project_id: str,
+            merge_request_id: str,
+            query: GitLabGetMRNotesQuerySchema
+    ) -> Response:
         return await self.get(
-            f"/api/v4/projects/{project_id}/merge_requests/{merge_request_id}/notes"
+            f"/api/v4/projects/{project_id}/merge_requests/{merge_request_id}/notes",
+            query=QueryParams(**query.model_dump())
         )
 
     @handle_http_error(client="GitLabMergeRequestsHTTPClient", exception=GitLabMergeRequestsHTTPClientError)
-    async def get_discussions_api(self, project_id: str, merge_request_id: str) -> Response:
+    async def get_discussions_api(
+            self,
+            project_id: str,
+            merge_request_id: str,
+            query: GitLabGetMRDiscussionsQuerySchema
+    ) -> Response:
         return await self.get(
-            f"/api/v4/projects/{project_id}/merge_requests/{merge_request_id}/discussions"
+            f"/api/v4/projects/{project_id}/merge_requests/{merge_request_id}/discussions",
+            query=QueryParams(**query.model_dump())
         )
 
     @handle_http_error(client="GitLabMergeRequestsHTTPClient", exception=GitLabMergeRequestsHTTPClientError)
@@ -71,7 +86,8 @@ class GitLabMergeRequestsHTTPClient(HTTPClient):
             project_id: str,
             merge_request_id: str
     ) -> GitLabGetMRNotesResponseSchema:
-        response = await self.get_notes_api(project_id, merge_request_id)
+        query = GitLabGetMRNotesQuerySchema(per_page=100)
+        response = await self.get_notes_api(project_id, merge_request_id, query)
         return GitLabGetMRNotesResponseSchema.model_validate_json(response.text)
 
     async def get_discussions(
@@ -79,7 +95,8 @@ class GitLabMergeRequestsHTTPClient(HTTPClient):
             project_id: str,
             merge_request_id: str
     ) -> GitLabGetMRDiscussionsResponseSchema:
-        response = await self.get_discussions_api(project_id, merge_request_id)
+        query = GitLabGetMRDiscussionsQuerySchema(per_page=100)
+        response = await self.get_discussions_api(project_id, merge_request_id, query)
         return GitLabGetMRDiscussionsResponseSchema.model_validate_json(response.text)
 
     async def create_note(

ai_review/clients/gitlab/mr/schema/discussions.py

@@ -17,6 +17,10 @@ class GitLabDiscussionPositionSchema(BaseModel):
     new_line: int
 
 
+class GitLabGetMRDiscussionsQuerySchema(BaseModel):
+    per_page: int
+
+
 class GitLabGetMRDiscussionsResponseSchema(RootModel[list[GitLabDiscussionSchema]]):
     root: list[GitLabDiscussionSchema]
 

ai_review/clients/gitlab/mr/schema/notes.py

@@ -6,6 +6,10 @@ class GitLabNoteSchema(BaseModel):
     body: str
 
 
+class GitLabGetMRNotesQuerySchema(BaseModel):
+    per_page: int
+
+
 class GitLabGetMRNotesResponseSchema(RootModel[list[GitLabNoteSchema]]):
     root: list[GitLabNoteSchema]
 

ai_review/clients/gitlab/mr/types.py

@@ -0,0 +1,35 @@
+from typing import Protocol
+
+from ai_review.clients.gitlab.mr.schema.changes import GitLabGetMRChangesResponseSchema
+from ai_review.clients.gitlab.mr.schema.discussions import (
+    GitLabGetMRDiscussionsResponseSchema,
+    GitLabCreateMRDiscussionResponseSchema,
+    GitLabCreateMRDiscussionRequestSchema,
+)
+from ai_review.clients.gitlab.mr.schema.notes import GitLabGetMRNotesResponseSchema, GitLabCreateMRNoteResponseSchema
+
+
+class GitLabMergeRequestsHTTPClientProtocol(Protocol):
+    async def get_changes(self, project_id: str, merge_request_id: str) -> GitLabGetMRChangesResponseSchema: ...
+
+    async def get_notes(self, project_id: str, merge_request_id: str) -> GitLabGetMRNotesResponseSchema: ...
+
+    async def get_discussions(
+            self,
+            project_id: str,
+            merge_request_id: str
+    ) -> GitLabGetMRDiscussionsResponseSchema: ...
+
+    async def create_note(
+            self,
+            body: str,
+            project_id: str,
+            merge_request_id: str,
+    ) -> GitLabCreateMRNoteResponseSchema: ...
+
+    async def create_discussion(
+            self,
+            project_id: str,
+            merge_request_id: str,
+            request: GitLabCreateMRDiscussionRequestSchema,
+    ) -> GitLabCreateMRDiscussionResponseSchema: ...

ai_review/clients/openai/client.py

@@ -26,7 +26,7 @@ class OpenAIHTTPClient(HTTPClient):
 def get_openai_http_client() -> OpenAIHTTPClient:
     logger = get_logger("OPENAI_HTTP_CLIENT")
     logger_event_hook = LoggerEventHook(logger=logger)
-    retry_transport = RetryTransport(transport=AsyncHTTPTransport())
+    retry_transport = RetryTransport(logger=logger, transport=AsyncHTTPTransport())
 
     client = AsyncClient(
         timeout=settings.llm.http_client.timeout,

ai_review/config.py

@@ -12,6 +12,7 @@ from ai_review.libs.config.base import (
     get_yaml_config_file_or_default,
     get_json_config_file_or_default
 )
+from ai_review.libs.config.core import CoreConfig
 from ai_review.libs.config.llm import LLMConfig
 from ai_review.libs.config.logger import LoggerConfig
 from ai_review.libs.config.prompt import PromptConfig
@@ -36,6 +37,7 @@ class Settings(BaseSettings):
 
     llm: LLMConfig
     vcs: VCSConfig
+    core: CoreConfig = CoreConfig()
     prompt: PromptConfig = PromptConfig()
     review: ReviewConfig = ReviewConfig()
     logger: LoggerConfig = LoggerConfig()

ai_review/libs/asynchronous/gather.py

@@ -1,14 +1,17 @@
 import asyncio
 from typing import Awaitable, Iterable, TypeVar
 
+from ai_review.config import settings
+
 T = TypeVar("T")
 
 
-async def bounded_gather(coroutines: Iterable[Awaitable[T]], concurrency: int = 32) -> tuple[T, ...]:
-    sem = asyncio.Semaphore(concurrency)
+async def bounded_gather(coroutines: Iterable[Awaitable[T]]) -> tuple[T, ...]:
+    sem = asyncio.Semaphore(settings.core.concurrency)
 
     async def wrap(coro: Awaitable[T]) -> T:
         async with sem:
             return await coro
 
-    return await asyncio.gather(*(wrap(coroutine) for coroutine in coroutines), return_exceptions=True)
+    results = await asyncio.gather(*(wrap(coroutine) for coroutine in coroutines), return_exceptions=True)
+    return tuple(results)

ai_review/libs/config/core.py

@@ -0,0 +1,5 @@
+from pydantic import BaseModel
+
+
+class CoreConfig(BaseModel):
+    concurrency: int = 7

ai_review/libs/http/event_hooks/logger.py

@@ -1,12 +1,15 @@
-from logging import Logger
+from typing import TYPE_CHECKING
 
 from httpx import Request, Response
 
 from ai_review.libs.http.event_hooks.base import BaseEventHook
 
+if TYPE_CHECKING:
+    from loguru import Logger
+
 
 class LoggerEventHook(BaseEventHook):
-    def __init__(self, logger: Logger):
+    def __init__(self, logger: "Logger"):
         self.logger = logger
 
     async def request(self, request: Request):

ai_review/libs/http/transports/retry.py

@@ -1,12 +1,17 @@
 import asyncio
 from http import HTTPStatus
+from typing import TYPE_CHECKING
 
 from httpx import Request, Response, AsyncBaseTransport
 
+if TYPE_CHECKING:
+    from loguru import Logger
+
 
 class RetryTransport(AsyncBaseTransport):
     def __init__(
             self,
+            logger: "Logger",
             transport: AsyncBaseTransport,
             max_retries: int = 5,
             retry_delay: float = 0.5,
@@ -17,18 +22,30 @@ class RetryTransport(AsyncBaseTransport):
                     HTTPStatus.INTERNAL_SERVER_ERROR,
             )
     ):
+        self.logger = logger
         self.transport = transport
         self.max_retries = max_retries
         self.retry_delay = retry_delay
         self.retry_status_codes = retry_status_codes
 
     async def handle_async_request(self, request: Request) -> Response:
-        response: Response | None = None
-        for _ in range(self.max_retries):
-            response = await self.transport.handle_async_request(request)
-            if response.status_code not in self.retry_status_codes:
-                return response
+        last_response: Response | None = None
+        for attempt in range(self.max_retries):
+            last_response = await self.transport.handle_async_request(request)
+            if last_response.status_code not in self.retry_status_codes:
+                return last_response
+
+            self.logger.warning(
+                f"Attempt {attempt}/{self.max_retries} failed "
+                f"with status={last_response.status_code} for {request.method} {request.url}. "
+                f"Retrying in {self.retry_delay:.1f}s..."
+            )
 
             await asyncio.sleep(self.retry_delay)
 
-        return response
+        self.logger.error(
+            f"All {self.max_retries} attempts failed for "
+            f"{request.method} {request.url} (last status={last_response.status_code})"
+        )
+
+        return last_response

ai_review/services/artifacts/service.py

@@ -6,11 +6,12 @@ from ai_review.config import settings
 from ai_review.libs.logger import get_logger
 from ai_review.services.artifacts.schema import LLMArtifactSchema
 from ai_review.services.artifacts.tools import make_artifact_id
+from ai_review.services.artifacts.types import ArtifactsServiceProtocol
 
 logger = get_logger("ARTIFACTS_SERVICE")
 
 
-class ArtifactsService:
+class ArtifactsService(ArtifactsServiceProtocol):
     @classmethod
     async def save_llm_interaction(cls, prompt: str, prompt_system: str, response: str | None = None) -> str | None:
         if not settings.artifacts.llm_enabled:

ai_review/services/artifacts/types.py

@@ -0,0 +1,20 @@
+from pathlib import Path
+from typing import Protocol
+
+
+class ArtifactsServiceProtocol(Protocol):
+    async def save_llm_interaction(
+            self,
+            prompt: str,
+            prompt_system: str,
+            response: str | None = None
+    ) -> str | None:
+        ...
+
+    async def save_artifact(
+            self,
+            file: Path,
+            content: str,
+            kind: str = "artifact"
+    ) -> Path | None:
+        ...

ai_review/services/cost/service.py

@@ -1,12 +1,13 @@
 from ai_review.config import settings
 from ai_review.libs.logger import get_logger
 from ai_review.services.cost.schema import CostReportSchema
+from ai_review.services.cost.types import CostServiceProtocol
 from ai_review.services.llm.types import ChatResultSchema
 
 logger = get_logger("COST_SERVICE")
 
 
-class CostService:
+class CostService(CostServiceProtocol):
     def __init__(self):
         self.pricing = settings.llm.load_pricing()
         self.reports: list[CostReportSchema] = []

ai_review/services/cost/types.py

@@ -0,0 +1,12 @@
+from typing import Protocol
+
+from ai_review.services.cost.schema import CostReportSchema
+from ai_review.services.llm.types import ChatResultSchema
+
+
+class CostServiceProtocol(Protocol):
+    def calculate(self, result: ChatResultSchema) -> CostReportSchema | None:
+        ...
+
+    def aggregate(self) -> CostReportSchema | None:
+        ...

ai_review/services/diff/service.py

@@ -16,12 +16,13 @@ from ai_review.services.diff.renderers import (
 )
 from ai_review.services.diff.schema import DiffFileSchema
 from ai_review.services.diff.tools import find_diff_file
+from ai_review.services.diff.types import DiffServiceProtocol
 from ai_review.services.git.types import GitServiceProtocol
 
 logger = get_logger("DIFF_SERVICE")
 
 
-class DiffService:
+class DiffService(DiffServiceProtocol):
     @classmethod
     def parse(cls, raw_diff: str) -> Diff:
         if not raw_diff.strip():

ai_review/services/diff/types.py

@@ -0,0 +1,28 @@
+from typing import Protocol
+
+from ai_review.libs.diff.models import Diff
+from ai_review.services.diff.schema import DiffFileSchema
+from ai_review.services.git.types import GitServiceProtocol
+
+
+class DiffServiceProtocol(Protocol):
+    def parse(self, raw_diff: str) -> Diff:
+        ...
+
+    def render_file(
+            self,
+            file: str,
+            raw_diff: str,
+            base_sha: str | None = None,
+            head_sha: str | None = None,
+    ) -> DiffFileSchema:
+        ...
+
+    def render_files(
+            self,
+            git: GitServiceProtocol,
+            files: list[str],
+            base_sha: str,
+            head_sha: str,
+    ) -> list[DiffFileSchema]:
+        ...

ai_review/services/hook/__init__.py

@@ -0,0 +1,5 @@
+from ai_review.services.hook.service import HookService
+
+hook = HookService()
+
+__all__ = ["hook"]