feat: mcp tools and resources v0

Author: rushilpatel0

src/codegen/cli/api/client.py

@@ -15,6 +15,7 @@
     PR_LOOKUP_ENDPOINT,
     RUN_ENDPOINT,
     RUN_ON_PR_ENDPOINT,
+    IMPROVE_ENDPOINT,
 )
 from codegen.cli.api.schemas import (
     AskExpertInput,
@@ -36,12 +37,15 @@
     RunCodemodOutput,
     RunOnPRInput,
     RunOnPRResponse,
+    ImproveCodemodInput,
+    ImproveCodemodResponse,
 )
 from codegen.cli.auth.session import CodegenSession
 from codegen.cli.codemod.convert import convert_to_ui
 from codegen.cli.env.global_env import global_env
 from codegen.cli.errors import InvalidTokenError, ServerError
 from codegen.cli.utils.codemods import Codemod
+from codegen.cli.utils.constants import ProgrammingLanguage
 from codegen.cli.utils.function_finder import DecoratedFunction
 
 InputT = TypeVar("InputT", bound=BaseModel)
@@ -248,3 +252,13 @@ def lookup_pr(self, repo_full_name: str, github_pr_number: int) -> PRSchema:
             PRLookupInput(input=PRLookupInput.BasePRLookupInput(repo_full_name=repo_full_name, github_pr_number=github_pr_number)),
             PRLookupResponse,
         )
+
+
+    def improve_codemod(self, codemod: str, task: str, concerns: list[str], context: dict[str, str], language: ProgrammingLanguage) -> ImproveCodemodResponse:
+        """Improve a codemod."""
+        return self._make_request(
+            "GET",
+            IMPROVE_ENDPOINT,
+            ImproveCodemodInput(input=ImproveCodemodInput.BaseImproveCodemodInput(codemod=codemod, task=task, concerns=concerns, context=context, language=language)),
+            ImproveCodemodResponse
+        )

src/codegen/cli/api/endpoints.py

@@ -10,3 +10,4 @@
 RUN_ON_PR_ENDPOINT = f"https://{MODAL_PREFIX}--cli-run-on-pull-request.modal.run"
 PR_LOOKUP_ENDPOINT = f"https://{MODAL_PREFIX}--cli-pr-lookup.modal.run"
 CODEGEN_SYSTEM_PROMPT_URL = "https://gist.githubusercontent.com/jayhack/15681a2ceaccd726f19e6fdb3a44738b/raw/17c08054e3931b3b7fdf424458269c9e607541e8/codegen-system-prompt.txt"
+IMPROVE_ENDPOINT = "https://codegen-sh-develop-rpatel--cli-improve.modal.run"

src/codegen/cli/api/schemas.py

@@ -234,3 +234,22 @@ class RunOnPRResponse(BaseModel):
     codemod_id: int = Field(..., description="ID of the codemod")
     codemod_run_id: int = Field(..., description="ID of the codemod run")
     web_url: str = Field(..., description="URL to view the test results")
+
+
+###########################################################################
+# IMPROVE
+###########################################################################
+
+class ImproveCodemodInput(BaseModel):
+    class BaseImproveCodemodInput(BaseModel):
+        codemod: str = Field(..., description="Source code of the codemod to improve")
+        task: str = Field(..., description="Task to which the codemod should implement to solve")
+        concerns: list[str] = Field(..., description="A list of issues that were discovered with the current codemod that need to be considered in the next iteration")
+        context: dict[str, str] = Field(..., description="Additional context for the codemod this can be a list of files that are related, additional information about the task, etc.")
+        language: ProgrammingLanguage = Field(..., description="Language of the codemod")
+
+    input: BaseImproveCodemodInput = Field(..., description="Input data for improvement")
+
+class ImproveCodemodResponse(BaseModel):
+    success: bool = Field(..., description="Whether the improvement was successful")
+    codemod_source: str = Field(..., description="Source code of the improved codemod")