Merge pull request #30 from elementsinteractive/jsonformatter

feat(#18): add json formatter and new `--output-format` option and config

Author: scastlara

README.md

@@ -296,6 +296,7 @@ lgtm uses a `.toml` file to configure how it works. It will autodetect a `lgtm.t
 - **model**: Choose which AI model you want lgtm to use.
 - **model_url**: When not using one of the specific supported models from the providers mentioned above, you can pass a custom url where the model is deployed.
 - **exclude**: Instruct lgtm to ignore certain files. This is important to reduce noise in reviews, but also to reduce the amount of tokens used for each review (and to avoid running into token limits). You can specify file patterns (`exclude = ["*.md", "package-lock.json"]`)
+- **output_format**: Format of the terminal output of lgtm. Can be `pretty` (default), `json`, and `markdown`.
 - **silent**: Do not print the review in the terminal.
 - **publish**: If `true`, it will post the review as comments on the PR page.
 - **ai_api_key**: API key to call the selected AI model. Can be given as a CLI argument, or as an environment variable (`LGTM_AI_API_KEY`).

src/lgtm_ai/__main__.py

@@ -2,7 +2,7 @@
 import logging
 from collections.abc import Callable
 from importlib.metadata import version
-from typing import get_args
+from typing import Any, assert_never, get_args
 
 import click
 import rich
@@ -13,10 +13,12 @@
     get_summarizing_agent_with_settings,
 )
 from lgtm_ai.ai.schemas import AgentSettings, CommentCategory, SupportedAIModels, SupportedAIModelsList
-from lgtm_ai.base.schemas import PRUrl
+from lgtm_ai.base.schemas import OutputFormat, PRUrl
 from lgtm_ai.config.handler import ConfigHandler, PartialConfig
+from lgtm_ai.formatters.base import Formatter
+from lgtm_ai.formatters.json import JsonFormatter
 from lgtm_ai.formatters.markdown import MarkDownFormatter
-from lgtm_ai.formatters.terminal import TerminalFormatter
+from lgtm_ai.formatters.pretty import PrettyFormatter
 from lgtm_ai.git_client.utils import get_git_client
 from lgtm_ai.review import CodeReviewer
 from lgtm_ai.review.guide import ReviewGuideGenerator
@@ -25,14 +27,15 @@
     parse_pr_url,
     validate_model_url,
 )
+from rich.console import Console
 from rich.logging import RichHandler
 
 __version__ = version("lgtm-ai")
 
 logging.basicConfig(
     format="%(message)s",
     datefmt="[%X]",
-    handlers=[RichHandler(rich_tracebacks=True, show_path=False)],
+    handlers=[RichHandler(rich_tracebacks=True, show_path=False, console=Console(stderr=True))],
 )
 logger = logging.getLogger("lgtm")
 
@@ -68,6 +71,7 @@ def _common_options[**P, T](func: Callable[P, T]) -> Callable[P, T]:
         help="Exclude files from the review. If not provided, all files in the PR will be reviewed. Uses UNIX-style wildcards.",
     )
     @click.option("--publish", is_flag=True, help="Publish the review or guide to the git service.")
+    @click.option("--output-format", type=click.Choice([format.value for format in OutputFormat]))
     @click.option("--silent", is_flag=True, help="Do not print the review or guide to the console.")
     @click.option(
         "--ai-retries",
@@ -104,6 +108,7 @@ def review(
     config: str | None,
     exclude: tuple[str, ...],
     publish: bool,
+    output_format: OutputFormat | None,
     silent: bool,
     ai_retries: int | None,
     verbose: int,
@@ -125,6 +130,7 @@ def review(
             model=model,
             model_url=model_url,
             publish=publish,
+            output_format=output_format,
             silent=silent,
             ai_retries=ai_retries,
         ),
@@ -146,10 +152,10 @@ def review(
 
     if not resolved_config.silent:
         logger.info("Printing review to console")
-        terminal_formatter = TerminalFormatter()
-        rich.print(terminal_formatter.format_review_summary_section(review))
+        formatter, printer = _get_formatter_and_printer(resolved_config.output_format)
+        printer(formatter.format_review_summary_section(review))
         if review.review_response.comments:
-            rich.print(terminal_formatter.format_review_comments_section(review.review_response.comments))
+            printer(formatter.format_review_comments_section(review.review_response.comments))
 
     if resolved_config.publish:
         logger.info("Publishing review to git service")
@@ -168,6 +174,7 @@ def guide(
     config: str | None,
     exclude: tuple[str, ...],
     publish: bool,
+    output_format: OutputFormat | None,
     silent: bool,
     ai_retries: int | None,
     verbose: int,
@@ -185,6 +192,7 @@ def guide(
             model=model,
             model_url=model_url,
             publish=publish,
+            output_format=output_format,
             silent=silent,
             ai_retries=ai_retries,
         ),
@@ -204,8 +212,8 @@ def guide(
 
     if not resolved_config.silent:
         logger.info("Printing review to console")
-        terminal_formatter = TerminalFormatter()
-        rich.print(terminal_formatter.format_guide(guide))
+        formatter, printer = _get_formatter_and_printer(resolved_config.output_format)
+        printer(formatter.format_guide(guide))
 
     if resolved_config.publish:
         logger.info("Publishing review guide to git service")
@@ -221,3 +229,15 @@ def _set_logging_level(logger: logging.Logger, verbose: int) -> None:
     else:
         logger.setLevel(logging.DEBUG)
     logger.info("Logging level set to %s", logging.getLevelName(logger.level))
+
+
+def _get_formatter_and_printer(output_format: OutputFormat) -> tuple[Formatter[Any], Callable[[Any], None]]:
+    """Get the formatter and the print method based on the output format."""
+    if output_format == OutputFormat.pretty:
+        return PrettyFormatter(), rich.print
+    elif output_format == OutputFormat.markdown:
+        return MarkDownFormatter(), print
+    elif output_format == OutputFormat.json:
+        return JsonFormatter(), print
+    else:
+        assert_never(output_format)

src/lgtm_ai/ai/schemas.py

@@ -138,8 +138,7 @@ class GuideResponse(BaseModel):
     references: Annotated[list[GuideReference], Field(description="References to external resources")]
 
 
-@dataclass(frozen=True)
-class PublishMetadata:
+class PublishMetadata(BaseModel):
     model_name: str
     usages: list[Usage]
 
@@ -152,17 +151,15 @@ def uuid(self) -> str:
         return uuid4().hex
 
 
-@dataclass(frozen=True, slots=True)
-class Review:
+class Review(BaseModel):
     """Represent a full code review performed by any AI agent."""
 
     pr_diff: PRDiff
     review_response: ReviewResponse
     metadata: PublishMetadata
 
 
-@dataclass(frozen=True, slots=True)
-class ReviewGuide:
+class ReviewGuide(BaseModel):
     """Represent a code review guide generated by the AI agent."""
 
     pr_diff: PRDiff

src/lgtm_ai/base/schemas.py

@@ -1,4 +1,5 @@
 from dataclasses import dataclass
+from enum import StrEnum
 from typing import Literal
 
 
@@ -8,3 +9,9 @@ class PRUrl:
     repo_path: str
     pr_number: int
     source: Literal["github", "gitlab"]
+
+
+class OutputFormat(StrEnum):
+    pretty = "pretty"
+    json = "json"
+    markdown = "markdown"

src/lgtm_ai/config/handler.py

@@ -7,6 +7,7 @@
 from typing import Any, ClassVar, Literal, cast, get_args, overload
 
 from lgtm_ai.ai.schemas import CommentCategory, SupportedAIModels
+from lgtm_ai.base.schemas import OutputFormat
 from lgtm_ai.config.constants import DEFAULT_AI_MODEL
 from lgtm_ai.config.exceptions import (
     ConfigFileNotFoundError,
@@ -31,6 +32,7 @@ class PartialConfig(BaseModel):
     categories: tuple[CommentCategory, ...] | None = None
     exclude: tuple[str, ...] | None = None
     publish: bool = False
+    output_format: OutputFormat | None = None
     silent: bool = False
     ai_retries: int | None = None
 
@@ -63,6 +65,9 @@ class ResolvedConfig(BaseModel):
     publish: bool = False
     """Publish the review to the git service as comments."""
 
+    output_format: OutputFormat = OutputFormat.pretty
+    """Output format for the review, defaults to pretty."""
+
     silent: bool = False
     """Suppress terminal output."""
 
@@ -125,6 +130,7 @@ def _parse_config_file(self) -> PartialConfig:
                 categories=config_data.get("categories", None),
                 exclude=config_data.get("exclude", None),
                 publish=config_data.get("publish", False),
+                output_format=config_data.get("output_format", None),
                 silent=config_data.get("silent", False),
                 ai_retries=config_data.get("ai_retries", None),
             )
@@ -181,6 +187,7 @@ def _parse_cli_args(self) -> PartialConfig:
             exclude=self.cli_args.exclude or None,
             ai_api_key=self.cli_args.ai_api_key or None,
             git_api_key=self.cli_args.git_api_key or None,
+            output_format=self.cli_args.output_format or None,
             silent=self.cli_args.silent,
             publish=self.cli_args.publish,
             ai_retries=self.cli_args.ai_retries or None,
@@ -221,6 +228,7 @@ def _resolve_from_multiple_sources(
             ),
             model_url=from_cli.model_url or from_file.model_url,
             publish=from_cli.publish or from_file.publish,
+            output_format=from_cli.output_format or from_file.output_format or OutputFormat.pretty,
             silent=from_cli.silent or from_file.silent,
             ai_retries=from_cli.ai_retries or from_file.ai_retries,
             git_api_key=self.resolver.resolve_string_field("git_api_key", from_cli=from_cli, from_env=from_env),

src/lgtm_ai/formatters/json.py

@@ -0,0 +1,23 @@
+from lgtm_ai.ai.schemas import Review, ReviewComment, ReviewGuide
+from lgtm_ai.formatters.base import Formatter
+
+
+class JsonFormatter(Formatter[str]):
+    def format_review_summary_section(self, review: Review, comments: list[ReviewComment] | None = None) -> str:
+        """Format the **whole** review as JSON."""
+        return review.model_dump_json(indent=2)
+
+    def format_review_comments_section(self, comments: list[ReviewComment]) -> str:
+        """No-op.
+
+        Formatting the comments section alone as JSON does not really make sense, so this is a no-op.
+        """
+        return ""
+
+    def format_review_comment(self, comment: ReviewComment, *, with_footer: bool = True) -> str:
+        """Format a single comment as JSON."""
+        return comment.model_dump_json(indent=2)
+
+    def format_guide(self, guide: ReviewGuide) -> str:
+        """Format the review guide as JSON."""
+        return guide.model_dump_json(indent=2)

src/lgtm_ai/formatters/terminal.py src/lgtm_ai/formatters/pretty.pysrc/lgtm_ai/formatters/terminal.py renamed to src/lgtm_ai/formatters/pretty.py

@@ -12,7 +12,7 @@
 logger = logging.getLogger("lgtm")
 
 
-class TerminalFormatter(Formatter[Panel | Layout]):
+class PrettyFormatter(Formatter[Panel | Layout]):
     def format_review_summary_section(self, review: Review, comments: list[ReviewComment] | None = None) -> Panel:
         if comments:
             logger.warning("Comments are not supported in the terminal formatter summary section")

src/lgtm_ai/git_client/gitlab.py

@@ -50,7 +50,7 @@ def get_diff_from_url(self, pr_url: PRUrl) -> PRDiff:
             raise PullRequestDiffError from err
 
         return PRDiff(
-            diff.id,
+            id=diff.id,
             diff=self._parse_gitlab_git_diff(diff.diffs),
             changed_files=[change["new_path"] for change in diff.diffs],
             target_branch=pr.target_branch,

src/lgtm_ai/git_client/schemas.py

@@ -1,25 +1,21 @@
-from dataclasses import dataclass
-
+from groq import BaseModel
 from lgtm_ai.git_parser.parser import DiffResult
 
 
-@dataclass(frozen=True, slots=True)
-class PRDiff:
+class PRDiff(BaseModel):
     id: int
     diff: list[DiffResult]
     changed_files: list[str]
     target_branch: str
     source_branch: str
 
 
-@dataclass(frozen=True, slots=True)
-class PRContextFileContents:
+class PRContextFileContents(BaseModel):
     file_path: str
     content: str
 
 
-@dataclass(slots=True)
-class PRContext:
+class PRContext(BaseModel):
     """Represents the context a reviewer might need when reviewing PRs.
 
     At the moment, it is just the contents of the files that are changed in the PR.
@@ -31,10 +27,9 @@ def __bool__(self) -> bool:
         return bool(self.file_contents)
 
     def add_file(self, file_path: str, content: str) -> None:
-        self.file_contents.append(PRContextFileContents(file_path, content))
+        self.file_contents.append(PRContextFileContents(file_path=file_path, content=content))
 
 
-@dataclass(frozen=True, slots=True)
-class PRMetadata:
+class PRMetadata(BaseModel):
     title: str
     description: str

src/lgtm_ai/review/reviewer.py

@@ -76,7 +76,7 @@ def review_pull_request(self, pr_url: PRUrl) -> Review:
         )
 
         return Review(
-            pr_diff,
-            final_res.output,
+            pr_diff=pr_diff,
+            review_response=final_res.output,
             metadata=PublishMetadata(model_name=self.config.model, usages=[initial_usage, final_usage]),
         )