자기 교정(Self-correction)을 갖춘 Codestral
자기 교정(Self-correction)을 갖춘 Codestral (Codestral with self-correction)
Codestral의 코드 생성 능력에 [AlphaCodium] 논문의 자기 교정 접근을 결합해, [코딩 질문에 대한 답을 반복적으로 구성]하는 방법을 배우는 문서예요. [LangGraph]를 사용해 구조화된 코드 생성 출력, 인라인 단위 테스트, 오류 피드백을 통한 자기 교정을 구현합니다.
출처: 문서
본문
Codestral은 fill-in-the-middle과 코드 완성 등 코드 생성 작업에 특화되고 최적화된 최첨단 생성 모델이에요. Codestral은 80개 이상의 프로그래밍 언어로 학습되어, 흔한 언어와 덜 흔한 언어 모두에서 좋은 성능을 발휘합니다.
Codestral의 코드 생성 능력과 [AlphaCodium] 논문에서 제시된 자기 교정 접근을 결합할 수 있어요. 즉 [코딩 질문에 대한 답을 반복적으로 구성]하는 방식이죠.
이 아이디어 중 일부를 [LangGraph]를 사용해 처음부터 구현할 거예요: 1) Codestral-instruct로부터 구조화된 코드 생성 출력 생성, 2) 임포트와 코드 실행이 동작하는지 인라인 단위 테스트 수행, 3) 오류를 Codestral에 다시 피드백해 자기 교정.
! pip install -U langchain_community langchain-mistralai langchain langgraph
LLM
Mistral API와 도구 사용(tool use)을 지원하는 Codestral instruct 모델을 사용할게요.
import os, getpass
def _set_env(var: str):
if not os.environ.get(var):
os.environ[var] = getpass.getpass(f"{var}: ")
_set_env("MISTRAL_API_KEY")
os.environ['TOKENIZERS_PARALLELISM'] = 'true'
선택적으로 트레이싱에 [LangSmith]를 사용할 수 있어요.
_set_env("LANGCHAIN_API_KEY")
os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_PROJECT"] = "mistral-cookbook"
코드 생성 (Code Generation)
구조화된 출력으로 테스트해요.
# Select LLM
from langchain_mistralai import ChatMistralAI
from langchain_core.prompts import ChatPromptTemplate
from pydantic import BaseModel, Field
# Mistral model
mistral_model = "codestral-latest"
llm = ChatMistralAI(model=mistral_model, temperature=0)
# Prompt
code_gen_prompt_claude = ChatPromptTemplate.from_messages(
[
(
"system",
"""You are a coding assistant. Ensure any code you provide can be executed with all required imports and variables \n
defined. Structure your answer: 1) a prefix describing the code solution, 2) the imports, 3) the functioning code block.
\n Here is the user question:""",
),
("placeholder", "{messages}"),
]
)
# Data model
class code(BaseModel):
"""Schema for code solutions to questions about LCEL."""
prefix: str = Field(description="Description of the problem and approach")
imports: str = Field(description="Code block import statements")
code: str = Field(description="Code block not including import statements")
# LLM
code_gen_chain = llm.with_structured_output(code, include_raw=False)
question = "Write a function for fibonacci."
messages = [("user", question)]
# Test
result = code_gen_chain.invoke(messages)
result
코드 생성 사슬은 prefix(문제·접근 설명), imports(임포트 문), code(임포트 제외 코드 블록) 세 부분의 구조화된 출력을 반환해요.
그래프 (Graph)
위 워크플로우를 [LangGraph]를 사용해 그래프로 구축해요.
그래프 상태 (Graph state)
그래프 state 스키마에는 우리가 원하는 키들이 포함돼요:
- 그래프의 각 노드에 전달
- 선택적으로 각 노드에서 수정
개념 문서는 [여기]를 참조하세요.
from typing import Annotated
from typing import Dict, TypedDict, List
from langgraph.graph.message import AnyMessage, add_messages
class GraphState(TypedDict):
"""
Represents the state of our graph.
Attributes:
error : Binary flag for control flow to indicate whether test error was tripped
messages : With user question, error messages, reasoning
generation : Code solution
iterations : Number of tries
"""
error: str
messages: Annotated[list[AnyMessage], add_messages]
generation: str
iterations: int
그래프 노드 (Graph nodes)
노드를 정의해요: generate(코드 해결책 생성), code_check(임포트·실행 검사), 그리고 조건부 엣지 decide_to_finish(오류가 없거나 최대 반복에 도달하면 종료).
from operator import itemgetter
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_core.runnables import RunnablePassthrough
from langchain_core.prompts import PromptTemplate
### Parameters
max_iterations = 3
### Nodes
def generate(state: GraphState):
"""
Generate a code solution
Args:
state (dict): The current graph state
Returns:
state (dict): New key added to state, generation
"""
print("---GENERATING CODE SOLUTION---")
# State
messages = state["messages"]
iterations = state["iterations"]
error = state.get("error", "")
# Solution
code_solution = code_gen_chain.invoke(messages)
messages += [
(
"assistant",
f"Here is my attempt to solve the problem: {code_solution.prefix} \n Imports: {code_solution.imports} \n Code: {code_solution.code}",
)
]
# Increment
iterations = iterations + 1
return {"generation": code_solution, "messages": messages, "iterations": iterations}
def code_check(state: GraphState):
"""
Check code
Args:
state (dict): The current graph state
Returns:
state (dict): New key added to state, error
"""
print("---CHECKING CODE---")
# State
messages = state["messages"]
code_solution = state["generation"]
iterations = state["iterations"]
# Get solution components
prefix = code_solution.prefix
imports = code_solution.imports
code = code_solution.code
# Check imports
try:
exec(imports)
except Exception as e:
print("---CODE IMPORT CHECK: FAILED---")
error_message = [("user", f"Your solution failed the import test. Here is the error: {e}. Reflect on this error and your prior attempt to solve the problem. (1) State what you think went wrong with the prior solution and (2) try to solve this problem again. Return the FULL SOLUTION. Use the code tool to structure the output with a prefix, imports, and code block:")]
messages += error_message
return {
"generation": code_solution,
"messages": messages,
"iterations": iterations,
"error": "yes",
}
# Check execution
try:
combined_code = f"{imports}\n{code}"
# Use a shared scope for exec
global_scope = {}
exec(combined_code, global_scope)
except Exception as e:
print("---CODE BLOCK CHECK: FAILED---")
error_message = [("user", f"Your solution failed the code execution test: {e}) Reflect on this error and your prior attempt to solve the problem. (1) State what you think went wrong with the prior solution and (2) try to solve this problem again. Return the FULL SOLUTION. Use the code tool to structure the output with a prefix, imports, and code block:")]
messages += error_message
return {
"generation": code_solution,
"messages": messages,
"iterations": iterations,
"error": "yes",
}
# No errors
print("---NO CODE TEST FAILURES---")
return {
"generation": code_solution,
"messages": messages,
"iterations": iterations,
"error": "no",
}
### Conditional edges
def decide_to_finish(state: GraphState):
"""
Determines whether to finish.
Args:
state (dict): The current graph state
Returns:
str: Next node to call
"""
error = state["error"]
iterations = state["iterations"]
if error == "no" or iterations == max_iterations:
print("---DECISION: FINISH---")
return "end"
else:
print("---DECISION: RE-TRY SOLUTION---")
return "generate"
그래프에 [체크포인터(checkpointer)]를 사용해 영속성을 추가할 거예요.
from IPython.display import Image, display
from langgraph.graph import END, StateGraph
# Define the graph
builder = StateGraph(GraphState)
# Define the nodes
builder.add_node("generate", generate) # generation solution
builder.add_node("check_code", code_check) # check code
# Build graph
builder.set_entry_point("generate")
builder.add_edge("generate", "check_code")
builder.add_conditional_edges(
"check_code",
decide_to_finish,
{
"end": END,
"generate": "generate",
},
)
graph = builder.compile()
display(Image(graph.get_graph(xray=True).draw_mermaid_png()))
간단한 질문으로 그래프를 실행해보면:
from langchain_core.messages import HumanMessage
question = "Write a Python program that prints 'Hello, World!' to the console."
for event in graph.stream({"messages": [HumanMessage(content=question)], "iterations": 0}, stream_mode="values"):
print(event)
가장 대표적인 예시로 두 명의 플레이어가 3x3 격자에서 틱택토 게임을 하는 Python 프로그램을 만드는 질문도 시도할 수 있어요 (유효하지 않은 이동 검사, 승자·무승부 판정, 2D 리스트 사용, 함수 모듈화, 입력 검증 등 요구사항).
question = """Create a Python program that allows two players to play a game of Tic-Tac-Toe. The game should be played on a 3x3 grid. The program should:
- Allow players to take turns to input their moves.
- Check for invalid moves (e.g., placing a marker on an already occupied space).
- Determine and announce the winner or if the game ends in a draw.
Requirements:
- Use a 2D list to represent the Tic-Tac-Toe board.
- Use functions to modularize the code.
- Validate player input.
- Check for win conditions and draw conditions after each move."""
for event in graph.stream({"messages": [HumanMessage(content=question)], "iterations": 0}, stream_mode="values"):
print(event)
그래프 흐름을 요약하면: generate가 구조화된 코드를 만들고, code_check가 임포트와 코드 실행을 검사하며, 실패하면 오류 메시지를 피드백해 generate로 되돌아가고, 성공하거나 max_iterations(여기선 3)에 도달하면 종료합니다.