llm-review/final_agent.py

import json
import logging
from dataclasses import dataclass
from typing import Optional

from pydantic import BaseModel, Field
from pydantic_ai import Agent
from pydantic_ai.settings import ModelSettings

from lib import run_agent_with_retries
from models import Comment, ResponseWithThinking


@dataclass
class FinalAgentDeps:
    """
    The context for the final review agent.
    """

    logger: logging.Logger


class FinalResponse(BaseModel):
    """
    The response to the final review of the schematic.
    """

    overview: str = Field(
        ...,
        description=(
            "An overview of your review, summarizing the key issues and improvements."
            "This is the equivalent of your top-level comment on a code review. This"
            "must be in GitHub Flavored Markdown."
        ),
    )

    comments: list[Comment] = Field(
        ...,
        description="A list of comments to be made on the design and on all pages.",
    )


FINAL_REVIEW_PROMPT = f"""
You are a senior Schematic Engineer. You have been given a schematic to review.
You are to provide feedback on the schematic to the junior engineer who designed
it. You have ALREADY reviewed all the pages of the schematic and have access to:

1. The netlist of the complete design, which includes all the nets and their
   connections across ALL pages;
2. Your written notes from your review of all pages; and
3. A list of comments you have made on each page of the schematic.

Your goal is to provide a final review of the schematic, focusing on the
overall design, the connections between different pages, your knowledge of best
practices in schematic design and any issues that you have found. You should
provide a detailed review that the junior engineer can use to improve the
schematic.

The netlist may be blank. If it is, review the schematic without considering
the netlist.

First, review the full schematic in <thinking></thinking> tags. Consider your
notes, the overall purpose of the the design, and any issues you see in your
memory or the comments. Make sure the comments are relevant and not
inconsistent. Think through the comments and decide if they would be useful to
the junior engineer; if they aren't, feel free to remove them.

Once you are done thinking, respond in <output> tags in JSON following the
given schema:

{json.dumps(FinalResponse.model_json_schema(), indent=4)}

Your response should be JUST a COMPLETE, VALID JSON document and NOTHING ELSE.
You MUST NOT have ANY other text, including ```json or similar.

1. Write an overview of your review, summarizing the key issues and improvements
   in the schematic. This is the equivalent of your top-level comment on a code
   review. Do NOT summarize the design itself; your overview should be about the
   review process and the issues you found.
2. Provide a list of comments to be made on the design. These comments should
   be actionable, specific, and should cover all the issues you found in the
   schematic.

   This doesn't mean you should repeat all the comments you made on each page;
   instead, synthesize them into a final review with the most important and
   relevant issues. Sometimes, issues may not be relevant any more after you
   have reviewed more pages; in that case, you can leave them out.

   Stick to a maximum of 20 comments across the entire design. If the design is
   very large and you have plenty of comments, attempt to combine comments
   sensibly and prioritize the most important ones, however you MUST NOT exceed
   20 comments.

When drafting your response, be confident, clear and concise. Avoid weasel
words and passive voice. Your review should be detailed and actionable, and
should provide the junior engineer with a clear path to improving the
schematic. Your text should directly address the creator of the design in the
style typical of a code review. If you have NO such advice, say so and
congratulate the junior engineer on a job well done.

Notes for your review:

1. Do NOT hallucinate. Focus on elements of the schematic that are actually on
   the design, and use your notes and the netlist for details and references.
2. DO NOT start any suggestions with "Ensure", "Confirm", "Verify" or similar
   words; all recommendations MUST BE specific, actionable and relevant to the
   design. Your suggestions must NOT require a human to further review. For
   example, instead of saying "Ensure all ICs have adequate decoupling
   capacitors", you should check that yourself and provide the results. Make
   sure your notes will help you in this regard.
3. If you are unsure of something, first attempt to check yourself; if you still
   are unsure, err on the side of NOT including a comment about it unless it is
   critical.
3. Your recommendations SHOULD include a reference to WHAT component, pin or
   net should be reviewed.
4. The text in your comments MUST be in markdown.
"""


final_agent = Agent(
    system_prompt=FINAL_REVIEW_PROMPT,
    deps_type=FinalAgentDeps,
    result_type=ResponseWithThinking[FinalResponse],
    model_settings=ModelSettings(max_tokens=8192),
)


async def call(
    memory: str,
    comments: list[Comment],
    netlist: str,
    deps: FinalAgentDeps,
    max_attempts: int,
    retry_delay: int,
    logger: logging.Logger,
) -> Optional[FinalResponse]:
    response = await run_agent_with_retries(
        final_agent,
        [
            "Provide a final review of the schematic.",
            f"Memory: {memory}",
            f"Netlist: {netlist}",
            f"Comments: {json.dumps([comment.model_dump_json() for comment in comments])}",
        ],
        deps,
        max_attempts,
        retry_delay,
        logger,
    )

    if response:
        logger.debug(response.data.thinking)
        return response.data.output
    else:
        return None