extractor.py

"""
Input: one source file's code.
Output: a Pydantic object listing its functions and their call relationships.
That's it. Pydantic itself filters out calls to Python builtins and common
stdlib modules (os, sys, json, etc.) — those are dropped automatically,
you never see them in the output.

Requirements: pydantic, llama-cpp-python, huggingface_hub
"""

import builtins
import json
from pathlib import Path
from typing import Literal, Optional

from pydantic import BaseModel, Field, field_validator

# Anything called on these is filtered out as noise, not real repo code.
STDLIB_MODULES = {
    "os", "sys", "json", "re", "math", "time", "datetime", "logging",
    "pathlib", "collections", "itertools", "functools", "typing",
    "asyncio", "threading", "subprocess", "random", "uuid", "hashlib",
    "traceback", "warnings", "copy", "io", "csv", "shutil",
}
BUILTIN_NAMES = {n for n in dir(builtins) if not n.startswith("_")}


class ExtractedFunction(BaseModel):
    name: str
    qualname: str = Field(description="Dotted path, e.g. 'ClassName.method_name'. Just 'name' if top-level.")
    kind: Literal["function", "method", "class"]
    start_line: int
    end_line: int
    docstring: Optional[str] = None


class ExtractedCall(BaseModel):
    caller_qualname: str = Field(description="qualname of the function/method that CONTAINS this call.")
    callee_name: str = Field(description="The real function name this call points to, after resolving any import alias.")

    @field_validator("callee_name")
    @classmethod
    def drop_builtins_and_stdlib(cls, v: str) -> str:
        base = v.split(".")[0]
        if v in BUILTIN_NAMES or base in STDLIB_MODULES:
            raise ValueError(f"'{v}' is a builtin/stdlib call, not repo code — should not be reported")
        return v


class FileExtractionResult(BaseModel):
    file: str
    functions: list[ExtractedFunction]
    calls: list[ExtractedCall]


PROMPT_TEMPLATE = """List every function/method/class defined in this file, and every \
call from one to another. Only report calls to functions DEFINED IN THIS REPO \
— skip calls to Python builtins (print, len, open, ...) and standard library \
modules (os, sys, json, ...) entirely, don't list them. \
If a call uses an import alias (e.g. "from x import y as z", then calling \
"z(...)"), report the REAL name ("y"), not the alias.

Respond with ONLY the JSON described by the schema. No prose.

FILE: {file_path}
SOURCE:
```python
{source}
```
"""


def build_prompt(file_path: str, source: str) -> str:
    return PROMPT_TEMPLATE.format(file_path=file_path, source=source)


def load_model():
    from huggingface_hub import hf_hub_download
    from llama_cpp import Llama
    model_path = hf_hub_download(
        repo_id="bartowski/Qwen2.5-Coder-7B-Instruct-GGUF",
filename="Qwen2.5-Coder-7B-Instruct-Q4_K_M.gguf",
    )
    return Llama(model_path=model_path, n_ctx=8192, n_threads=4, verbose=False)


def extract(llm, file_path: str, source: str) -> FileExtractionResult:
    from llama_cpp import LlamaGrammar
    grammar = LlamaGrammar.from_json_schema(json.dumps(FileExtractionResult.model_json_schema()))
    response = llm.create_chat_completion(
        messages=[{"role": "user", "content": build_prompt(file_path, source)}],
        grammar=grammar,
        temperature=0,
    )
    return FileExtractionResult.model_validate_json(response["choices"][0]["message"]["content"])


def main():
    import sys
    file_path = sys.argv[1]
    source = Path(file_path).read_text(encoding="utf-8")
    llm = load_model()
    result = extract(llm, file_path, source)
    print(result.model_dump_json(indent=2))


if __name__ == "__main__":
    main()