# Note this file gets uploaded to the OpenAI API as a grader
from ast_grep_py import SgRoot
from pydantic import BaseModel, Field  # type: ignore
from typing import Any, List, Optional
import re

SUPPORTED_LANGUAGES = ['typescript', 'javascript', 'ts', 'js']

class CodeBlock(BaseModel):
    language: str = Field(
        description="Programming language of the code block (e.g., 'python', 'javascript')",
        examples=["python", "javascript", "typescript"]
    )
    path: str = Field(
        description="Target file path where the code should be written",
        examples=["main.py", "src/app.js", "index.html"]
    )
    code: str = Field(
        description="Actual code content extracted from the code block"
    )

class ASTGrepPattern(BaseModel):
    file_path_mask: str = Field(..., description="The file path pattern to match against")
    pattern: str = Field(..., description="The main AST grep pattern to search for")
    additional_greps: Optional[List[str]] = Field(
        default=None,
        description="Additional patterns that must also be present in the matched code"
    )

def extract_code_blocks(llm_output: str) -> List[CodeBlock]:
    # Regular expression to match code blocks with optional language and path
    try:
        pattern = r"```(\w+\s+)?([\w./-]+)?\n([\s\S]*?)\n```"
        matches = list(re.finditer(pattern, llm_output, re.DOTALL))

        print(f"Found {len(matches)} code blocks in the LLM output")

        # Check if any code blocks were found
        if not matches:
            raise Exception("No code blocks found in the LLM response")

        code_blocks: list[CodeBlock] = []
        for match in matches:
            language = match.group(1) or ""
            path = match.group(2) or ""
            code = match.group(3)

            # Clean the path and language
            path = path.strip()
            language = language.strip()

            # If path is relative (doesn't start with /), prefix with /home/user/testbed/
            if path and not path.startswith("/"):
                original_path = path
                path = f"/home/user/testbed/{path}"
                print(
                    f"Converting relative path '{original_path}' to absolute path '{path}'"
                )

            code_blocks.append(
                CodeBlock(language=language, path=path, code=code.strip())
            )

        # Check for missing language or path in code blocks
        missing_language = [
            i for i, block in enumerate(code_blocks) if not block.language
        ]
        missing_path = [i for i, block in enumerate(code_blocks) if not block.path]

        if missing_language:
            print(
                f"WARNING: Code blocks at positions {missing_language} are missing language identifiers"
            )
            raise Exception(
                f"Code blocks at positions {missing_language} are missing language identifiers"
            )

        if missing_path:
            print(
                f"WARNING: Code blocks at positions {missing_path} are missing file paths"
            )
            raise Exception(
                f"Code blocks at positions {missing_path} are missing file paths"
            )

        paths = [block.path for block in code_blocks if block.path]
        print(
            f"Successfully extracted {len(code_blocks)} code blocks with paths: {', '.join(paths)}"
        )

    except Exception as e:
        print(f"Error extracting code blocks: {str(e)}")
        raise

    return code_blocks


def calculate_ast_grep_score(code_blocks: List[CodeBlock], ast_greps: Any) -> float:
    # Convert ast_greps to list if it's a dict
    if isinstance(ast_greps, dict):
        ast_greps = [ast_greps]

    # Parse each grep pattern into the Pydantic model
    parsed_patterns: List[ASTGrepPattern] = []
    for grep in ast_greps:
        try:
            pattern = ASTGrepPattern(**grep)
            parsed_patterns.append(pattern)
        except Exception as e:
            print(f"Error parsing AST grep pattern: {e}")
            return 0.0

    if not parsed_patterns:
        return 0.0

    total_score = 0.0
    pattern_count = len(parsed_patterns)

    # Filter code blocks to only include TypeScript and JavaScript files
    supported_blocks = [
        block for block in code_blocks
        if block.language.lower() in SUPPORTED_LANGUAGES
    ]

    if not supported_blocks:
        print("No TypeScript or JavaScript code blocks found to analyze")
        return 0.0

    for pattern in parsed_patterns:
        # Find matching code blocks based on path prefix
        matching_blocks = [
            block for block in supported_blocks
            if block.path.startswith(pattern.file_path_mask)
        ]

        if not matching_blocks:
            print(f"No matching code blocks found for path prefix: {pattern.file_path_mask}")
            continue

        pattern_found = False
        for block in matching_blocks:
            try:
                # Create AST root for the code block
                root = SgRoot(block.code, block.language)
                node = root.root()

                # Check main pattern
                matches = node.find(pattern=pattern.pattern)
                if not matches:
                    continue

                # If we have additional greps, check them too
                if pattern.additional_greps:
                    all_additional_found = True
                    for additional_grep in pattern.additional_greps:
                        if additional_grep not in block.code:
                            all_additional_found = False
                            break

                    if not all_additional_found:
                        continue

                # If we get here, we found a match with all required patterns
                pattern_found = True
                break

            except Exception as e:
                print(f"Error processing code block {block.path}: {e}")
                continue

        if pattern_found:
            total_score += 1.0

    # Return average score across all patterns
    return total_score / pattern_count if pattern_count > 0 else 0.0

def grade_format(output_text: str) -> float:
        # Find <plan> and </plan> tags
    plan_start = output_text.find('<plan>')
    plan_end = output_text.find('</plan>')

    # Find <code> and </code> tags
    code_start = output_text.find('<code>')
    code_end = output_text.find('</code>')

    reward = 0.0

    if plan_start == -1 or plan_end == -1 or code_start == -1 or code_end == -1:
        print(f'missing plan or code tags. format reward: {reward}')
        return reward
    reward += 0.1 # total: 0.1

    if not (plan_start < plan_end < code_start < code_end):
        print(f'tags present but not in the correct order. format reward: {reward}')
        return reward
    reward += 0.1 # total: 0.2

    # Check if there are any stray tags
    plan_tags = re.findall(r'</?plan>', output_text)
    code_tags = re.findall(r'</?code>', output_text)

    if len(plan_tags) != 2 or len(code_tags) != 2:
        print(f'found stray plan or code tags. format reward: {reward}')
        return reward
    reward += 0.2 # total: 0.4

    # Extract content after </code> tag
    after_tags = output_text[code_end + len('</code>'):].strip()
    if after_tags:
        print(f'found text after code tags. format reward: {reward}')
        return reward
    reward += 0.2 # total: 0.6

    # Extract content inside <plan> tags
    plan_content = output_text[plan_start + len('<plan>'):plan_end].strip()
    if not plan_content:
        print(f'no plan content found. format reward: {reward}')
        return reward
    reward += 0.1 # total: 0.7

    # Extract content inside <code> tags
    code_content = output_text[code_start + len('<code>'):code_end].strip()
    if not code_content:
        print(f'no code content found. format reward: {reward}')
        return reward
    reward += 0.1 # total: 0.8

    # Extract content between </plan> and <code> tags
    between_tags = output_text[plan_end + len('</plan>'):code_start].strip()
    if between_tags:
        print(f'found text between plan and code tags. format reward: {reward}')
        return reward
    reward += 0.2 # total: 1.0

    if reward == 1.0:
        print(f'global format reward: {reward}')

    return reward

def grade(sample: Any, item: Any) -> float:
    try:
        output_text = sample["output_text"]

        format_reward = grade_format(output_text)
        if format_reward < 1.0:
            return format_reward

        # Extract code content for grading
        code_start = output_text.find('<code>')
        code_end = output_text.find('</code>')
        code_to_grade: str = output_text[code_start + len('<code>'):code_end].strip()
        code_blocks: List[CodeBlock] = []
        try:
            code_blocks = extract_code_blocks(code_to_grade)
        except Exception as e:
            print(f'error extracting code blocks: {e}')
            return 0.5

        ast_greps = item["reference_answer"]["ast_greps"]
        ast_grep_score = calculate_ast_grep_score(code_blocks, ast_greps)

        return (format_reward + ast_grep_score) / 2.0
    except Exception as e:
        print(f"Error during grading: {str(e)}")
        return 0.0
````

Results

> Looking at the total reward (format and AST Grep) together, Runloop has seen
> improvements of on average **12%** of the RFT model compared to the base
> o3-mini model on the benchmark.
>
> They implement two types of tests, one providing explicit content from the
> integration guides (assessing reasoning and instruction following) and one
> without (assessing knowledge recall). Both variants saw improvement of over
> **8%**.
>
> âOpenAIs RFT platform gives us access to the best generalized reasoning models
> in the world, with the toolset to supercharge that reasoning on problem
> domains important to our business.â
>
> âRunloop

#### Correct handling of conflicts and dupes in a schedule manager

Use case

> **Company**: Milo helps busy parents manage chaotic family schedules by
> converting messy inputsâlike text convos with to-dos, school newsletter PDFs,
> weekly reminders, sports schedule emailsâinto reliable calendar and list
> actions.
>
> **Problem to solve**: Base GPT-4o prompting and SFT fell short of trust
> thresholds.
>
> **Objective**: Milo used RFT to properly create coding tasks like event vs.
> list classification, recurrence rule generation, accurate updates and deletes,
> conflict detection, and strict output formatting. They defined a grader that
> checked whether generated item objects were complete, categorized correctly,
> and were a duplicate or had a calendar conflict.

Results

> Results showed performance improvements across the board, with average
> correctness scores **increasing from 0.86 to 0.91**, while the most
> challenging scenarios improved from **0.46 to 0.71** (where a perfect
> score=1).
>
> "Accuracy isn't just a metricâit's peace of mind for busy parents. These are
> still early days but with such important improvements in base performance,
> we're able to push more aggressively into complex reasoning needs."
>
> "Navigating and supporting family dynamics involves understanding nuanced
> implications of the data. Take conflictsâknowing soccer for Ethan conflicts
> with Ella's recital because Dad has to drive both kids goes deeper than simple
> overlapping times."
>
> âMilo, AI scheduling tool for families

### 2\. Pull facts into a clean format

These tasks typically involve subtle distinctions that demand clear
classification guidelines. Successful framing requires explicit and hierarchical
labeling schemes defined through consensus by domain experts. Without consistent
agreement, grading signals become noisy, weakening RFT effectiveness.

#### Assigning ICD-10 medical codes

Use case

> **Company**: Ambience is an AI platform that eliminates administrative burden
> for clinicians and ensures accurate, compliant documentation across 100+
> specialties, helping physicians focus on patient care while increasing
> documentation quality and reducing compliance risk for health systems.
>
> **Problem to solve**: ICD-10 coding is one of the most intricate
> administrative tasks in medicine. After every patient encounter, clinicians
> must map each diagnosis to one of ~70,000 codesânavigating payor-specific
> rules on specificity, site-of-care, and mutually exclusive pairings. Errors
> can trigger audits and fines that stretch into nine figures.
>
> **Objective**: Using reinforcement fine-tuning on OpenAI frontier models,
> Ambience wanted to train a reasoning system that listens to the visit audio,
> pulls in relevant EHR context, and recommends ICD-10 codes with accuracy
> exceeding expert clinicians.

Results

> Ambience achieved model improvements that can lead human experts.
>
> On a gold-panel test set spanning hundreds of encounters, reinforcement
> fine-tuning moved the model from trailing humans to leading them by **12
> pointsâeliminating roughly one quarter of the coding errors trained physicians
> make**:
>
> - o3-mini (base): 0.39 (-6 pts)
> - Physician baseline: 0.45
> - RFT-tuned o3-mini: 0.57 (+12 pts)
>
> The result is a real-time, point-of-care coding support that can raise
> reimbursement integrity while reducing compliance risk.
>
> âAccurate ICD-10 selection is mission-critical for compliant documentation.
> RFT unlocked a new level of coding precision we hadnât seen from any
> foundation model and set a new bar for automated coding.â
>
> âAmbience Healthcare

#### Extracting excerpts to support legal claims

Use case

> **Company**: Harvey is building AI that legal teams trustâand that trust
> hinges on retrieving precisely the right evidence from a sprawling corpora of
> contracts, statutes, and case law. Legal professionals arenât satisfied with
> models that merely generate plausible-sounding summaries or paraphrased
> answers. They demand verifiable citationsâpassages that can be traced directly
> back to source documents.
>
> **Problem to solve**: Harveyâs clients use its models to triage litigation
> risk, construct legal arguments, and support due diligence for legal
> professionalsâall tasks where a single missed or misquoted sentence can flip
> an outcome. Models must be able to parse long, dense legal documents and
> extract only the portions that matter. In practice, these inputs are often
> messy and inconsistent: some claims are vague, while others hinge on rare
> legal doctrines buried deep in boilerplate.
>
> **Objective**: The taskâs requirements are to interpret nuanced legal claims,
> navigate long-form documents, and select on-point support with verbatim
> excerpts.

Prompt

```text
## Instructions
You will be provided with a question and a text excerpt. Identify any passages in the text that are directly relevant to answering the question.
- If there are no relevant passages, return an empty list.
- Passages must be copied **exactly** from the text. Do not paraphrase or summarize.
## Excerpt
"""{text_excerpt}"""
```

Grader

```python
from rapidfuzz import fuzz