laas.py

import logging
from typing import Any, List, Optional

import requests
from langchain_core.callbacks import CallbackManagerForLLMRun
from langchain_core.language_models import BaseChatModel, BaseLanguageModel
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
from langchain_core.pydantic_v1 import BaseModel, Field, SecretStr, root_validator

logger = logging.getLogger(__name__)


class ChatLaaS(BaseChatModel):
    laas_api_base: Optional[str] = Field(
        default="https://api-laas.wanted.co.kr/api/preset", alias="base_url"
    )
    laas_api_key: Optional[SecretStr] = Field(default=None, alias="api_key")
    laas_project: Optional[str] = Field(default=None, alias="project")
    laas_hash: Optional[str] = Field(default=None, alias="hash")
    timeout: Optional[float] = Field(default=60.0)

    _ROLE_MAP = {
        "human": "user",
        "ai": "assistant",
    }

    @property
    def _llm_type(self) -> str:
        """Return type of chat model."""
        return "laas-chat"

    @classmethod
    def is_lc_serializable(cls) -> bool:
        """Return whether this model can be serialized by Langchain."""
        return False

    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> ChatResult:
        try:
            body = {
                "hash": self.laas_hash,
                "messages": [
                    {
                        "role": self._ROLE_MAP.get(msg.type, msg.type),
                        "content": msg.content,
                    }
                    for msg in messages
                    if msg.content.strip()  # This filters out empty or whitespace-only content
                ],
                **kwargs,
            }
            print(body)
            # return

            headers = {
                "Content-Type": "application/json",
                "apiKey": self.laas_api_key.get_secret_value(),
                "project": self.laas_project,
            }

            response = requests.post(
                f"{self.laas_api_base}/chat/completions",
                headers=headers,
                json=body,
                timeout=self.timeout,
            ).json()

            # Extract the content from the API response
            content = response["choices"][0]["message"]["content"]
            message = AIMessage(id=response["id"], content=content)
            generation = ChatGeneration(message=message)
            return ChatResult(generations=[generation])
        except Exception as e:
            raise