Langgraph 클래스 기반 코드 리팩토링

최근 langgraph를 사용할 때는 단순히 def 함수로 로직을 구현했던 경험이 있었다.

그런데 이번에 투자 AI 모임 코드를 보면서 각 모듈이 클래스를 기반으로 설계되어 있다는 점이 눈에 띄었다.

기존 함수 기반 구현 vs 클래스 기반 구현

함수 기반 코드는 짧고 간단한 작업을 수행할 때 빠르게 작성할 수 있지만 다음과 같은 한계가 있다.

1. 상태 관리의 어려움: 여러 함수에 걸쳐 데이터를 주고받으며 상태를 관리해야 하는 경우 복잡성이 증가함.
2. 코드 중복: 반복되는 로직을 분리하여 관리하기 어렵고, 동일한 코드가 여러 곳에 등장할 수 있음.
3. 유지보수 난이도: 프로젝트 규모가 커질수록 어떤 함수가 어느 역할을 수행하는지 파악하기 어려워짐.

반면 클래스 기반 구현은 관련 변수와 함수를 하나의 객체로 묶어 캡슐화함으로써 다음과 같은 장점을 제공한다.

1. 모듈화: 클래스 내부에 상태와 메서드를 묶어 관리하여 코드 구조를 명확하게 함.
2. 재사용성: 상속과 다형성을 통해 반복되는 로직을 최소화할 수 있음.
3. 테스트 용이성: 객체 단위의 테스트가 가능해지고 Mocking이나 스텁을 활용해 세밀한 단위 테스트 작성이 쉬워짐.

 

최근 후임자에게 인수인계를 할때 상태관리에 대해 설명하기 복잡한 경우가 있었는데

class 기반 구축이 되어있었다면 어땠을까를 다시한번 생각하게 되었다.

코드 예시

class RewrittenFactorAgent:

    def __init__(self, llm):
        self.llm = llm

    def make_ast(self, state: FactorAgentState):
        print(f"{state['hypothesis']}")
        prompt = ChatPromptTemplate.from_messages(
            [
                ("system", Prompt["system"]),
                ("human", Prompt["human"]),
            ]
        )
        chain = prompt | self.llm
 		# ...생략...

        state["market"] = market_result
        return state

    def get_market_prices(self, tickers):
        # ...생략...
        return {}