2222import re
2323import time
2424import unicodedata
25- from collections .abc import Callable , Iterator
25+ from collections .abc import Callable , Iterator , Sequence
2626from dataclasses import dataclass , field
2727from typing import cast
2828
@@ -941,6 +941,7 @@ class _TP4Candidate:
941941 content : str
942942 start_line : int = 1
943943 end_line : int = 1
944+ context : str = ""
944945
945946
946947_TP4_MARKDOWN_TYPES = frozenset ({"markdown" , "text" })
@@ -964,21 +965,55 @@ class _TP4Candidate:
964965}
965966_TP4_FENCE_OPEN_RE = re .compile (r"^[ ]{0,3}(`{3,}|~{3,})[ \t]*([^ \t]+)?[ \t]*$" )
966967_TP4_FENCE_CLOSE_RE = re .compile (r"^[ ]{0,3}(`{3,}|~{3,})[ \t]*$" )
968+ TP4_MAX_CONTEXT_CHARS = 2_048
969+ _TP4_MAX_CONTEXT_LINES = 8
970+
971+
972+ def _tp4_fence_context (preceding : Sequence [str ]) -> str :
973+ """Return the bounded prose that introduces a fence, if any.
974+
975+ A fenced block is often introduced by text that changes how it should be
976+ read, such as a heading marking it as an example that must not be run.
977+ The extractor cannot see that framing when only the fence body is sent, so
978+ a short trailing window of non-blank prose is retained with the code.
979+ """
980+ context_lines : list [str ] = []
981+ for raw in reversed (preceding ):
982+ stripped = raw .strip ()
983+ if not stripped :
984+ if context_lines :
985+ break
986+ continue
987+ context_lines .append (stripped )
988+ if len (context_lines ) >= _TP4_MAX_CONTEXT_LINES :
989+ break
990+ if not context_lines :
991+ return ""
992+ context = "\n " .join (reversed (context_lines ))
993+ if len (context ) > TP4_MAX_CONTEXT_CHARS :
994+ context = context [- TP4_MAX_CONTEXT_CHARS :]
995+ return context
967996
968997
969998def _iter_tp4_markdown_fences (
970999 content : str ,
971- ) -> Iterator [tuple [str , str , int , int ]]:
972- """Yield exactly labeled, non-empty executable fences from bounded Markdown/text."""
1000+ ) -> Iterator [tuple [str , str , int , int , str ]]:
1001+ """Yield labeled, non-empty executable fences and their leading context.
1002+
1003+ Yields ``(language, body, start_line, end_line, context)`` where ``context``
1004+ is the bounded prose immediately preceding the opening fence.
1005+ """
9731006 lines = content .splitlines (keepends = True )
9741007 active : tuple [str , int , str ] | None = None
9751008 body : list [str ] = []
9761009 body_start = 0
1010+ preceding : list [str ] = []
9771011 for line_number , line in enumerate (lines , start = 1 ):
9781012 stripped = line .rstrip ("\r \n " )
9791013 if active is None :
9801014 opening = _TP4_FENCE_OPEN_RE .fullmatch (stripped )
9811015 if opening is None :
1016+ preceding .append (line )
9821017 continue
9831018 delimiter , label = opening .groups ()
9841019 active = (delimiter [0 ], len (delimiter ), label .casefold () if label else "" )
@@ -993,9 +1028,16 @@ def _iter_tp4_markdown_fences(
9931028 language = _TP4_MARKDOWN_EXECUTABLE_LABELS .get (active [2 ])
9941029 body_text = "" .join (body )
9951030 if language is not None and body_text .strip ():
996- yield language , body_text , body_start , line_number - 1
1031+ yield (
1032+ language ,
1033+ body_text ,
1034+ body_start ,
1035+ line_number - 1 ,
1036+ _tp4_fence_context (preceding ),
1037+ )
9971038 active = None
9981039 body = []
1040+ preceding = []
9991041 continue
10001042 body .append (line )
10011043
@@ -1316,9 +1358,18 @@ def plan_candidate(candidate: _TP4Candidate) -> None:
13161358 )
13171359 )
13181360 break
1361+ # The context block is part of the prompt, so it must also
1362+ # participate in the same token and byte budgets as the code.
1363+ context_block = (
1364+ f"Document context (verbatim, not part of the code):\n { candidate .context } \n "
1365+ if candidate .context
1366+ else ""
1367+ )
13191368 prompt = (
13201369 prefix
1321- + f"### { candidate .path } ({ candidate .language } )\n { chunk .content } "
1370+ + f"### { candidate .path } ({ candidate .language } )\n "
1371+ + context_block
1372+ + f"{ chunk .content } "
13221373 + _TP4_PROMPT_SUFFIX
13231374 )
13241375 if estimate_tokens (prompt ) > batch_input_tokens :
@@ -1461,8 +1512,10 @@ def iter_candidates() -> Iterator[_TP4Candidate]:
14611512 return
14621513 if not retained :
14631514 continue
1464- for language , body , start_line , end_line in _iter_tp4_markdown_fences (retained ):
1465- yield _TP4Candidate (path , language , body , start_line , end_line )
1515+ for language , body , start_line , end_line , context in _iter_tp4_markdown_fences (
1516+ retained
1517+ ):
1518+ yield _TP4Candidate (path , language , body , start_line , end_line , context )
14661519 if stop_reason is not None :
14671520 break
14681521
0 commit comments