Skip to content

Commit e672c0f

Browse files
Update
1 parent 9d6b09d commit e672c0f

11 files changed

+84
-184
lines changed

src/migrations/versions/52d271c30ee5_migration.py src/migrations/versions/b2a0463d9b91_migration.py

+6-5
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,16 @@
11
"""Migration
22
3-
Revision ID: 52d271c30ee5
3+
Revision ID: b2a0463d9b91
44
Revises: e0fcdc14251c
5-
Create Date: 2025-02-28 23:53:40.954667
5+
Create Date: 2025-03-03 17:34:37.242326
66
77
"""
88
from alembic import op
99
import sqlalchemy as sa
1010

1111

1212
# revision identifiers, used by Alembic.
13-
revision = '52d271c30ee5'
13+
revision = 'b2a0463d9b91'
1414
down_revision = 'e0fcdc14251c'
1515
branch_labels = None
1616
depends_on = None
@@ -26,12 +26,13 @@ def upgrade():
2626
sa.Column('repo', sa.String(), nullable=False),
2727
sa.Column('run_id', sa.Integer(), nullable=False),
2828
sa.Column('iterations', sa.Integer(), nullable=False),
29+
sa.Column('original_pr_url', sa.String(), nullable=False),
2930
sa.ForeignKeyConstraint(['run_id'], ['run_state.id'], ondelete='CASCADE'),
3031
sa.PrimaryKeyConstraint('id'),
31-
sa.UniqueConstraint('provider', 'pr_id', 'repo', 'owner')
32+
sa.UniqueConstraint('provider', 'pr_id', 'repo', 'owner', 'original_pr_url')
3233
)
3334
with op.batch_alter_table('codegen_unit_test_generation_pr_context_to_run_id', schema=None) as batch_op:
34-
batch_op.create_index('ix_autofix_repo_owner_pr_id', ['owner', 'repo', 'pr_id'], unique=False)
35+
batch_op.create_index('ix_autofix_repo_owner_pr_id', ['owner', 'repo', 'pr_id', 'original_pr_url'], unique=False)
3536

3637
# ### end Alembic commands ###
3738

src/seer/app.py

+6-3
Original file line numberDiff line numberDiff line change
@@ -290,7 +290,9 @@ def codegen_pr_review_state_endpoint(
290290

291291
@json_api(blueprint, "/v1/automation/codegen/retry-unit-tests")
292292
def codegen_retry_unit_tests_endpoint(data: CodegenBaseRequest) -> CodegenUnitTestsResponse:
293-
return codegen_retry_unittest(data)
293+
raise NotImplementedError("Retry unit tests is not implemented yet.")
294+
# TODO: Finish implementation of retry-unit-tests
295+
# return codegen_retry_unittest(data)
294296

295297

296298
@json_api(blueprint, "/v1/automation/codecov-request")
@@ -308,8 +310,9 @@ def codecov_request_endpoint(
308310
return codegen_pr_review_endpoint(data.data)
309311
elif data.request_type == "unit-tests":
310312
return codegen_unit_tests_endpoint(data.data)
311-
elif data.request_type == "retry-unit-tests":
312-
return codegen_retry_unit_tests_endpoint(data.data)
313+
# TODO: Finish implementation of retry-unit-tests
314+
# elif data.request_type == "retry-unit-tests":
315+
# return codegen_retry_unit_tests_endpoint(data.data)
313316

314317
raise ValueError(f"Unsupported request_type: {data.request_type}")
315318

src/seer/automation/codebase/repo_client.py

+10-3
Original file line numberDiff line numberDiff line change
@@ -622,11 +622,17 @@ def get_pr_head_sha(self, pr_url: str) -> str:
622622
data.raise_for_status() # Raise an exception for HTTP errors
623623
return data.json()["head"]["sha"]
624624

625-
def post_unit_test_reference_to_original_pr(self, original_pr_url: str, unit_test_pr_url: str):
625+
def post_unit_test_reference_to_original_pr(
626+
self,
627+
original_pr_url: str,
628+
unit_test_pr_url: str,
629+
type: str = RepoClientType.CODECOV_UNIT_TEST,
630+
):
626631
original_pr_id = int(original_pr_url.split("/")[-1])
627632
repo_name = original_pr_url.split("github.com/")[1].split("/pull")[0]
628633
url = f"https://api.github.com/repos/{repo_name}/issues/{original_pr_id}/comments"
629-
comment = f"Sentry has generated a new [PR]({unit_test_pr_url}) with unit tests for this PR. View the new PR({unit_test_pr_url}) to review the changes."
634+
gh_app = "Sentry" if type == RepoClientType.CODECOV_UNIT_TEST else "Codecov"
635+
comment = f"{gh_app} has generated a new [PR]({unit_test_pr_url}) with unit tests for this PR. View the new PR({unit_test_pr_url}) to review the changes."
630636
params = {"body": comment}
631637
headers = self._get_auth_headers()
632638
response = requests.post(url, headers=headers, json=params)
@@ -637,7 +643,8 @@ def post_unit_test_not_generated_message_to_original_pr(self, original_pr_url: s
637643
original_pr_id = int(original_pr_url.split("/")[-1])
638644
repo_name = original_pr_url.split("github.com/")[1].split("/pull")[0]
639645
url = f"https://api.github.com/repos/{repo_name}/issues/{original_pr_id}/comments"
640-
comment = "Sentry has determined that unit tests already exist on this PR or that they are not necessary."
646+
gh_app = "Sentry" if type == RepoClientType.CODECOV_UNIT_TEST else "Codecov"
647+
comment = f"{gh_app} has determined that unit tests are not necessary for this PR."
641648
params = {"body": comment}
642649
headers = self._get_auth_headers()
643650
response = requests.post(url, headers=headers, json=params)

src/seer/automation/codegen/codegen_context.py

+8-29
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,12 @@
33
from seer.automation.agent.models import Message
44
from seer.automation.codebase.repo_client import RepoClient, RepoClientType
55
from seer.automation.codegen.codegen_event_manager import CodegenEventManager
6-
from seer.automation.codegen.models import CodegenContinuation, UnitTestRunMemory
6+
from seer.automation.codegen.models import CodegenContinuation
77
from seer.automation.codegen.state import CodegenContinuationState
88
from seer.automation.models import RepoDefinition
99
from seer.automation.pipeline import PipelineContext
1010
from seer.automation.state import DbStateRunTypes
11-
from seer.db import DbRunMemory, Session
11+
from seer.db import DbPrContextToUnitTestGenerationRunIdMapping
1212

1313
logger = logging.getLogger(__name__)
1414

@@ -80,30 +80,9 @@ def get_file_contents(
8080

8181
return file_contents
8282

83-
def store_memory(self, key: str, memory: list[Message]):
84-
with Session() as session:
85-
memory_record = (
86-
session.query(DbRunMemory).where(DbRunMemory.run_id == self.run_id).one_or_none()
87-
)
88-
89-
if not memory_record:
90-
memory_model = UnitTestRunMemory(run_id=self.run_id)
91-
else:
92-
memory_model = UnitTestRunMemory.from_db_model(memory_record)
93-
94-
memory_model.memory[key] = memory
95-
memory_record = memory_model.to_db_model()
96-
97-
session.merge(memory_record)
98-
session.commit()
99-
100-
def get_memory(self, key: str) -> list[Message]:
101-
with Session() as session:
102-
memory_record = (
103-
session.query(DbRunMemory).where(DbRunMemory.run_id == self.run_id).one_or_none()
104-
)
105-
106-
if not memory_record:
107-
return []
108-
109-
return UnitTestRunMemory.from_db_model(memory_record).memory.get(key, [])
83+
def get_unit_test_memory(self, owner: str, repo: str, pr_id: int) -> list[Message]:
84+
return DbPrContextToUnitTestGenerationRunIdMapping.objects.filter(
85+
owner=self.request.owner,
86+
repo=self.request.repo_definition.name,
87+
pr_id=self.request.pr_id,
88+
).first()

src/seer/automation/codegen/models.py

+2-13
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import datetime
22
from enum import Enum
3-
from typing import Literal
3+
from typing import Literal, Optional
44

55
from pydantic import BaseModel, Field
66

@@ -47,6 +47,7 @@ class CodeUnitTestOutput(BaseComponentOutput):
4747
class CodegenBaseRequest(BaseModel):
4848
repo: RepoDefinition
4949
pr_id: int # The PR number
50+
codecov_status: Optional[dict] = None
5051

5152

5253
class CodegenUnitTestsRequest(CodegenBaseRequest):
@@ -211,15 +212,3 @@ class CodecovTaskRequest(BaseModel):
211212
data: CodegenUnitTestsRequest | CodegenPrReviewRequest | CodegenRelevantWarningsRequest
212213
external_owner_id: str
213214
request_type: Literal["unit-tests", "pr-review", "relevant-warnings", "retry-unit-tests"]
214-
215-
216-
class UnitTestRunMemory(BaseModel):
217-
run_id: int
218-
memory: dict[str, list[Message]] = Field(default_factory=dict)
219-
220-
def to_db_model(self) -> DbRunMemory:
221-
return DbRunMemory(run_id=self.run_id, value=self.model_dump(mode="json"))
222-
223-
@classmethod
224-
def from_db_model(cls, model: DbRunMemory) -> "UnitTestRunMemory":
225-
return cls.model_validate(model.value)

src/seer/automation/codegen/pr_review_step.py

+1
Original file line numberDiff line numberDiff line change
@@ -86,4 +86,5 @@ def _invoke(self, **kwargs):
8686
except ValueError as e:
8787
self.logger.error(f"Error publishing pr review for {pr.url}: {e}")
8888
return
89+
8990
self.context.event_manager.mark_completed()

src/seer/automation/codegen/retry_unittest_coding_component.py

+1-88
Original file line numberDiff line numberDiff line change
@@ -34,91 +34,4 @@ class RetryUnitTestCodingComponent(BaseComponent[CodeUnitTestRequest, CodeUnitTe
3434
def invoke(
3535
self, request: CodeUnitTestRequest, generated_run_id: int, llm_client: LlmClient = injected
3636
) -> CodeUnitTestOutput | None:
37-
with BaseTools(self.context, repo_client_type=RepoClientType.CODECOV_UNIT_TEST) as tools:
38-
agent = LlmAgent(
39-
tools=tools.get_tools(),
40-
config=AgentConfig(interactive=False),
41-
)
42-
43-
past_memory = self.context.get_memory("unit-tests", generated_run_id)
44-
print(past_memory)
45-
46-
# codecov_client_params = request.codecov_client_params
47-
48-
# code_coverage_data = CodecovClient.fetch_coverage(
49-
# repo_name=codecov_client_params["repo_name"],
50-
# pullid=codecov_client_params["pullid"],
51-
# owner_username=codecov_client_params["owner_username"],
52-
# )
53-
54-
# test_result_data = CodecovClient.fetch_test_results_for_commit(
55-
# repo_name=codecov_client_params["repo_name"],
56-
# owner_username=codecov_client_params["owner_username"],
57-
# latest_commit_sha=codecov_client_params["head_sha"],
58-
# )
59-
60-
# # GIVEN UNIT TEST INFORMATION, REGENERATE UNIT TESTS
61-
62-
# existing_test_design_response = llm_client.generate_text(
63-
# model=AnthropicProvider.model("claude-3-5-sonnet@20240620"),
64-
# prompt=CodingUnitTestPrompts.format_find_unit_test_pattern_step_msg(
65-
# diff_str=request.diff
66-
# ),
67-
# )
68-
69-
# formatted_plan_response = llm_client.generate_text(
70-
# model=AnthropicProvider.model("claude-3-5-sonnet@20240620"),
71-
# prompt=CodingUnitTestPrompts.format_plan_step_msg(
72-
# diff_str=request.diff,
73-
# has_coverage_info=code_coverage_data,
74-
# has_test_result_info=test_result_data,
75-
# ),
76-
# )
77-
78-
# final_response = agent.run(
79-
# run_config=RunConfig(
80-
# prompt=CodingUnitTestPrompts.format_unit_test_msg(
81-
# diff_str=request.diff,
82-
# test_design_hint=f"{existing_test_design_response}\n\n{formatted_plan_response}",
83-
# ),
84-
# system_prompt=CodingUnitTestPrompts.format_system_msg(),
85-
# model=AnthropicProvider.model("claude-3-5-sonnet@20240620"),
86-
# run_name="Retry Generate Unit Tests",
87-
# ),
88-
# )
89-
90-
# if not final_response:
91-
# return None
92-
# plan_steps_content = extract_text_inside_tags(final_response, "plan_steps")
93-
94-
# if len(plan_steps_content) == 0:
95-
# raise ValueError("Failed to extract plan_steps from the planning step of LLM")
96-
97-
# coding_output = PlanStepsPromptXml.from_xml(
98-
# f"<plan_steps>{escape_multi_xml(plan_steps_content, ['diff', 'description', 'commit_message'])}</plan_steps>"
99-
# ).to_model()
100-
101-
# if not coding_output.tasks:
102-
# raise ValueError("No tasks found in coding output")
103-
# file_changes: list[FileChange] = []
104-
# for task in coding_output.tasks:
105-
# repo_client = self.context.get_repo_client(
106-
# task.repo_name, type=RepoClientType.CODECOV_UNIT_TEST
107-
# )
108-
# if task.type == "file_change":
109-
# file_content, _ = repo_client.get_file_content(task.file_path)
110-
# if not file_content:
111-
# logger.warning(f"Failed to get content for {task.file_path}")
112-
# continue
113-
114-
# changes, _ = task_to_file_change(task, file_content)
115-
# file_changes += changes
116-
# elif task.type == "file_delete":
117-
# change = task_to_file_delete(task)
118-
# file_changes.append(change)
119-
# elif task.type == "file_create":
120-
# change = task_to_file_create(task)
121-
# file_changes.append(change)
122-
# else:
123-
# logger.warning(f"Unsupported task type: {task.type}")
124-
return CodeUnitTestOutput(diffs=[]) #
37+
pass

src/seer/automation/codegen/retry_unittest_step.py

+37-37
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@
99
AUTOFIX_EXECUTION_SOFT_TIME_LIMIT_SECS,
1010
)
1111

12-
# from seer.automation.codegen.retry_unit_test_coding_component import RetryUnitTestCodingComponent
1312
from seer.automation.codebase.repo_client import RepoClientType
1413
from seer.automation.codegen.models import CodeUnitTestRequest
1514
from seer.automation.codegen.retry_unittest_coding_component import RetryUnitTestCodingComponent
@@ -19,12 +18,13 @@
1918
from seer.automation.pipeline import PipelineStepTaskRequest
2019
from seer.automation.state import DbStateRunTypes
2120
from seer.automation.utils import determine_mapped_unit_test_run_id
21+
from seer.db import DbPrContextToUnitTestGenerationRunIdMapping
2222

2323

2424
class RetryUnittestStepRequest(PipelineStepTaskRequest):
2525
pr_id: int
2626
repo_definition: RepoDefinition
27-
# codecov_status: dict
27+
codecov_status: dict
2828

2929

3030
@celery_app.task(
@@ -50,50 +50,50 @@ def _instantiate_request(request: dict[str, Any]) -> RetryUnittestStepRequest:
5050

5151
@staticmethod
5252
def get_task():
53-
x = retry_unittest_task
54-
return x
53+
return retry_unittest_task
5554

5655
@observe(name="Codegen - Retry Unittest Step")
5756
@ai_track(description="Codegen - Retry Unittest Step")
5857
def _invoke(self, **kwargs):
5958
self.logger.info("Executing Codegen - Retry Unittest Step")
6059
self.context.event_manager.mark_running()
61-
# TODO: IF STATUS CHECK HAS PASSED OR WE HAVE MORE THAN 3 COMMITS, SKIP UNIT TEST GENERATION:
62-
63-
repo_client = self.context.get_repo_client(type=RepoClientType.CODECOV_UNIT_TEST)
60+
repo_client = self.context.get_repo_client(
61+
type=RepoClientType.CODECOV_PR_REVIEW
62+
) # Codecov-ai GH app
6463
pr = repo_client.repo.get_pull(self.request.pr_id)
65-
diff_content = repo_client.get_pr_diff_content(pr.url)
64+
codecov_status = self.request.codecov_status["check_run"]["conclusion"]
6665

67-
latest_commit_sha = repo_client.get_pr_head_sha(pr.url)
66+
if codecov_status == "success":
67+
saved_memory = DbPrContextToUnitTestGenerationRunIdMapping.objects.filter(
68+
owner=self.request.owner,
69+
repo=self.request.repo_definition.name,
70+
pr_id=self.request.pr_id,
71+
).first()
6872

69-
codecov_client_params = {
70-
"repo_name": self.request.repo_definition.name,
71-
"pullid": self.request.pr_id,
72-
"owner_username": self.request.repo_definition.owner,
73-
"head_sha": latest_commit_sha,
74-
}
75-
try:
76-
unittest_output = RetryUnitTestCodingComponent(self.context).invoke(
77-
CodeUnitTestRequest(
78-
diff=diff_content,
79-
codecov_client_params=codecov_client_params,
80-
),
81-
generated_run_id=determine_mapped_unit_test_run_id(
82-
owner=self.request.repo_definition.owner,
83-
repo_name=self.request.repo_definition.name,
84-
pr_id=self.request.pr_id,
85-
),
73+
repo_client.post_unit_test_reference_to_original_pr(
74+
saved_memory.original_pr_url, pr.html_url
8675
)
87-
88-
if unittest_output:
89-
for file_change in unittest_output.diffs:
90-
self.context.event_manager.append_file_change(file_change)
91-
generator = GeneratedTestsPullRequestCreator(unittest_output.diffs, pr, repo_client)
92-
generator.create_github_pull_request()
93-
else:
94-
repo_client.post_unit_test_not_generated_message_to_original_pr(pr.html_url)
76+
self.context.event_manager.mark_completed()
77+
else:
78+
past_run = DbPrContextToUnitTestGenerationRunIdMapping.objects.filter(
79+
owner=self.request.owner,
80+
repo=self.request.repo_definition.name,
81+
pr_id=self.request.pr_id,
82+
).first()
83+
if not past_run:
9584
return
85+
if past_run.iterations == 3:
86+
# TODO: Fetch the "best" run and update the PR
87+
return
88+
else:
89+
# TODO: Retry test generation
90+
pass
91+
self.context.event_manager.mark_completed()
9692

97-
except ValueError:
98-
repo_client.post_unit_test_not_generated_message_to_original_pr(pr.html_url)
99-
self.context.event_manager.mark_completed()
93+
def get_mapping(owner, repo, pr_id):
94+
try:
95+
return DbPrContextToUnitTestGenerationRunIdMapping.objects.get(
96+
owner=owner, repo=repo, pr_id=pr_id
97+
)
98+
except DbPrContextToUnitTestGenerationRunIdMapping.DoesNotExist:
99+
return None

src/seer/automation/codegen/tasks.py

+1
Original file line numberDiff line numberDiff line change
@@ -158,6 +158,7 @@ def codegen_retry_unittest(request: CodegenBaseRequest, app_config: AppConfig =
158158
run_id=cur_state.run_id,
159159
pr_id=request.pr_id,
160160
repo_definition=request.repo,
161+
codecov_status=request.codecov_status,
161162
)
162163
RetryUnittestStep.get_signature(
163164
retry_unittest_request, queue=app_config.CELERY_WORKER_QUEUE

0 commit comments

Comments
 (0)