langchain 공부

ChatOllama 내부 코드 분석 1

필만이 2024. 11. 17. 23:13

배경

  • ollama의 내부 코드 살펴보고자함

코드

import json
from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union, cast

from langchain_core._api import deprecated
from langchain_core.callbacks import (
    AsyncCallbackManagerForLLMRun,
    CallbackManagerForLLMRun,
)
from langchain_core.language_models.chat_models import BaseChatModel, LangSmithParams
from langchain_core.messages import (
    AIMessage,
    AIMessageChunk,
    BaseMessage,
    ChatMessage,
    HumanMessage,
    SystemMessage,
)
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult

from langchain_community.llms.ollama import OllamaEndpointNotFoundError, _OllamaCommon

@deprecated("0.0.3", alternative="_chat_stream_response_to_chat_generation_chunk")
def _stream_response_to_chat_generation_chunk(
    stream_response: str,
) -> ChatGenerationChunk:
    """
    이전 버전과의 호환성을 위해 남겨진 함수로,
    JSON 형식의 스트림 응답을 ChatGenerationChunk 객체로 변환합니다.

    **Note**: 이 함수는 "0.0.3" 버전 이후로 폐기(deprecated)되었으며,
    `_chat_stream_response_to_chat_generation_chunk`를 대신 사용할 것을 권장합니다.

    Args:
        stream_response (str): 스트림으로 전달된 JSON 응답 문자열.

    Returns:
        ChatGenerationChunk: 메시지와 생성 정보가 포함된 ChatGenerationChunk 객체.
    """
    # JSON 문자열을 파싱하여 Python 딕셔너리로 변환
    parsed_response = json.loads(stream_response)

    # 생성 정보(generation_info)를 결정
    # "done" 키가 True인 경우, 응답 전체를 생성 정보로 저장
    generation_info = parsed_response if parsed_response.get("done") is True else None

    # ChatGenerationChunk 객체를 반환
    return ChatGenerationChunk(
        # 메시지 부분은 AIMessageChunk 객체로 래핑하여, "response" 키의 값을 설정
        message=AIMessageChunk(content=parsed_response.get("response", "")),
        # "done" 값에 따라 생성 정보 설정
        generation_info=generation_info,
    )


def _chat_stream_response_to_chat_generation_chunk(
    stream_response: str,
) -> ChatGenerationChunk:
    """
    JSON 형식의 스트림 응답을 ChatGenerationChunk 객체로 변환합니다.

    이 함수는 `_stream_response_to_chat_generation_chunk`의 최신 버전으로,
    새로운 JSON 구조를 지원합니다.

    Args:
        stream_response (str): 스트림으로 전달된 JSON 응답 문자열.

    Returns:
        ChatGenerationChunk: 메시지와 생성 정보가 포함된 ChatGenerationChunk 객체.
    """
    # JSON 문자열을 파싱하여 Python 딕셔너리로 변환
    parsed_response = json.loads(stream_response)

    # 생성 정보(generation_info)를 결정
    # "done" 키가 True인 경우, 응답 전체를 생성 정보로 저장
    # 목적 : 작업 완료 상태에서만 유효한 정보를 처리함으로써, 불완전하거나 불필요한 데이터가 나중에 처리 로직에 영향을 미치지 않도록 보장
    generation_info = parsed_response if parsed_response.get("done") is True else None

    # ChatGenerationChunk 객체를 반환
    return ChatGenerationChunk(
        # 메시지 부분은 AIMessageChunk 객체로 래핑하여, "message.content" 키의 값을 설정
        message=AIMessageChunk(
            # 딕셔너리에서 "message" 키의 값을 가져옵니다.만약 "message" 키가 존재하지 않으면, 기본값 {}(빈 딕셔너리)를 반환
            # 만약 "content" 키가 존재하지 않으면, 기본값 ""(빈 문자열)를 반환
            content=parsed_response.get("message", {}).get("content", "")
        ),
        # "done" 값에 따라 생성 정보 설정
        generation_info=generation_info,
    )

class ChatOllama(BaseChatModel, _OllamaCommon):
    """
    Ollama를 사용하여 로컬에서 대규모 언어 모델(LLM)을 실행하는 Chat 모델 클래스.

    Ollama는 로컬에서 LLM을 실행할 수 있는 플랫폼으로, 이를 사용하여 채팅 기반 애플리케이션을 구현합니다.

    **Example**:
        .. code-block:: python

            from langchain_community.chat_models import ChatOllama
            ollama = ChatOllama(model="llama2")

    """

    @property
    def _llm_type(self) -> str:
        """
        이 Chat 모델의 유형을 반환합니다.

        Returns:
            str: 모델의 유형. 여기서는 "ollama-chat".
        """
        return "ollama-chat"

    @classmethod
    def is_lc_serializable(cls) -> bool:
        """
        이 모델이 LangChain에서 직렬화 가능한지 여부를 반환.

        Returns:
            bool: False. 이 모델은 LangChain 직렬화를 지원하지 않음.
        """
        return False

    def _get_ls_params(
        self, stop: Optional[List[str]] = None, **kwargs: Any
    ) -> LangSmithParams:
        """
        LangSmith 추적에 사용될 표준 파라미터를 생성합니다.

        Args:
            stop (Optional[List[str]]): 중단할 토큰의 리스트.
            kwargs: 추가적인 파라미터.

        Returns:
            LangSmithParams: LangSmith에 사용될 파라미터 객체.
        """
        params = self._get_invocation_params(stop=stop, **kwargs)  # 호출 파라미터 가져오기
        ls_params = LangSmithParams(
            ls_provider="ollama",  # 제공자 이름
            ls_model_name=self.model,  # 모델 이름
            ls_model_type="chat",  # 모델 유형
            ls_temperature=params.get("temperature", self.temperature),  # 생성 온도값
        )
        if ls_max_tokens := params.get("num_predict", self.num_predict):  # 생성 최대 토큰
            ls_params["ls_max_tokens"] = ls_max_tokens
        if ls_stop := stop or params.get("stop", None) or self.stop:  # 중단 조건
            ls_params["ls_stop"] = ls_stop
        return ls_params

    @deprecated("0.0.3", alternative="_convert_messages_to_ollama_messages")
    def _format_message_as_text(self, message: BaseMessage) -> str:
        """
        주어진 메시지를 텍스트 형식으로 변환합니다.

        **Note**: 이 메서드는 "0.0.3" 버전 이후로 폐기되었습니다. 
        `_convert_messages_to_ollama_messages`를 대신 사용하는 것을 권장합니다.

        Args:
            message (BaseMessage): 변환할 메시지. 
                - `ChatMessage`, `HumanMessage`, `AIMessage`, `SystemMessage` 등의 타입으로 전달될 수 있음.

        Returns:
            str: 텍스트 형식으로 변환된 메시지.

        Raises:
            ValueError: 지원되지 않는 메시지 타입이 전달된 경우 예외 발생.
        """
        if isinstance(message, ChatMessage):
            # ChatMessage 타입의 메시지를 처리
            # 역할(role)과 내용을 포함하여 메시지를 포맷팅
            # 예: "User: Hello"
            message_text = f"\n\n{message.role.capitalize()}: {message.content}"
        elif isinstance(message, HumanMessage):
            # HumanMessage 타입의 메시지를 처리
            # HumanMessage는 사용자가 입력한 메시지를 나타냄
            if isinstance(message.content, List):  # 콘텐츠가 리스트인 경우
                # 첫 번째 요소를 가져옴
                first_content = cast(List[Dict], message.content)[0]  # 타입 힌트를 명시적으로 변환
                content_type = first_content.get("type")  # 콘텐츠 타입 확인

                if content_type == "text":
                    # 콘텐츠 타입이 "text"인 경우, 해당 텍스트를 포맷
                    # 예: "[INST] Some text [/INST]"
                    message_text = f"[INST] {first_content['text']} [/INST]"
                elif content_type == "image_url":
                    # 콘텐츠 타입이 "image_url"인 경우, URL을 반환
                    # 예: "https://example.com/image.jpg"
                    message_text = first_content["image_url"]["url"]
            else:
                # 콘텐츠가 리스트가 아닌 일반 텍스트인 경우, 포맷팅
                # 예: "[INST] Hello [/INST]"
                message_text = f"[INST] {message.content} [/INST]"
        elif isinstance(message, AIMessage):
            # AIMessage 타입의 메시지를 처리
            # AIMessage는 AI 모델의 응답 메시지를 나타냄
            # 예: "This is the AI's response."
            message_text = f"{message.content}"
        elif isinstance(message, SystemMessage):
            # SystemMessage 타입의 메시지를 처리
            # 시스템 메시지는 특별한 명령이나 설정을 나타냄
            # 예: "<<SYS>> System info <<SYS>>"
            message_text = f"<<SYS>> {message.content} <</SYS>>"
        else:
            # 알 수 없는 메시지 타입이 전달된 경우 예외를 발생시킴
            # 사용자는 적합한 메시지 타입을 제공해야 함
            raise ValueError(f"Got unknown type {message}")

        # 포맷된 메시지를 반환
        return message_text


    def _format_messages_as_text(self, messages: List[BaseMessage]) -> str:
        """
        여러 메시지를 텍스트 형식으로 변환합니다.

        Args:
            messages (List[BaseMessage]): 변환할 메시지 리스트.

        Returns:
            str: 변환된 텍스트.
        """
        return "\n".join(
            [self._format_message_as_text(message) for message in messages]
        )

    def _convert_messages_to_ollama_messages(
        self, messages: List[BaseMessage]
    ) -> List[Dict[str, Union[str, List[str]]]]:
        """
        메시지를 Ollama 형식에 맞게 변환합니다.

        Args:
            messages (List[BaseMessage]): 변환할 메시지 리스트.

        Returns:
            List[Dict[str, Union[str, List[str]]]]: Ollama에 적합한 메시지 리스트.
        """
        ollama_messages: List = []
        for message in messages:
            # 역할(role)을 결정
            role = ""
            if isinstance(message, HumanMessage):
                role = "user"
            elif isinstance(message, AIMessage):
                role = "assistant"
            elif isinstance(message, SystemMessage):
                role = "system"
            else:
                raise ValueError("Received unsupported message type for Ollama.")

            # 메시지 콘텐츠(content)와 이미지(images)를 초기화
            content = ""
            images = []
            if isinstance(message.content, str):
                # 콘텐츠가 문자열인 경우 직접 할당
                content = message.content
            else:
                # 콘텐츠가 리스트 형식인 경우 각 요소 처리
                for content_part in cast(List[Dict], message.content):
                    if content_part.get("type") == "text":
                        content += f"\n{content_part['text']}"
                    elif content_part.get("type") == "image_url":
                        image_url = None
                        temp_image_url = content_part.get("image_url")
                        if isinstance(temp_image_url, str):
                            image_url = temp_image_url
                        elif (
                            isinstance(temp_image_url, dict) and "url" in temp_image_url
                        ):
                            image_url = temp_image_url["url"]
                        else:
                            raise ValueError(
                                "Only string image_url or dict with string 'url' "
                                "inside content parts are supported."
                            )

                        # 이미지 URL 처리
                        image_url_components = image_url.split(",")
                        # "data:image/jpeg;base64,<image>"와 같은 포맷 지원
                        if len(image_url_components) > 1:
                            images.append(image_url_components[1])
                        else:
                            images.append(image_url_components[0])
                    else:
                        # 지원되지 않는 콘텐츠 타입에 대한 에러 처리
                        raise ValueError(
                            "Unsupported message content type. "
                            "Must either have type 'text' or type 'image_url' "
                            "with a string 'image_url' field."
                        )

            # Ollama 메시지 포맷에 맞게 변환
            ollama_messages.append(
                {
                    "role": role,  # 메시지 역할(user, assistant, system)
                    "content": content,  # 메시지 내용
                    "images": images,  # 포함된 이미지 리스트
                }
            )

        return ollama_messages