160 lines
6.3 KiB
Python
160 lines
6.3 KiB
Python
import json
|
|
import logging
|
|
from typing import Optional
|
|
|
|
from pydantic import BaseModel, Field
|
|
from pydantic_ai import Agent, BinaryContent
|
|
from pydantic_ai.settings import ModelSettings
|
|
|
|
from lib import run_agent_with_retries
|
|
from models import Comment, ResponseWithThinking
|
|
|
|
|
|
class StepResponse(BaseModel):
|
|
"""
|
|
The response to a single step in the design review process.
|
|
"""
|
|
|
|
memory: str = Field(
|
|
...,
|
|
description=(
|
|
"The memory to be used for the next page in the review, including "
|
|
"your notes from this page."
|
|
),
|
|
)
|
|
|
|
comments: list[Comment] = Field(
|
|
...,
|
|
description="A list of comments to be made on the design ONLY for this page.",
|
|
)
|
|
|
|
|
|
STEP_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 currently are reviewing a single page of the schematic, and are
|
|
proceeding in a page-by-page manner. For this page, you have access to:
|
|
|
|
1. The netlist of the complete design, which includes all the nets and their
|
|
connections across ALL pages;
|
|
2. A high-resolution image of the schematic page you are reviewing.
|
|
3. A JSON representation of the schematic page you are reviewing.
|
|
4. Written notes (memory) from your review so far.
|
|
5. A set of comments you have made on previous pages of the schematic.
|
|
|
|
Your goal is to review this page, keeping in mind the overall design, your
|
|
knowledge of best practices in schematic design, and your memory so far. The
|
|
netlist may be blank. If it is, review the schematic without considering the
|
|
netlist.
|
|
|
|
First, review the page in <thinking></thinking> tags. Make sure to think about
|
|
the purpose of the design and your insights and knowledge of engineering. You
|
|
should also think about the overall design and how this page fits into it.
|
|
Then, identify key components, nets and connections, noting down potential
|
|
issues or improvements along the way.
|
|
|
|
Once you understand the design, you can start thinking about your output. For
|
|
this, start by making a shortlist of issues or improvements you want to comment
|
|
on. On each element of the shortlist, write a paragraph starting with your
|
|
reasoning behind the issue or comment, then a summary, and finally a
|
|
recommendation if relevant. This should look something like:
|
|
|
|
According to its spec sheet, this chip requires a 100K resistor to be in
|
|
series at pin 6. However, I see that the schematic has a 10K resistor on
|
|
U2.6. This could lead to instability in the output signal. I should
|
|
recommend replacing the 10K resistor (R18) with a 100K resistor.
|
|
|
|
Sometimes, you may be missing details that are on other pages. In that case,
|
|
make sure to make a note to yourself in your memory such that you will be able
|
|
to verify whether or not there is an issue and therefore make a relevant
|
|
comment. You can also create a comment that notes that this may be invalidated
|
|
in another page; if you run into such comments, double check those and remove
|
|
them if they are no longer valid.
|
|
|
|
In your list of elements, you SHOULD be specific: if you find an issue with a
|
|
component, refer to it by its designator; if you find an issue with a pin,
|
|
refer to it in the form `U1A.1` where `U1A` is the designator of the component
|
|
and `1` is the pin number; do NOT refer to the component itself.
|
|
|
|
Once you are done thinking, respond in <output> tags in JSON following the
|
|
given schema:
|
|
|
|
{json.dumps(StepResponse.model_json_schema(), indent=4)}
|
|
|
|
Your output MUST be JUST a COMPLETE, VALID JSON document and NOTHING ELSE.
|
|
You MUST NOT have ANY other text, including ```json or similar.
|
|
|
|
1. Your memory should be comprehensive, detailed and accurate, and should be
|
|
sufficient to conduct the final review of this page WITHOUT access to the
|
|
JSON or the image.
|
|
2. Your comments should be actionable, specific, and should cover all the
|
|
issues you find on the page.
|
|
|
|
Eventually, once you have reviewed all pages, you will provide a final review
|
|
covering the entire schematic, so make sure to keep track of all the important
|
|
issues you find. At that time, you will only have access to the netlist and your
|
|
written notes, so you MUST ensure that your notes are detailed enough to cover
|
|
all the issues you find, with potential suggestions for solutions.
|
|
|
|
When drafting your response, be confident, clear and concise. Avoid weasel
|
|
words and passive voice. Your comments 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.
|
|
|
|
Notes for your review:
|
|
|
|
1. Do NOT hallucinate. Focus on elements of the schematic that are actually on
|
|
the page, and make sure to include details to keep track of these elements
|
|
in your notes.
|
|
2. Do NOT provide feedback on the overall design of the schematic yet. You will
|
|
do that once you have reviewed all pages.
|
|
3. Use as much written memory as you need to keep track of details.
|
|
4. 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.
|
|
"""
|
|
|
|
|
|
step_agent = Agent(
|
|
system_prompt=STEP_PROMPT,
|
|
result_type=ResponseWithThinking[StepResponse],
|
|
model_settings=ModelSettings(max_tokens=8192),
|
|
)
|
|
|
|
|
|
async def call(
|
|
memory: str,
|
|
comments: list[Comment],
|
|
page_json: str,
|
|
page_image: bytes,
|
|
netlist: str,
|
|
page_number: int,
|
|
total_pages: int,
|
|
max_attempts: int,
|
|
retry_delay: int,
|
|
logger: logging.Logger,
|
|
) -> Optional[StepResponse]:
|
|
response = await run_agent_with_retries(
|
|
step_agent,
|
|
[
|
|
f"Review this page. You are on page {page_number} of {total_pages}.",
|
|
f"Page JSON: {page_json}",
|
|
BinaryContent(page_image, media_type="image/png"),
|
|
f"Netlist: {netlist}",
|
|
f"Memory: {memory}",
|
|
f"Comments: {json.dumps([comment.model_dump_json() for comment in comments])}",
|
|
],
|
|
deps=None,
|
|
max_attempts=max_attempts,
|
|
retry_delay=retry_delay,
|
|
logger=logger,
|
|
)
|
|
|
|
if response:
|
|
logger.debug(response.data.thinking)
|
|
return response.data.output
|
|
else:
|
|
return None
|