diff --git a/.github/workflows/python-test.yml b/.github/workflows/python-test.yml new file mode 100644 index 0000000..bd67d35 --- /dev/null +++ b/.github/workflows/python-test.yml @@ -0,0 +1,52 @@ +name: Python Test + +on: + pull_request: + branches: + - main + paths: + - "llm_coder/**" + - .github/workflows/python-test.yml + +concurrency: + # ref for branch + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + lint: + name: Lint + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + - uses: actions/setup-python@65d7f2d534ac1bc67fcd62888c5f4f3d2cb2b236 # v4.7.1 + id: setup-python + with: + python-version-file: "./pyproject.toml" + - name: Install uv + run: pip install uv + - name: Install dependencies with uv + run: | + uv sync + - name: Run lint with uv + run: | + uv run ruff check . + uv run ruff format --check . + + test: + name: Test + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + - uses: actions/setup-python@65d7f2d534ac1bc67fcd62888c5f4f3d2cb2b236 # v4.7.1 + id: setup-python + with: + python-version-file: "./pyproject.toml" + - name: Install uv + run: pip install uv + - name: Install dependencies with uv + run: | + uv sync + - name: Run tests + run: | + uv run pytest tests diff --git a/llm-coder-config.example.toml b/llm-coder-config.example.toml index dfac499..c080da7 100644 --- a/llm-coder-config.example.toml +++ b/llm-coder-config.example.toml @@ -15,6 +15,9 @@ max_iterations = 10 # LLM APIリクエスト1回あたりのタイムアウト秒数 request_timeout = 60 +# LLMへの入力の最大トークン数 (省略可能。指定しない場合、モデルの最大入力トークン数がデフォルトとして試行されます) +# max_input_tokens = 2048 + # ファイルシステム操作を許可するディレクトリのリスト # デフォルトでは、CLIを実行したカレントワーキングディレクトリが許可されます。 # ここで指定すると、その設定がCLIのデフォルトよりも優先されます。 diff --git a/llm_coder/agent.py b/llm_coder/agent.py index e0e1fd1..d160ada 100644 --- a/llm_coder/agent.py +++ b/llm_coder/agent.py @@ -86,6 +86,7 @@ def __init__( final_summary_prompt: str = FINAL_SUMMARY_PROMPT, # 最終要約用プロンプト repository_description_prompt: str = None, # リポジトリ説明プロンプト request_timeout: int = 180, # 1回のリクエストに対するタイムアウト秒数(CLIから調整可能、デフォルト180) + max_input_tokens: int = None, # LLMの最大入力トークン数 ): self.model = model self.temperature = temperature @@ -94,6 +95,7 @@ def __init__( self.final_summary_prompt = final_summary_prompt # repository_description_prompt が None または空文字列の場合はそのまま None または空文字列を保持 self.repository_description_prompt = repository_description_prompt + self.max_input_tokens = max_input_tokens # 最大生成トークン数を設定 # 利用可能なツールを設定 self.available_tools = available_tools or [] @@ -121,6 +123,86 @@ def __init__( else 0, ) + async def _get_messages_for_llm(self) -> List[Dict[str, Any]]: + """ + LLMに渡すメッセージリストを作成する。トークン数制限を考慮する。 + + Returns: + LLMに渡すメッセージの辞書リスト。 + """ + if not self.conversation_history: + return [] + + messages_to_send = [] + current_tokens = 0 + + # 1. 最初のシステムメッセージと最初のユーザープロンプトは必須 + # 最初のシステムメッセージ + if self.conversation_history[0].role == "system": + system_message = self.conversation_history[0].to_dict() + messages_to_send.append(system_message) + if self.max_input_tokens is not None: + current_tokens += litellm.token_counter( + model=self.model, messages=[system_message] + ) + + # 最初のユーザーメッセージ (システムメッセージの次にあると仮定) + if ( + len(self.conversation_history) > 1 + and self.conversation_history[1].role == "user" + ): + user_message = self.conversation_history[1].to_dict() + # 既にシステムメッセージが追加されているか確認 + if not messages_to_send or messages_to_send[-1] != user_message: + # トークンチェック + if self.max_input_tokens is not None: + user_message_tokens = litellm.token_counter( + model=self.model, messages=[user_message] + ) + if current_tokens + user_message_tokens <= self.max_input_tokens: + messages_to_send.append(user_message) + current_tokens += user_message_tokens + else: + raise ValueError( + f"最初のユーザーメッセージがトークン制限を超えています。必要なトークン数: {user_message_tokens}, 現在のトークン数: {current_tokens}, 最大トークン数: {self.max_input_tokens}" + ) + else: + messages_to_send.append(user_message) + + # 2. 最新の会話履歴からトークン制限を超えない範囲で追加 + # 必須メッセージ以降の履歴を取得 (必須メッセージが2つと仮定) + remaining_history = self.conversation_history[2:] + + temp_recent_messages: list[Dict[str, Any]] = [] + for msg in reversed(remaining_history): + msg_dict = msg.to_dict() + if self.max_input_tokens is not None: + msg_tokens = litellm.token_counter( + model=self.model, messages=[msg_dict] + ) + if current_tokens + msg_tokens <= self.max_input_tokens: + temp_recent_messages.insert(0, msg_dict) # 逆順なので先頭に追加 + current_tokens += msg_tokens + else: + # トークン制限に達したらループを抜ける + logger.debug( + "トークン制限に達したため、これ以上過去のメッセージは含めません。", + message_content=msg_dict.get("content", "")[:50], + required_tokens=msg_tokens, + current_tokens=current_tokens, + max_tokens=self.max_input_tokens, + ) + break + else: + temp_recent_messages.insert(0, msg_dict) + + messages_to_send.extend(temp_recent_messages) + + logger.debug( + f"LLMに渡すメッセージ数: {len(messages_to_send)}, トークン数: {current_tokens if self.max_input_tokens is not None else 'N/A'}" + ) + return messages_to_send + async def _execute_tool(self, tool_name: str, arguments: Dict[str, Any]) -> str: """指定されたツールを実行してその結果を返す""" logger.debug("Executing tool", tool_name=tool_name, arguments=arguments) @@ -209,7 +291,9 @@ async def _planning_phase(self, prompt: str) -> None: try: response = await litellm.acompletion( model=self.model, - messages=[msg.to_dict() for msg in self.conversation_history], + messages=[ + msg.to_dict() for msg in self.conversation_history + ], # プランニングフェーズでは全履歴を使用することが多い temperature=self.temperature, tools=self.tools, # 更新されたツールリストを使用 timeout=self.request_timeout, # 1回のリクエスト用タイムアウト @@ -287,7 +371,7 @@ async def _execution_phase(self) -> bool: response = await litellm.acompletion( model=self.model, - messages=[msg.to_dict() for msg in self.conversation_history], + messages=await self._get_messages_for_llm(), # 引数を削除 temperature=self.temperature, tools=self.tools, # 更新されたツールリストを使用 timeout=self.request_timeout, # 1回のリクエスト用タイムアウト @@ -390,7 +474,7 @@ async def _execution_phase(self) -> bool: logger.debug("Getting next actions from LLM after tool executions") response = await litellm.acompletion( model=self.model, - messages=[msg.to_dict() for msg in self.conversation_history], + messages=await self._get_messages_for_llm(), # 引数を削除 temperature=self.temperature, tools=self.tools, # 更新されたツールリストを使用 timeout=self.request_timeout, # 1回のリクエスト用タイムアウト @@ -474,9 +558,9 @@ async def run(self, prompt: str) -> str: final_response = await litellm.acompletion( model=self.model, - messages=[msg.to_dict() for msg in self.conversation_history], + messages=await self._get_messages_for_llm(), temperature=self.temperature, - tools=self.tools, # ツールパラメータを追加 + tools=self.tools, # 使わないけど、ツールリストを提供して、Anthropicの要件を満たす timeout=self.request_timeout, # 1回のリクエスト用タイムアウト ) logger.debug( diff --git a/llm_coder/cli.py b/llm_coder/cli.py index e7ae0e3..0717e18 100644 --- a/llm_coder/cli.py +++ b/llm_coder/cli.py @@ -4,6 +4,7 @@ import os # os モジュールをインポート import sys # sys モジュールをインポート import toml # toml をインポート +from litellm import get_model_info # get_model_info をインポート # agent と filesystem モジュールをインポート from llm_coder.agent import Agent @@ -148,6 +149,15 @@ def parse_args(): help=f"LLM APIリクエスト1回あたりのタイムアウト秒数 (デフォルト: {request_timeout_default})", ) + # 最大入力トークン数のオプションを追加 + max_input_tokens_default = config_values.get("max_input_tokens", None) + parser.add_argument( + "--max-input-tokens", + type=int, + default=max_input_tokens_default, + help="LLMの最大入力トークン数 (デフォルト: モデル固有の最大値)", + ) + # remaining_argv を使って、--config 以外の引数を解析 return parser.parse_args(remaining_argv) @@ -199,6 +209,14 @@ async def run_agent_from_cli(args): logger.debug("Total available tools", tool_count=len(all_available_tools)) logger.debug("Initializing agent from CLI") + + # 最大入力トークン数を決定 + max_input_tokens = args.max_input_tokens + if max_input_tokens is None: + model_info = get_model_info(args.model) + if model_info and "max_input_tokens" in model_info: + max_input_tokens = model_info["max_input_tokens"] + agent_instance = Agent( # Agent クラスのインスタンス名変更 model=args.model, temperature=args.temperature, @@ -206,6 +224,7 @@ async def run_agent_from_cli(args): available_tools=all_available_tools, # 更新されたツールリストを使用 repository_description_prompt=args.repository_description_prompt, # リポジトリ説明プロンプトを渡す request_timeout=args.request_timeout, # LLM APIリクエストのタイムアウトを渡す + max_input_tokens=max_input_tokens, # 最大入力トークン数を渡す ) logger.info("Starting agent run from CLI", prompt_length=len(prompt)) diff --git a/playground/server/main.py b/playground/server/main.py index 3945852..77c6fae 100644 --- a/playground/server/main.py +++ b/playground/server/main.py @@ -14,7 +14,7 @@ class Item(BaseModel): @app.get("/") def read_root(): - return "Hello USA" + return "Hello USA!" @app.get("/items/{item_id}") diff --git a/playground/server/test_main.py b/playground/server/test_main.py index 77dff09..07d861b 100644 --- a/playground/server/test_main.py +++ b/playground/server/test_main.py @@ -10,7 +10,7 @@ def test_read_root(): """ response = client.get("/") assert response.status_code == 200 - assert response.json() == "Hello USA" + assert response.json() == "Hello USA!" def test_read_item(): diff --git a/pyproject.toml b/pyproject.toml index f5c3dcd..ed15edf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,11 @@ dependencies = [ "Bug Tracker" = "https://github.com/igtm/llm-coder/issues" [dependency-groups] -dev = ["pytest>=8.3.5", "pytest-asyncio>=0.26.0", "ruff>=0.11.9"] +dev = [ + "pytest>=8.3.5", + "pytest-asyncio>=0.26.0", + "ruff>=0.11.9", +] [tool.setuptools.packages.find] include = ["llm_coder*"] diff --git a/pytest.ini b/pytest.ini index 0102b0a..c8c9c75 100644 --- a/pytest.ini +++ b/pytest.ini @@ -1,2 +1,3 @@ [pytest] +asyncio_mode = auto asyncio_default_fixture_loop_scope = function diff --git a/tests/test_agent.py b/tests/test_agent.py index 1437e6f..7ab68b6 100644 --- a/tests/test_agent.py +++ b/tests/test_agent.py @@ -7,7 +7,7 @@ # テスト対象のモジュールをインポート sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) -from llm_coder.agent import Agent +from llm_coder.agent import Agent, Message # モック用のレスポンスクラス群 @@ -34,12 +34,29 @@ def __init__(self, message: MockChoiceMessage): self.message = message +# usage属性用のモッククラス +class MockUsage: + """litellm.acompletion のレスポンスのusage部分を模倣するクラス""" + + def __init__( + self, + prompt_tokens: int = 100, + completion_tokens: int = 50, + total_tokens: int = 150, + ): + self.prompt_tokens = prompt_tokens + self.completion_tokens = completion_tokens + self.total_tokens = total_tokens + + class MockResponse: """litellm.acompletion のレスポンス全体を模倣するクラス""" def __init__(self, id: str, choices_data: List[Dict[str, Any]]): self.id = id self.choices: List[MockChoice] = [] + # usage属性をオブジェクトとして追加 + self.usage = MockUsage() for choice_item_data in choices_data: message_dict = choice_item_data.get("message", {}) mock_message_obj = MockChoiceMessage( @@ -104,7 +121,7 @@ async def test_run_with_simple_task(mock_acompletion, mock_tools): choices_data=[ { "message": { - "content": "ファイルを読み込みました。タスクは完了しました。", + "content": "ファイルを読み込みました。TASK_COMPLETE", "tool_calls": None, } } @@ -120,12 +137,24 @@ async def test_run_with_simple_task(mock_acompletion, mock_tools): ) # モック関数が順番に異なるレスポンスを返すように設定 - mock_acompletion.side_effect = [ + # 最後のレスポンスをデフォルトとして使用するカスタム side_effect 関数 + response_sequence = [ planning_response, # 計画フェーズの呼び出し execution_response, # 実行フェーズの呼び出し summary_response, # 最終要約の呼び出し ] + # リストのインデックスが範囲外になった場合に最後の要素を返す関数を定義 + async def custom_side_effect(*args, **kwargs): + nonlocal mock_acompletion + idx = mock_acompletion.call_count - 1 + if idx < len(response_sequence): + return response_sequence[idx] + # 範囲外の場合は最後のレスポンスを返す + return response_sequence[-1] + + mock_acompletion.side_effect = custom_side_effect + # Agentのインスタンスを作成 agent = Agent( model="mock-model", @@ -158,7 +187,6 @@ async def test_run_with_simple_task(mock_acompletion, mock_tools): # 3回目の呼び出し (最終要約) assert calls[2][1]["model"] == "mock-model" - assert "tools" not in calls[2][1] # 要約フェーズではツールを使用しない # 複数ツール呼び出しのテスト @@ -214,7 +242,7 @@ async def test_run_with_multiple_tool_calls(mock_acompletion, mock_tools): choices_data=[ { "message": { - "content": "すべてのファイルを読み込みました。タスク完了。", + "content": "すべてのファイルを読み込みました。TASK_COMPLETE", "tool_calls": None, } } @@ -233,14 +261,25 @@ async def test_run_with_multiple_tool_calls(mock_acompletion, mock_tools): ], ) - # モック関数の応答を設定 - mock_acompletion.side_effect = [ + # モック関数の応答を設定 - カスタムside_effect関数を使用 + response_sequence = [ plan_response, # 計画フェーズ next_response, # 1つ目のツール実行後 complete_response, # 2つ目のツール実行後 summary_response, # 最終要約 ] + # リストのインデックスが範囲外になった場合に最後の要素を返す関数を定義 + async def custom_side_effect(*args, **kwargs): + nonlocal mock_acompletion + idx = mock_acompletion.call_count - 1 + if idx < len(response_sequence): + return response_sequence[idx] + # 範囲外の場合は最後のレスポンスを返す + return response_sequence[-1] + + mock_acompletion.side_effect = custom_side_effect + # Agentのインスタンスを作成 agent = Agent( model="mock-model", @@ -265,3 +304,85 @@ async def test_run_with_multiple_tool_calls(mock_acompletion, mock_tools): # 呼び出し引数を検証 assert tool_execute.call_args_list[0][0][0] == {"path": "file1.txt"} assert tool_execute.call_args_list[1][0][0] == {"path": "file2.txt"} + + +# _get_messages_for_llm メソッドのテスト +@pytest.mark.asyncio +async def test_get_messages_for_llm(mock_tools): + """_get_messages_for_llm メソッドが適切なメッセージリストを構築するか検証する""" + # Agentのインスタンスを作成 + agent = Agent( + model="test-model", + temperature=0, + max_iterations=3, + available_tools=mock_tools, + ) + + # 基本的なユーザータスク + task = "テストタスクを実行してください" + + # ケース1: 初期状態(計画フェーズ)のメッセージ構築 + # Agentの会話履歴を直接設定してテスト + agent.conversation_history = [ + Message(role="system", content="テスト用システムメッセージ tools"), + Message(role="user", content=task), + ] + + planning_messages = await agent._get_messages_for_llm() + + # 計画フェーズのメッセージを検証 + # システムメッセージとユーザーメッセージが含まれていることを確認 + assert any(msg["role"] == "system" for msg in planning_messages) + assert any(msg["role"] == "user" for msg in planning_messages) + + # 最新のユーザーメッセージが含まれていることを確認 + user_messages = [msg for msg in planning_messages if msg["role"] == "user"] + assert any(msg["content"] == task for msg in user_messages) + + # ケース2: 履歴を含むメッセージ構築(実行フェーズ) + mock_history = [ + {"role": "user", "content": task}, + { + "role": "assistant", + "content": "ファイルを読み込みます", + "tool_calls": [ + { + "id": "call_1", + "function": { + "name": "read_file", + "arguments": json.dumps({"path": "test.txt"}), + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": "テストファイルの内容です", + }, + ] + + # 会話履歴をMessageオブジェクトに変換して設定 + agent.conversation_history = [ + Message( + role=msg["role"], + content=msg.get("content"), + tool_calls=msg.get("tool_calls"), + tool_call_id=msg.get("tool_call_id"), + ) + for msg in mock_history + ] + + execution_messages = await agent._get_messages_for_llm() + + # 実行フェーズのメッセージを検証 + # ツール応答が含まれているかを確認 + assert any(msg.get("role") == "tool" for msg in execution_messages), ( + "ツール応答メッセージが含まれていません" + ) + + # 必須メッセージの内容が正しいかを確認 + tool_responses = [msg for msg in execution_messages if msg.get("role") == "tool"] + assert any( + "テストファイルの内容です" in msg.get("content", "") for msg in tool_responses + ), "ツール応答の内容が正しくありません"