langchain/libs/community/langchain_community/chat_models/symblai_nebula.py

import json
import os
from json import JSONDecodeError
from typing import Any, AsyncIterator, Dict, Iterator, List, Optional

import requests
from aiohttp import ClientSession
from langchain_core.callbacks import (
    AsyncCallbackManagerForLLMRun,
    CallbackManagerForLLMRun,
)
from langchain_core.language_models.chat_models import (
    BaseChatModel,
    agenerate_from_stream,
    generate_from_stream,
)
from langchain_core.messages import AIMessage, AIMessageChunk, BaseMessage
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
from langchain_core.utils import convert_to_secret_str
from pydantic import ConfigDict, Field, SecretStr


def _convert_role(role: str) -> str:
    map = {"ai": "assistant", "human": "human", "chat": "human"}
    if role in map:
        return map[role]
    else:
        raise ValueError(f"Unknown role type: {role}")


def _format_nebula_messages(messages: List[BaseMessage]) -> Dict[str, Any]:
    system = ""
    formatted_messages = []
    for message in messages[:-1]:
        if message.type == "system":
            if isinstance(message.content, str):
                system = message.content
            else:
                raise ValueError("System prompt must be a string")
        else:
            formatted_messages.append(
                {
                    "role": _convert_role(message.type),
                    "text": message.content,
                }
            )

    text = messages[-1].content
    formatted_messages.append({"role": "human", "text": text})
    return {"system_prompt": system, "messages": formatted_messages}


class ChatNebula(BaseChatModel):
    """`Nebula` chat large language model - https://docs.symbl.ai/docs/nebula-llm

    API Reference: https://docs.symbl.ai/reference/nebula-chat

    To use, set the environment variable ``NEBULA_API_KEY``,
    or pass it as a named parameter to the constructor.
    To request an API key, visit https://platform.symbl.ai/#/login
    Example:
        .. code-block:: python

            from langchain_community.chat_models import ChatNebula
            from langchain_core.messages import SystemMessage, HumanMessage

            chat = ChatNebula(max_new_tokens=1024, temperature=0.5)

            messages = [
            SystemMessage(
                content="You are a helpful assistant."
            ),
            HumanMessage(
                "Answer the following question. How can I help save the world."
            ),
            ]
            chat.invoke(messages)
    """

    max_new_tokens: int = 1024
    """Denotes the number of tokens to predict per generation."""

    temperature: Optional[float] = 0
    """A non-negative float that tunes the degree of randomness in generation."""

    streaming: bool = False

    nebula_api_url: str = "https://api-nebula.symbl.ai"

    nebula_api_key: Optional[SecretStr] = Field(None, description="Nebula API Token")

    model_config = ConfigDict(
        populate_by_name=True,
        arbitrary_types_allowed=True,
    )

    def __init__(self, **kwargs: Any) -> None:
        if "nebula_api_key" in kwargs:
            api_key = convert_to_secret_str(kwargs.pop("nebula_api_key"))
        elif "NEBULA_API_KEY" in os.environ:
            api_key = convert_to_secret_str(os.environ["NEBULA_API_KEY"])
        else:
            api_key = None
        super().__init__(nebula_api_key=api_key, **kwargs)  # type: ignore[call-arg]

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

    @property
    def _api_key(self) -> str:
        if self.nebula_api_key:
            return self.nebula_api_key.get_secret_value()
        return ""

    def _stream(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> Iterator[ChatGenerationChunk]:
        """Call out to Nebula's chat endpoint."""
        url = f"{self.nebula_api_url}/v1/model/chat/streaming"
        headers = {
            "ApiKey": self._api_key,
            "Content-Type": "application/json",
        }
        formatted_data = _format_nebula_messages(messages=messages)
        payload: Dict[str, Any] = {
            "max_new_tokens": self.max_new_tokens,
            "temperature": self.temperature,
            **formatted_data,
            **kwargs,
        }

        payload = {k: v for k, v in payload.items() if v is not None}
        json_payload = json.dumps(payload)

        response = requests.request(
            "POST", url, headers=headers, data=json_payload, stream=True
        )
        response.raise_for_status()

        for chunk_response in response.iter_lines():
            chunk_decoded = chunk_response.decode()[6:]
            try:
                chunk = json.loads(chunk_decoded)
            except JSONDecodeError:
                continue
            token = chunk["delta"]
            cg_chunk = ChatGenerationChunk(message=AIMessageChunk(content=token))
            if run_manager:
                run_manager.on_llm_new_token(token, chunk=cg_chunk)
            yield cg_chunk

    async def _astream(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> AsyncIterator[ChatGenerationChunk]:
        url = f"{self.nebula_api_url}/v1/model/chat/streaming"
        headers = {"ApiKey": self._api_key, "Content-Type": "application/json"}
        formatted_data = _format_nebula_messages(messages=messages)
        payload: Dict[str, Any] = {
            "max_new_tokens": self.max_new_tokens,
            "temperature": self.temperature,
            **formatted_data,
            **kwargs,
        }

        payload = {k: v for k, v in payload.items() if v is not None}
        json_payload = json.dumps(payload)

        async with ClientSession() as session:
            async with session.post(  # type: ignore[call-arg]
                url, data=json_payload, headers=headers, stream=True
            ) as response:
                response.raise_for_status()
                async for chunk_response in response.content:
                    chunk_decoded = chunk_response.decode()[6:]
                    try:
                        chunk = json.loads(chunk_decoded)
                    except JSONDecodeError:
                        continue
                    token = chunk["delta"]
                    cg_chunk = ChatGenerationChunk(
                        message=AIMessageChunk(content=token)
                    )
                    if run_manager:
                        await run_manager.on_llm_new_token(token, chunk=cg_chunk)
                    yield cg_chunk

    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> ChatResult:
        if self.streaming:
            stream_iter = self._stream(
                messages, stop=stop, run_manager=run_manager, **kwargs
            )
            return generate_from_stream(stream_iter)

        url = f"{self.nebula_api_url}/v1/model/chat"
        headers = {"ApiKey": self._api_key, "Content-Type": "application/json"}
        formatted_data = _format_nebula_messages(messages=messages)
        payload: Dict[str, Any] = {
            "max_new_tokens": self.max_new_tokens,
            "temperature": self.temperature,
            **formatted_data,
            **kwargs,
        }

        payload = {k: v for k, v in payload.items() if v is not None}
        json_payload = json.dumps(payload)

        response = requests.request("POST", url, headers=headers, data=json_payload)
        response.raise_for_status()
        data = response.json()

        return ChatResult(
            generations=[ChatGeneration(message=AIMessage(content=data["messages"]))],
            llm_output=data,
        )

    async def _agenerate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> ChatResult:
        if self.streaming:
            stream_iter = self._astream(
                messages, stop=stop, run_manager=run_manager, **kwargs
            )
            return await agenerate_from_stream(stream_iter)

        url = f"{self.nebula_api_url}/v1/model/chat"
        headers = {"ApiKey": self._api_key, "Content-Type": "application/json"}
        formatted_data = _format_nebula_messages(messages=messages)
        payload: Dict[str, Any] = {
            "max_new_tokens": self.max_new_tokens,
            "temperature": self.temperature,
            **formatted_data,
            **kwargs,
        }

        payload = {k: v for k, v in payload.items() if v is not None}
        json_payload = json.dumps(payload)

        async with ClientSession() as session:
            async with session.post(
                url, data=json_payload, headers=headers
            ) as response:
                response.raise_for_status()
                data = await response.json()

                return ChatResult(
                    generations=[
                        ChatGeneration(message=AIMessage(content=data["messages"]))
                    ],
                    llm_output=data,
                )