"""Writer chat wrapper.""" from __future__ import annotations import json import logging from typing import ( Any, AsyncIterator, Callable, Dict, Iterator, List, Literal, Optional, Sequence, Tuple, Type, Union, ) from langchain_core.callbacks import ( AsyncCallbackManagerForLLMRun, CallbackManagerForLLMRun, ) from langchain_core.language_models import LanguageModelInput from langchain_core.language_models.chat_models import ( BaseChatModel, ) from langchain_core.messages import ( AIMessage, AIMessageChunk, BaseMessage, ChatMessage, HumanMessage, SystemMessage, ToolMessage, ) from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult from langchain_core.runnables import Runnable from langchain_core.utils import get_from_dict_or_env from langchain_core.utils.function_calling import convert_to_openai_tool from pydantic import BaseModel, ConfigDict, Field, SecretStr, model_validator logger = logging.getLogger(__name__) class ChatWriter(BaseChatModel): """Writer chat model. To use, you should have the ``writer-sdk`` Python package installed, and the environment variable ``WRITER_API_KEY`` set with your API key or pass 'api_key' init param. Example: .. code-block:: python from langchain_community.chat_models import ChatWriter chat = ChatWriter( api_key="your key" model="palmyra-x-004" ) """ client: Any = Field(default=None, exclude=True) #: :meta private: async_client: Any = Field(default=None, exclude=True) #: :meta private: api_key: Optional[SecretStr] = Field(default=None) """Writer API key.""" model_name: str = Field(default="palmyra-x-004", alias="model") """Model name to use.""" temperature: float = 0.7 """What sampling temperature to use.""" model_kwargs: Dict[str, Any] = Field(default_factory=dict) """Holds any model parameters valid for `create` call not explicitly specified.""" n: int = 1 """Number of chat completions to generate for each prompt.""" max_tokens: Optional[int] = None """Maximum number of tokens to generate.""" model_config = ConfigDict(populate_by_name=True) @property def _llm_type(self) -> str: """Return type of chat model.""" return "writer-chat" @property def _identifying_params(self) -> Dict[str, Any]: """Get the identifying parameters.""" return { "model_name": self.model_name, "temperature": self.temperature, **self.model_kwargs, } @property def _default_params(self) -> Dict[str, Any]: """Get the default parameters for calling Writer API.""" return { "model": self.model_name, "temperature": self.temperature, "n": self.n, "max_tokens": self.max_tokens, **self.model_kwargs, } @model_validator(mode="before") @classmethod def validate_environment(cls, values: Dict) -> Any: """Validates that api key is passed and creates Writer clients.""" try: from writerai import AsyncClient, Client except ImportError as e: raise ImportError( "Could not import writerai python package. " "Please install it with `pip install writerai`." ) from e if not values.get("client"): values.update( { "client": Client( api_key=get_from_dict_or_env( values, "api_key", "WRITER_API_KEY" ) ) } ) if not values.get("async_client"): values.update( { "async_client": AsyncClient( api_key=get_from_dict_or_env( values, "api_key", "WRITER_API_KEY" ) ) } ) if not ( type(values.get("client")) is Client and type(values.get("async_client")) is AsyncClient ): raise ValueError( "'client' attribute must be with type 'Client' and " "'async_client' must be with type 'AsyncClient' from 'writerai' package" ) return values def _create_chat_result(self, response: Any) -> ChatResult: generations = [] for choice in response.choices: message = self._convert_writer_to_langchain(choice.message) gen = ChatGeneration( message=message, generation_info=dict(finish_reason=choice.finish_reason), ) generations.append(gen) token_usage = {} if response.usage: token_usage = response.usage.__dict__ llm_output = { "token_usage": token_usage, "model_name": self.model_name, "system_fingerprint": response.system_fingerprint, } return ChatResult(generations=generations, llm_output=llm_output) @staticmethod def _convert_langchain_to_writer(message: BaseMessage) -> dict: """Convert a LangChain message to a Writer message dict.""" message_dict = {"role": "", "content": message.content} if isinstance(message, ChatMessage): message_dict["role"] = message.role elif isinstance(message, HumanMessage): message_dict["role"] = "user" elif isinstance(message, AIMessage): message_dict["role"] = "assistant" if message.tool_calls: message_dict["tool_calls"] = [ { "id": tool["id"], "type": "function", "function": {"name": tool["name"], "arguments": tool["args"]}, } for tool in message.tool_calls ] elif isinstance(message, SystemMessage): message_dict["role"] = "system" elif isinstance(message, ToolMessage): message_dict["role"] = "tool" message_dict["tool_call_id"] = message.tool_call_id else: raise ValueError(f"Got unknown message type: {type(message)}") if message.name: message_dict["name"] = message.name return message_dict @staticmethod def _convert_writer_to_langchain(response_message: Any) -> BaseMessage: """Convert a Writer message to a LangChain message.""" if not isinstance(response_message, dict): response_message = json.loads( json.dumps(response_message, default=lambda o: o.__dict__) ) role = response_message.get("role", "") content = response_message.get("content") if not content: content = "" if role == "user": return HumanMessage(content=content) elif role == "assistant": additional_kwargs = {} if tool_calls := response_message.get("tool_calls", []): additional_kwargs["tool_calls"] = tool_calls return AIMessageChunk(content=content, additional_kwargs=additional_kwargs) elif role == "system": return SystemMessage(content=content) elif role == "tool": return ToolMessage( content=content, tool_call_id=response_message.get("tool_call_id", ""), name=response_message.get("name", ""), ) else: return ChatMessage(content=content, role=role) def _convert_messages_to_writer( self, messages: List[BaseMessage], stop: Optional[List[str]] = None ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: """Convert a list of LangChain messages to List of Writer dicts.""" params = { "model": self.model_name, "temperature": self.temperature, "n": self.n, **self.model_kwargs, } if stop: params["stop"] = stop if self.max_tokens is not None: params["max_tokens"] = self.max_tokens message_dicts = [self._convert_langchain_to_writer(m) for m in messages] return message_dicts, params def _stream( self, messages: List[BaseMessage], stop: Optional[List[str]] = None, run_manager: Optional[CallbackManagerForLLMRun] = None, **kwargs: Any, ) -> Iterator[ChatGenerationChunk]: message_dicts, params = self._convert_messages_to_writer(messages, stop) params = {**params, **kwargs, "stream": True} response = self.client.chat.chat(messages=message_dicts, **params) for chunk in response: delta = chunk.choices[0].delta if not delta or not delta.content: continue chunk = self._convert_writer_to_langchain( {"role": "assistant", "content": delta.content} ) chunk = ChatGenerationChunk(message=chunk) if run_manager: run_manager.on_llm_new_token(chunk.text) yield chunk async def _astream( self, messages: List[BaseMessage], stop: Optional[List[str]] = None, run_manager: Optional[AsyncCallbackManagerForLLMRun] = None, **kwargs: Any, ) -> AsyncIterator[ChatGenerationChunk]: message_dicts, params = self._convert_messages_to_writer(messages, stop) params = {**params, **kwargs, "stream": True} response = await self.async_client.chat.chat(messages=message_dicts, **params) async for chunk in response: delta = chunk.choices[0].delta if not delta or not delta.content: continue chunk = self._convert_writer_to_langchain( {"role": "assistant", "content": delta.content} ) chunk = ChatGenerationChunk(message=chunk) if run_manager: await run_manager.on_llm_new_token(chunk.text) yield chunk def _generate( self, messages: List[BaseMessage], stop: Optional[List[str]] = None, run_manager: Optional[CallbackManagerForLLMRun] = None, **kwargs: Any, ) -> ChatResult: message_dicts, params = self._convert_messages_to_writer(messages, stop) params = {**params, **kwargs} response = self.client.chat.chat(messages=message_dicts, **params) return self._create_chat_result(response) async def _agenerate( self, messages: List[BaseMessage], stop: Optional[List[str]] = None, run_manager: Optional[AsyncCallbackManagerForLLMRun] = None, **kwargs: Any, ) -> ChatResult: message_dicts, params = self._convert_messages_to_writer(messages, stop) params = {**params, **kwargs} response = await self.async_client.chat.chat(messages=message_dicts, **params) return self._create_chat_result(response) def bind_tools( self, tools: Sequence[Union[Dict[str, Any], Type[BaseModel], Callable]], *, tool_choice: Optional[Union[str, Literal["auto", "none"]]] = None, **kwargs: Any, ) -> Runnable[LanguageModelInput, BaseMessage]: """Bind tools to the chat model. Args: tools: Tools to bind to the model tool_choice: Which tool to require ('auto', 'none', or specific tool name) **kwargs: Additional parameters to pass to the chat model Returns: A runnable that will use the tools """ formatted_tools = [convert_to_openai_tool(tool) for tool in tools] if tool_choice: kwargs["tool_choice"] = ( (tool_choice) if tool_choice in ("auto", "none") else {"type": "function", "function": {"name": tool_choice}} ) return super().bind(tools=formatted_tools, **kwargs)