AtCoderStudyBooster 0.4.0__py3-none-any.whl → 0.4.2__py3-none-any.whl

atcdr/ai.py

@@ -0,0 +1,434 @@
+import json
+import random
+import time
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional
+
+import rich_click as click
+from openai import BadRequestError, OpenAI
+from rich.console import Console
+from rich.live import Live
+from rich.markup import escape
+from rich.panel import Panel
+from rich.syntax import Syntax
+from rich.table import Table
+from rich.text import Text
+
+from atcdr.test import (
+    LabeledTestCase,
+    ResultStatus,
+    TestCase,
+    TestRunner,
+)
+from atcdr.util.fileops import add_file_selector
+from atcdr.util.filetype import (
+    COMPILED_LANGUAGES,
+    FILE_EXTENSIONS,
+    INTERPRETED_LANGUAGES,
+    Lang,
+    str2lang,
+)
+from atcdr.util.i18n import _
+from atcdr.util.openai import set_api_key
+from atcdr.util.parse import ProblemHTML
+
+
+def render_result_for_GPT(test: TestRunner) -> tuple[str, bool]:
+    results = list(test)
+    match test.info.summary:
+        case ResultStatus.CE:
+            return f'Compile Error \n {test.info.compiler_message}', False
+        case _:
+            message_for_gpt = ''.join(
+                (
+                    f'\n{r.label} => {r.result.passed.value}, Execution Time : {r.result.executed_time}\n'
+                    f'\nInput :\n{r.testcase.input}\nOutput :\n{r.result.output}\nExpected :\n{r.testcase.output}\n'
+                    if r.result.passed == ResultStatus.WA
+                    else f'\n{r.label} => {r.result.passed.value}\nInput :\n{r.testcase.input}\nOutput :\n{r.result.output}\n'
+                )
+                for r in results
+            )
+            return message_for_gpt, False
+
+
+def display_test_results(console: Console, test: TestRunner) -> None:
+    results = list(test)
+
+    table = Table(title='🧪 Test Results')
+    table.add_column('Test Case', style='cyan', no_wrap=True)
+    table.add_column('Status', justify='center', no_wrap=True)
+    table.add_column('Input', style='dim', max_width=30)
+    table.add_column('Output', style='yellow', max_width=30)
+    table.add_column('Expected', style='green', max_width=30)
+
+    for r in results:
+        if r.result.passed == ResultStatus.AC:
+            status = '[green]✅ AC[/green]'
+        elif r.result.passed == ResultStatus.WA:
+            status = '[red]❌ WA[/red]'
+        elif r.result.passed == ResultStatus.TLE:
+            status = '[yellow]⏰ TLE[/yellow]'
+        elif r.result.passed == ResultStatus.RE:
+            status = '[red]💥 RE[/red]'
+        else:
+            status = f'[red]{r.result.passed.value}[/red]'
+
+        input_preview = escape(
+            r.testcase.input.strip()[:50] + '...'
+            if len(r.testcase.input.strip()) > 50
+            else r.testcase.input.strip()
+        )
+        output_preview = escape(
+            r.result.output.strip()[:50] + '...'
+            if len(r.result.output.strip()) > 50
+            else r.result.output.strip()
+        )
+        expected_preview = escape(
+            r.testcase.output.strip()[:50] + '...'
+            if len(r.testcase.output.strip()) > 50
+            else r.testcase.output.strip()
+        )
+
+        table.add_row(r.label, status, input_preview, output_preview, expected_preview)
+
+    console.print(table)
+
+
+def create_func(labeled_cases: list[LabeledTestCase], model: str):
+    def test_example_case(code: str, language: str) -> str:
+        language_enum: Lang = str2lang(language)
+        source_path = Path(f'{model}{FILE_EXTENSIONS[language_enum]}')
+        source_path.write_text(code, encoding='utf-8')
+        test = TestRunner(str(source_path), labeled_cases)
+        message_for_gpt, _ = render_result_for_GPT(test)
+        return message_for_gpt
+
+    def execute_code(input: Optional[str], code: str, language: str) -> str:
+        language_enum: Lang = str2lang(language)
+        random_name = random.randint(0, 100_000_000)
+        source_path = Path(f'tmp{random_name}{FILE_EXTENSIONS[language_enum]}')
+        source_path.write_text(code, encoding='utf-8')
+        labeled_cases = [LabeledTestCase('case by gpt', TestCase(input or '', ''))]
+        test = TestRunner(str(source_path), labeled_cases)
+        labeled_result = next(test)
+        source_path.unlink(missing_ok=True)
+        return labeled_result.result.output
+
+    return test_example_case, execute_code
+
+
+def solve_problem(path: Path, lang: Lang, model: str) -> None:
+    console = Console()
+    content = path.read_text(encoding='utf-8')
+    html = ProblemHTML(content)
+    md = html.make_problem_markdown('en')
+    labeled_cases = html.load_labeled_testcase()
+
+    test_example_case, execute_code = create_func(labeled_cases, model)
+
+    # Responses API 形式のツール定義（トップレベル）
+    TOOLS = [
+        {
+            'type': 'function',
+            'name': 'test_example_case',
+            'description': 'Run the given source code against example test cases and return a summarized result.',
+            'parameters': {
+                'type': 'object',
+                'properties': {
+                    'code': {'type': 'string'},
+                    'language': {
+                        'type': 'string',
+                        'enum': [
+                            lang.value
+                            for lang in (COMPILED_LANGUAGES + INTERPRETED_LANGUAGES)
+                        ],
+                    },
+                },
+                'required': ['code', 'language'],
+                'additionalProperties': False,
+            },
+            'strict': True,
+        },
+        {
+            'type': 'function',
+            'name': 'execute_code',
+            'description': 'Execute the given source code with a single input and return the actual output.',
+            'parameters': {
+                'type': 'object',
+                'properties': {
+                    'input': {'type': 'string'},
+                    'code': {'type': 'string'},
+                    'language': {
+                        'type': 'string',
+                        'enum': [
+                            lang.value
+                            for lang in (COMPILED_LANGUAGES + INTERPRETED_LANGUAGES)
+                        ],
+                    },
+                },
+                'required': ['input', 'code', 'language'],
+                'additionalProperties': False,
+            },
+            'strict': True,
+        },
+    ]
+
+    client = OpenAI()
+    if set_api_key() is None:
+        console.print('[red]OpenAI API key is not set.[/red]')
+        return
+
+    system_prompt = f"""You are a competitive programming assistant for {lang.value}.
+The user will provide problems in Markdown format.
+Read the problem carefully and output a complete, correct, and efficient solution in {lang.value}.
+Use standard input and output. Do not omit any code.
+Always pay close attention to algorithmic complexity (time and space).
+Choose the most optimal algorithms and data structures so that the solution runs within time limits even for the largest possible inputs.
+
+Use the provided tool test_example_case to run the example test cases from the problem statement.
+If tests do not pass, fix the code and repeat.
+The last tested code will be automatically saved to a local file on the user's computer.
+You do not need to include the final source code in your response.
+Simply confirm to the user that all tests passed, or briefly explain if they did not.
+Once you run test_example_case, the exact code you tested will already be saved locally on the user's machine, so sending it again in the response is unnecessary."""
+
+    # ツール名→ローカル実装のディスパッチ
+    tool_impl: Dict[str, Callable[..., Any]] = {
+        'test_example_case': test_example_case,
+        'execute_code': lambda **p: execute_code(
+            p.get('input', ''),  # ← 空なら空文字に
+            p.get('code', ''),
+            p.get('language', lang.value),
+        ),
+    }
+
+    console.print(f'Solving :{path} Language: {lang.value} / Model: {model}')
+
+    context_msgs = [
+        {'role': 'system', 'content': system_prompt},
+        {'role': 'user', 'content': md},
+    ]
+    turn = 1
+    assistant_text = Text()
+
+    def call_model():
+        try:
+            return client.responses.create(
+                model=model,
+                input=context_msgs,
+                tools=TOOLS,
+                tool_choice='auto',
+                include=['reasoning.encrypted_content'],
+                store=False,
+            )
+        except BadRequestError as e:
+            body = getattr(getattr(e, 'response', None), 'json', lambda: None)()
+            console.print(
+                Panel.fit(f'{e}\n\n{body}', title='API Error', border_style='red')
+            )
+            raise
+
+    while True:
+        start_time = time.time()
+        with Live(
+            Panel(
+                f'[bold blue]🤔 Thinking... (turn {turn})[/bold blue]\n[dim]Elapsed: 0.0s[/dim]',
+                border_style='blue',
+            ),
+            console=console,
+            refresh_per_second=10,
+        ) as live:
+
+            def update_timer():
+                elapsed = time.time() - start_time
+                live.update(
+                    Panel(
+                        f'[bold blue]🤔 Thinking... (turn {turn})[/bold blue]\n[dim]Elapsed: {elapsed:.1f}s[/dim]',
+                        border_style='blue',
+                    )
+                )
+
+            import threading
+
+            resp = None
+            error = None
+
+            def model_call():
+                nonlocal resp, error
+                try:
+                    resp = call_model()
+                except Exception as e:
+                    error = e
+
+            thread = threading.Thread(target=model_call)
+            thread.start()
+
+            while thread.is_alive():
+                update_timer()
+                time.sleep(0.1)
+
+            thread.join()
+
+            if error:
+                raise error
+
+            elapsed = time.time() - start_time
+            live.update(
+                Panel(
+                    f'[bold green]✓ Completed thinking (turn {turn})[/bold green]\n[dim]Time taken: {elapsed:.1f}s[/dim]',
+                    border_style='green',
+                )
+            )
+
+        # Display token usage
+        if resp and hasattr(resp, 'usage') and resp.usage:
+            usage = resp.usage
+            input_tokens = getattr(usage, 'input_tokens', 0)
+            output_tokens = getattr(usage, 'output_tokens', 0)
+            total_tokens = getattr(usage, 'total_tokens', 0)
+
+            # Check for cached tokens
+            cached_tokens = 0
+            if hasattr(usage, 'input_tokens_details'):
+                details = usage.input_tokens_details
+                if hasattr(details, 'cached_tokens'):
+                    cached_tokens = details.cached_tokens
+
+            # Check for reasoning tokens
+            reasoning_tokens = 0
+            if hasattr(usage, 'output_tokens_details'):
+                details = usage.output_tokens_details
+                if hasattr(details, 'reasoning_tokens'):
+                    reasoning_tokens = details.reasoning_tokens
+
+            token_msg = f'[dim]Tokens - Input: {input_tokens:,}'
+            if cached_tokens > 0:
+                token_msg += f' (cached: {cached_tokens:,})'
+            token_msg += f' | Output: {output_tokens:,}'
+            if reasoning_tokens > 0:
+                token_msg += f' (reasoning: {reasoning_tokens:,})'
+            token_msg += f' | Total: {total_tokens:,}[/dim]'
+            console.print(token_msg)
+
+        if resp and getattr(resp, 'output_text', None):
+            assistant_text.append(resp.output_text)
+
+            output_content = str(resp.output_text).strip()
+            if any(
+                keyword in output_content
+                for keyword in [
+                    'def ',
+                    'class ',
+                    'import ',
+                    'from ',
+                    '#include',
+                    'public class',
+                ]
+            ):
+                try:
+                    syntax = Syntax(
+                        output_content, lang, theme='monokai', line_numbers=True
+                    )
+                    console.print(
+                        Panel(
+                            syntax,
+                            title='Assistant Output (Code)',
+                            border_style='green',
+                        )
+                    )
+                except Exception:
+                    console.print(
+                        Panel(
+                            assistant_text,
+                            title='Assistant Output',
+                            border_style='green',
+                        )
+                    )
+            else:
+                console.print(
+                    Panel(
+                        assistant_text, title='Assistant Output', border_style='green'
+                    )
+                )
+
+        if resp and hasattr(resp, 'output'):
+            context_msgs += resp.output
+
+            # function_call を収集
+            calls: List[dict] = []
+            for o in resp.output:
+                if getattr(o, 'type', '') == 'function_call':
+                    try:
+                        args = json.loads(o.arguments or '{}')
+                    except Exception:
+                        args = {}
+                    call_id = getattr(o, 'call_id', None) or getattr(
+                        o, 'id'
+                    )  # ★ ここがポイント
+                    calls.append({'name': o.name, 'call_id': call_id, 'args': args})
+        else:
+            calls = []
+
+        if not calls:
+            console.print(
+                Panel.fit('✅ Done (no more tool calls).', border_style='green')
+            )
+            break
+
+        # ツールを実行し、function_call_output を context に積む
+        for c in calls:
+            args_str = json.dumps(c['args'], ensure_ascii=False) if c['args'] else ''
+            console.print(
+                Panel.fit(
+                    f"Tool: [bold]{c['name']}[/bold]\nargs: {args_str}",
+                    title=f"function_call ({c['call_id']})",
+                    border_style='cyan',
+                )
+            )
+
+            impl = tool_impl.get(c['name'])
+            if not impl:
+                out = f"[ERROR] Unknown tool: {c['name']}"
+            else:
+                try:
+                    with console.status(f"Running {c['name']}...", spinner='dots'):
+                        out = impl(**c['args']) if c['args'] else impl()
+                except TypeError:
+                    out = impl(
+                        **{
+                            k: v
+                            for k, v in c['args'].items()
+                            if k in impl.__code__.co_varnames
+                        }
+                    )
+                except Exception as e:
+                    out = f"[Tool '{c['name']}' error] {e}"
+
+            console.print(
+                Panel(
+                    str(out) or '(no output)',
+                    title=f"{c['name']} result",
+                    border_style='magenta',
+                )
+            )
+
+            context_msgs.append(
+                {
+                    'type': 'function_call_output',
+                    'call_id': c['call_id'],
+                    'output': str(out),
+                }
+            )
+
+        turn += 1
+
+
+@click.command(short_help=_('cmd_generate'), help=_('cmd_generate'))
+@add_file_selector('files', filetypes=[Lang.HTML])
+@click.option('--lang', default='Python', help=_('opt_output_lang'))
+@click.option('--model', default='gpt-5-mini', help=_('opt_model'))
+def ai(files, lang, model):
+    """HTMLファイルからコード生成またはテンプレート出力を行います。"""
+    lang_enum: Lang = str2lang(lang)
+    for path in files:
+        solve_problem(Path(path), lang_enum, model)

atcdr/cli.py

@@ -8,8 +8,8 @@ from rich.table import Table
 from rich.traceback import install
 from rich_click import RichGroup
 
+from atcdr.ai import ai
 from atcdr.download import download
-from atcdr.generate import generate
 from atcdr.login import login
 from atcdr.logout import logout
 from atcdr.markdown import markdown
@@ -76,7 +76,7 @@ def cli():
 cli.add_command(test, aliases=['t'])
 cli.add_command(download, aliases=['d'])
 cli.add_command(open_files, 'open', aliases=['o'])
-cli.add_command(generate, aliases=['g'])
+cli.add_command(ai)
 cli.add_command(markdown, aliases=['md'])
 cli.add_command(submit, aliases=['s'])
 cli.add_command(login)

atcdr/download.py

@@ -11,6 +11,7 @@ from rich import print
 from rich.prompt import Prompt
 
 from atcdr.util.filetype import FILE_EXTENSIONS, Lang
+from atcdr.util.i18n import _, i18n
 from atcdr.util.parse import ProblemHTML
 from atcdr.util.problem import Contest, Problem
 from atcdr.util.session import load_session
@@ -25,32 +26,37 @@ class Downloader:
         retry_attempts = 3
         retry_wait = 1  # 1 second
 
-        for _ in range(retry_attempts):
+        for attempt in range(retry_attempts):
             response = session.get(problem.url)
             if response.status_code == 200:
                 return ProblemHTML(response.text)
             elif response.status_code == 429:
                 print(
-                    f'[bold yellow][Error {response.status_code}][/bold yellow] 再試行します。{problem}'
+                    f'[bold yellow][Error {response.status_code}][/bold yellow] '
+                    + _('retry_problem', problem)
                 )
                 time.sleep(retry_wait)
             elif 300 <= response.status_code < 400:
                 print(
-                    f'[bold yellow][Error {response.status_code}][/bold yellow] リダイレクトが発生しました。{problem}'
+                    f'[bold yellow][Error {response.status_code}][/bold yellow] '
+                    + _('redirect_occurred', problem)
                 )
             elif 400 <= response.status_code < 500:
                 print(
-                    f'[bold red][Error {response.status_code}][/bold red] 問題が見つかりません。{problem}'
+                    f'[bold red][Error {response.status_code}][/bold red] '
+                    + _('problem_not_found', problem)
                 )
                 break
             elif 500 <= response.status_code < 600:
                 print(
-                    f'[bold red][Error {response.status_code}][/bold red] サーバーエラーが発生しました。{problem}'
+                    f'[bold red][Error {response.status_code}][/bold red] '
+                    + _('server_error', problem)
                 )
                 break
             else:
                 print(
-                    f'[bold red][Error {response.status_code}][/bold red] {problem}に対応するHTMLファイルを取得できませんでした。'
+                    f'[bold red][Error {response.status_code}][/bold red] '
+                    + _('html_fetch_failed', problem)
                 )
                 break
         return ProblemHTML('')
@@ -68,7 +74,7 @@ def save_problem(problem: Problem, path: Path, session: requests.Session) -> Non
     problem_content = downloader.get(problem)
 
     if not problem_content:
-        print(f'[bold red][Error][/] {problem}の保存に失敗しました')
+        print('[bold red][Error][/] ' + _('save_failed', problem))
         return
 
     # ディレクトリ作成（pathをそのまま使用）
@@ -80,26 +86,26 @@ def save_problem(problem: Problem, path: Path, session: requests.Session) -> Non
     # HTMLファイル保存
     html_path = path / (title + FILE_EXTENSIONS[Lang.HTML])
     html_path.write_text(problem_content.html, encoding='utf-8')
-    print(f'[bold green][+][/bold green] ファイルを保存しました: {html_path}')
+    print('[bold green][+][/bold green] ' + _('file_saved', html_path))
 
     # Markdownファイル保存
-    md = problem_content.make_problem_markdown('ja')
+    md = problem_content.make_problem_markdown(i18n.language)
     md_path = path / (title + FILE_EXTENSIONS[Lang.MARKDOWN])
     md_path.write_text(md, encoding='utf-8')
-    print(f'[bold green][+][/bold green] ファイルを保存しました: {md_path}')
+    print('[bold green][+][/bold green] ' + _('file_saved', md_path))
 
 
 def interactive_download(session) -> None:
-    CONTEST = '1. コンテストの問題を解きたい'
-    ONE_FILE = '2. 1問だけダウンロードする'
-    END = '3. 終了する'
+    CONTEST = '1. ' + _('solve_contest_problems')
+    ONE_FILE = '2. ' + _('download_one_problem')
+    END = '3. ' + _('exit')
 
     choice = q.select(
-        message='AtCoderの問題のHTMLファイルをダウンロードします',
+        message=_('download_atcoder_html'),
         qmark='',
         pointer='❯❯❯',
         choices=[CONTEST, ONE_FILE, END],
-        instruction='\n 十字キーで移動,[enter]で実行',
+        instruction='\n ' + _('navigate_with_arrows'),
         style=q.Style(
             [
                 ('question', 'fg:#2196F3 bold'),
@@ -112,7 +118,7 @@ def interactive_download(session) -> None:
     ).ask()
 
     if choice == CONTEST:
-        name = Prompt.ask('コンテスト名を入力してください (例: abc012, abs, typical90)')
+        name = Prompt.ask(_('input_contest_name'))
         try:
             contest = Contest(name, session)
             for problem in contest.problems:
@@ -121,18 +127,18 @@ def interactive_download(session) -> None:
             print(f'[red][Error][/red] {e}')
 
     elif choice == ONE_FILE:
-        name = Prompt.ask('コンテスト名を入力してください (例: abc012, abs, typical90)')
+        name = Prompt.ask(_('input_contest_name'))
         try:
             contest = Contest(name, session)
             problem = q.select(
-                message='どの問題をダウンロードしますか?',
+                message=_('which_problem_download'),
                 qmark='',
                 pointer='❯❯❯',
                 choices=[
                     q.Choice(title=f'{p.label:10} | {p.url}', value=p)
                     for p in contest.problems
                 ],
-                instruction='\n 十字キーで移動,[enter]で実行',
+                instruction='\n ' + _('navigate_with_arrows'),
                 style=q.Style(
                     [
                         ('question', 'fg:#2196F3 bold'),
@@ -148,9 +154,9 @@ def interactive_download(session) -> None:
             print(f'[red][Error][/red] {e}')
 
     elif choice == END:
-        print('[bold red]終了します[/]')
+        print('[bold red]' + _('exiting') + '[/]')
     else:
-        print('[bold red]無効な選択です[/]')
+        print('[bold red]' + _('invalid_selection') + '[/]')
 
 
 def plan_download(
@@ -179,7 +185,7 @@ def plan_download(
                 for problem in contest.problems
             ]
         else:
-            raise ValueError('コンテスト名を指定してください')
+            raise ValueError(_('specify_contest_name'))
     elif len(groups) == 2:
         result = []
         for i, j in product(groups[0], groups[1]):
@@ -195,16 +201,15 @@ def plan_download(
                         result.append((problem, Path(i) / j.name))
         return result
     else:
-        raise ValueError('ダウンロードの引数が正しくありません')
+        raise ValueError(_('invalid_download_args'))
 
 
-@click.command(short_help='AtCoder の問題をダウンロード')
+@click.command(short_help=_('cmd_download'), help=_('cmd_download'))
 @click.argument('args', nargs=-1)
 def download(args: List[str]) -> None:
     """
-    例:
-        download abc{001..012} {A..C}
-        download {A..E} abc{001..012}
+    download abc{001..012} {A..C}
+    download {A..E} abc{001..012}
     """
     session = load_session()