Add context to tool calling

.github/workflows/publish.yaml

@@ -4,9 +4,9 @@ on:
   push:
     tags:
       - 'v[0-9]+.[0-9]+.[0-9]+'
-      - 'v[0-9]+.[0-9]+.[0-9]+a[0-9]*'
-      - 'v[0-9]+.[0-9]+.[0-9]+b[0-9]*'
-      - 'v[0-9]+.[0-9]+.[0-9]+rc[0-9]*'
+      - 'v[0-9]+.[0-9]+.[0-9]+a[0-9]+'
+      - 'v[0-9]+.[0-9]+.[0-9]+b[0-9]+'
+      - 'v[0-9]+.[0-9]+.[0-9]+rc[0-9]+'
 
 permissions:
   id-token: write

.github/workflows/test.yaml

@@ -5,13 +5,14 @@ on:
 
 jobs:
   check:
-    name: Ruff check
+    name: Ruff & uv check
     runs-on: ubuntu-latest
 
     steps:
       - uses: actions/checkout@v4
       - uses: cashapp/activate-hermit@v1
       - run: |
+          uv lock --check
           uv run ruff check .
           uv run ruff format --check
 
@@ -34,4 +35,4 @@ jobs:
       - uses: actions/checkout@v4
       - uses: cashapp/activate-hermit@v1
       - run: |
-          uv run pytest
+          uv run pytest

projects/mini/tools.py

@@ -1,7 +1,9 @@
 from pydantic_ai import Tool
 
+from ai_migrate.context import ToolCallContext
 
-def add(x: int, y: int) -> int:
+
+def add(ctx: ToolCallContext, x: int, y: int) -> int:
     """Add two numbers"""
     return x + y
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ai-migrate-tools"
-version = "0.1.4a1"
+version = "0.1.4a2"
 description = "LLM-powered code migrations at scale"
 license = "Apache-2.0"
 readme = "README.md"

src/ai_migrate/context.py

@@ -0,0 +1,13 @@
+from dataclasses import dataclass
+from pathlib import Path
+
+from pydantic_ai import RunContext
+
+
+@dataclass
+class MigrationContext:
+    target_files: list[str]
+    target_dir: Path | None
+
+
+ToolCallContext = RunContext[MigrationContext]

src/ai_migrate/migrate.py

@@ -12,9 +12,9 @@
 
 from pydantic_ai.messages import ToolCallPart
 from pydantic_ai.tools import Tool
-from pydantic_ai import RunContext
 
 from ai_migrate.llm_providers import DefaultClient
+from .context import MigrationContext, ToolCallContext
 from .fake_llm_client import FakeLLMClient
 from .git_identity import environment_variables
 from .manifest import FileGroup, FileEntry, Manifest
@@ -31,9 +31,6 @@ def log(*args, **kwargs):
     print(*args, **kwargs, file=LOG_STREAM.get(), flush=True)
 
 
-NoneContext = RunContext[None]
-
-
 @dataclass
 class FileContent:
     name: str
@@ -201,7 +198,7 @@ def migrate_prompt(example: MigrationExample) -> list[dict]:
 
 
 async def handle_tool_calls(
-    tools: list[Tool], tool_calls: list[dict[str, Any]]
+    tools: list[Tool], tool_calls: list[dict[str, Any]], context: MigrationContext
 ) -> list[dict[str, str]]:
     tool_results = []
     tools_by_name = {tool.name: tool for tool in tools}
@@ -215,8 +212,8 @@ async def handle_tool_calls(
                 log("[agent] Running tool", tool.name, f"{tool_call_message=}")
                 result = await tool._run(
                     tool_call_message,
-                    NoneContext(
-                        deps=None,
+                    ToolCallContext(
+                        deps=context,
                         model=None,
                         usage=None,
                         prompt="",
@@ -265,7 +262,12 @@ def combine_examples_into_conversation(
 
 
 async def call_llm(
-    client: DefaultClient, messages: list, tools: list[Tool], temperature=0.1
+    client: DefaultClient,
+    messages: list,
+    tools: list[Tool],
+    *,
+    context: MigrationContext,
+    temperature=0.1,
 ) -> tuple[dict, list[dict]]:
     """Call LLM for completions
 
@@ -283,7 +285,7 @@ async def call_llm(
         if not (tool_calls := assistant_message.get("tool_calls")):
             return response, messages
 
-        tool_results = await handle_tool_calls(tools, tool_calls)
+        tool_results = await handle_tool_calls(tools, tool_calls, context)
 
         messages.append(assistant_message)
         for result in tool_results:
@@ -606,7 +608,13 @@ async def _run(
             iteration_messages = iteration_messages[1:]
             messages = build_messages(messages, iteration_messages)
 
-        response, messages = await call_llm(client, messages, tools or [])
+        context = MigrationContext(
+            target_files=target_files,
+            target_dir=Path(target_dir) if target_dir else None,
+        )
+        response, messages = await call_llm(
+            client, messages, tools or [], context=context
+        )
 
         response_text = response["choices"][0]["message"]["content"]
         parsed_result = extract_code_blocks(response_text)