langchain 공부

InMemoryDocstore 내부 코드 분석

필만이 2024. 11. 7. 23:23

배경

  • langchain 에 쓰이는 도구들을 깊게 이해하고 사용하기 위해 도구들을 하나씩 분석

코드분석

  • add 메서드: {"doc1": doc1, "doc2": doc2} 형태의 딕셔너리를 사용해 새 문서를 추가합니다. 중복된 ID가 없으면 성공적으로 추가되며, 중복된 ID가 있으면 오류를 발생
  • delete 메서드: 지정된 ID의 문서를 삭제합니다. 존재하지 않는 ID를 삭제하려고 시도하면 오류가 발생
  • search 메서드: doc2라는 ID로 문서를 검색합니다. 문서가 존재하면 해당 문서 객체가 반환되고, 존재하지 않으면 "ID not found" 메시지를 반환
"""Simple in memory docstore in the form of a dict."""

from typing import Dict, List, Optional, Union

from langchain_core.documents import Document

from langchain_community.docstore.base import AddableMixin, Docstore


# AddableMixin은 문서를 추가하는 기능을 지원하는 믹스인 클래스입니다. 믹스인은 일반적으로 특정 기능을 제공하기 위해 
# 여러 클래스와 함께 사용할 수 있는 작은 클래스이며, 이 경우는 문서를 추가하는 기능을 담당합니다. 
# InMemoryDocstore는 이 믹스인을 통해 문서 추가 기능을 상속받아 사용할 수 있습니다.

# 따라서 InMemoryDocstore 클래스는 Docstore의 기본 저장소 기능과 AddableMixin의 추가 기능을 모두 상속받아, 
# 메모리 내에서 문서의 저장과 관리를 가능하게 합니다.
class InMemoryDocstore(Docstore, AddableMixin):
    """dict 형태로 메모리에 저장되는 간단한 문서 저장소 클래스."""

    def __init__(self, _dict: Optional[Dict[str, Document]] = None):
        """dict 초기화.

        Args:
            _dict: 문서 ID(str)를 키로, 문서 객체(Document)를 값으로 가지는 딕셔너리. 
                   지정하지 않으면 빈 딕셔너리로 초기화됩니다.
        """
        # 초기화된 딕셔너리가 있으면 그것을 사용, 없으면 빈 딕셔너리로 초기화
        self._dict = _dict if _dict is not None else {}

    def add(self, texts: Dict[str, Document]) -> None:
        """문서를 메모리 딕셔너리에 추가.

        Args:
            texts: ID(str)를 키로 하고 문서 객체(Document)를 값으로 가지는 딕셔너리.

        Raises:
            ValueError: 이미 존재하는 ID를 추가하려고 할 때 발생.
        """
        # 추가하려는 문서 ID가 이미 존재하는지 확인
        # 해시 테이블을 사용해 교집합(intersection)을 구함.
        overlapping = set(texts).intersection(self._dict)
        if overlapping:
            # 이미 존재하는 ID가 있으면 오류 발생
            raise ValueError(f"Tried to add ids that already exist: {overlapping}")
        # 딕셔너리에 새로운 문서들을 병합
        self._dict = {**self._dict, **texts}

    def delete(self, ids: List) -> None:
        """지정한 ID의 문서를 메모리 딕셔너리에서 삭제.

        Args:
            ids: 삭제할 문서 ID들의 리스트.

        Raises:
            ValueError: 삭제하려는 ID가 존재하지 않는 경우 발생.
        """
        # 삭제하려는 ID 중 실제로 존재하는 ID들 확인
        overlapping = set(ids).intersection(self._dict)
        if not overlapping:
            # 삭제하려는 ID가 전부 존재하지 않으면 오류 발생
            raise ValueError(f"Tried to delete ids that do not exist: {ids}")

        # 존재하는 각 ID에 대해 딕셔너리에서 해당 문서 삭제
        for _id in ids:
            self._dict.pop(_id)

    def search(self, search: str) -> Union[str, Document]:
        """문서 ID를 통해 문서를 검색.

        Args:
            search: 검색할 문서의 ID.

        Returns:
            Document: 검색된 문서 객체를 반환.
            str: 해당 ID가 존재하지 않으면 오류 메시지를 반환.
        """
        # 딕셔너리에 검색할 ID가 존재하지 않으면 오류 메시지 반환
        if search not in self._dict:
            return f"ID {search} not found."
        else:
            # 존재하면 해당 문서 객체 반환
            return self._dict[search]