codekingpro/portable-devtools

1import json2from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union, cast3 4from langchain_core._api import deprecated5from langchain_core.callbacks import (6    AsyncCallbackManagerForLLMRun,7    CallbackManagerForLLMRun,8)9from langchain_core.language_models.chat_models import BaseChatModel, LangSmithParams10from langchain_core.messages import (11    AIMessage,12    AIMessageChunk,13    BaseMessage,14    ChatMessage,15    HumanMessage,16    SystemMessage,17)18from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult19 20from langchain_community.llms.ollama import OllamaEndpointNotFoundError, _OllamaCommon21 22 23@deprecated("0.0.3", alternative="_chat_stream_response_to_chat_generation_chunk")24def _stream_response_to_chat_generation_chunk(25    stream_response: str,26) -> ChatGenerationChunk:27    """Convert a stream response to a generation chunk."""28    parsed_response = json.loads(stream_response)29    generation_info = parsed_response if parsed_response.get("done") is True else None30    return ChatGenerationChunk(31        message=AIMessageChunk(content=parsed_response.get("response", "")),32        generation_info=generation_info,33    )34 35 36def _chat_stream_response_to_chat_generation_chunk(37    stream_response: str,38) -> ChatGenerationChunk:39    """Convert a stream response to a generation chunk."""40    parsed_response = json.loads(stream_response)41    generation_info = parsed_response if parsed_response.get("done") is True else None42    return ChatGenerationChunk(43        message=AIMessageChunk(44            content=parsed_response.get("message", {}).get("content", "")45        ),46        generation_info=generation_info,47    )48 49 50@deprecated(51    since="0.3.1",52    removal="1.0.0",53    alternative_import="langchain_ollama.ChatOllama",54)55class ChatOllama(BaseChatModel, _OllamaCommon):56    """Ollama locally runs large language models.57 58    To use, follow the instructions at https://ollama.ai/.59 60    Example:61        .. code-block:: python62 63            from langchain_community.chat_models import ChatOllama64            ollama = ChatOllama(model="llama2")65    """66 67    @property68    def _llm_type(self) -> str:69        """Return type of chat model."""70        return "ollama-chat"71 72    @classmethod73    def is_lc_serializable(cls) -> bool:74        """Return whether this model can be serialized by Langchain."""75        return False76 77    def _get_ls_params(78        self, stop: Optional[List[str]] = None, **kwargs: Any79    ) -> LangSmithParams:80        """Get standard params for tracing."""81        params = self._get_invocation_params(stop=stop, **kwargs)82        ls_params = LangSmithParams(83            ls_provider="ollama",84            ls_model_name=self.model,85            ls_model_type="chat",86            ls_temperature=params.get("temperature", self.temperature),87        )88        if ls_max_tokens := params.get("num_predict", self.num_predict):89            ls_params["ls_max_tokens"] = ls_max_tokens90        if ls_stop := stop or params.get("stop", None) or self.stop:91            ls_params["ls_stop"] = ls_stop92        return ls_params93 94    @deprecated("0.0.3", alternative="_convert_messages_to_ollama_messages")95    def _format_message_as_text(self, message: BaseMessage) -> str:96        if isinstance(message, ChatMessage):97            message_text = f"\n\n{message.role.capitalize()}: {message.content}"98        elif isinstance(message, HumanMessage):99            if isinstance(message.content, List):100                first_content = cast(List[Dict], message.content)[0]101                content_type = first_content.get("type")102                if content_type == "text":103                    message_text = f"[INST] {first_content['text']} [/INST]"104                elif content_type == "image_url":105                    message_text = first_content["image_url"]["url"]106            else:107                message_text = f"[INST] {message.content} [/INST]"108        elif isinstance(message, AIMessage):109            message_text = f"{message.content}"110        elif isinstance(message, SystemMessage):111            message_text = f"<<SYS>> {message.content} <</SYS>>"112        else:113            raise ValueError(f"Got unknown type {message}")114        return message_text115 116    def _format_messages_as_text(self, messages: List[BaseMessage]) -> str:117        return "\n".join(118            [self._format_message_as_text(message) for message in messages]119        )120 121    def _convert_messages_to_ollama_messages(122        self, messages: List[BaseMessage]123    ) -> List[Dict[str, Union[str, List[str]]]]:124        ollama_messages: List = []125        for message in messages:126            role = ""127            if isinstance(message, HumanMessage):128                role = "user"129            elif isinstance(message, AIMessage):130                role = "assistant"131            elif isinstance(message, SystemMessage):132                role = "system"133            else:134                raise ValueError("Received unsupported message type for Ollama.")135 136            content = ""137            images = []138            if isinstance(message.content, str):139                content = message.content140            else:141                for content_part in cast(List[Dict], message.content):142                    if content_part.get("type") == "text":143                        content += f"\n{content_part['text']}"144                    elif content_part.get("type") == "image_url":145                        image_url = None146                        temp_image_url = content_part.get("image_url")147                        if isinstance(temp_image_url, str):148                            image_url = content_part["image_url"]149                        elif (150                            isinstance(temp_image_url, dict) and "url" in temp_image_url151                        ):152                            image_url = temp_image_url["url"]153                        else:154                            raise ValueError(155                                "Only string image_url or dict with string 'url' "156                                "inside content parts are supported."157                            )158 159                        image_url_components = image_url.split(",")160                        # Support data:image/jpeg;base64,<image> format161                        # and base64 strings162                        if len(image_url_components) > 1:163                            images.append(image_url_components[1])164                        else:165                            images.append(image_url_components[0])166 167                    else:168                        raise ValueError(169                            "Unsupported message content type. "170                            "Must either have type 'text' or type 'image_url' "171                            "with a string 'image_url' field."172                        )173 174            ollama_messages.append(175                {176                    "role": role,177                    "content": content,178                    "images": images,179                }180            )181 182        return ollama_messages183 184    def _create_chat_stream(185        self,186        messages: List[BaseMessage],187        stop: Optional[List[str]] = None,188        **kwargs: Any,189    ) -> Iterator[str]:190        payload = {191            "model": self.model,192            "messages": self._convert_messages_to_ollama_messages(messages),193        }194        yield from self._create_stream(195            payload=payload, stop=stop, api_url=f"{self.base_url}/api/chat", **kwargs196        )197 198    async def _acreate_chat_stream(199        self,200        messages: List[BaseMessage],201        stop: Optional[List[str]] = None,202        **kwargs: Any,203    ) -> AsyncIterator[str]:204        payload = {205            "model": self.model,206            "messages": self._convert_messages_to_ollama_messages(messages),207        }208        async for stream_resp in self._acreate_stream(209            payload=payload, stop=stop, api_url=f"{self.base_url}/api/chat", **kwargs210        ):211            yield stream_resp212 213    def _chat_stream_with_aggregation(214        self,215        messages: List[BaseMessage],216        stop: Optional[List[str]] = None,217        run_manager: Optional[CallbackManagerForLLMRun] = None,218        verbose: bool = False,219        **kwargs: Any,220    ) -> ChatGenerationChunk:221        final_chunk: Optional[ChatGenerationChunk] = None222        for stream_resp in self._create_chat_stream(messages, stop, **kwargs):223            if stream_resp:224                chunk = _chat_stream_response_to_chat_generation_chunk(stream_resp)225                if final_chunk is None:226                    final_chunk = chunk227                else:228                    final_chunk += chunk229                if run_manager:230                    run_manager.on_llm_new_token(231                        chunk.text,232                        chunk=chunk,233                        verbose=verbose,234                    )235        if final_chunk is None:236            raise ValueError("No data received from Ollama stream.")237 238        return final_chunk239 240    async def _achat_stream_with_aggregation(241        self,242        messages: List[BaseMessage],243        stop: Optional[List[str]] = None,244        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,245        verbose: bool = False,246        **kwargs: Any,247    ) -> ChatGenerationChunk:248        final_chunk: Optional[ChatGenerationChunk] = None249        async for stream_resp in self._acreate_chat_stream(messages, stop, **kwargs):250            if stream_resp:251                chunk = _chat_stream_response_to_chat_generation_chunk(stream_resp)252                if final_chunk is None:253                    final_chunk = chunk254                else:255                    final_chunk += chunk256                if run_manager:257                    await run_manager.on_llm_new_token(258                        chunk.text,259                        chunk=chunk,260                        verbose=verbose,261                    )262        if final_chunk is None:263            raise ValueError("No data received from Ollama stream.")264 265        return final_chunk266 267    def _generate(268        self,269        messages: List[BaseMessage],270        stop: Optional[List[str]] = None,271        run_manager: Optional[CallbackManagerForLLMRun] = None,272        **kwargs: Any,273    ) -> ChatResult:274        """Call out to Ollama's generate endpoint.275 276        Args:277            messages: The list of base messages to pass into the model.278            stop: Optional list of stop words to use when generating.279 280        Returns:281            Chat generations from the model282 283        Example:284            .. code-block:: python285 286                response = ollama([287                    HumanMessage(content="Tell me about the history of AI")288                ])289        """290 291        final_chunk = self._chat_stream_with_aggregation(292            messages,293            stop=stop,294            run_manager=run_manager,295            verbose=self.verbose,296            **kwargs,297        )298        chat_generation = ChatGeneration(299            message=AIMessage(content=final_chunk.text),300            generation_info=final_chunk.generation_info,301        )302        return ChatResult(generations=[chat_generation])303 304    async def _agenerate(305        self,306        messages: List[BaseMessage],307        stop: Optional[List[str]] = None,308        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,309        **kwargs: Any,310    ) -> ChatResult:311        """Call out to Ollama's generate endpoint.312 313        Args:314            messages: The list of base messages to pass into the model.315            stop: Optional list of stop words to use when generating.316 317        Returns:318            Chat generations from the model319 320        Example:321            .. code-block:: python322 323                response = ollama([324                    HumanMessage(content="Tell me about the history of AI")325                ])326        """327 328        final_chunk = await self._achat_stream_with_aggregation(329            messages,330            stop=stop,331            run_manager=run_manager,332            verbose=self.verbose,333            **kwargs,334        )335        chat_generation = ChatGeneration(336            message=AIMessage(content=final_chunk.text),337            generation_info=final_chunk.generation_info,338        )339        return ChatResult(generations=[chat_generation])340 341    def _stream(342        self,343        messages: List[BaseMessage],344        stop: Optional[List[str]] = None,345        run_manager: Optional[CallbackManagerForLLMRun] = None,346        **kwargs: Any,347    ) -> Iterator[ChatGenerationChunk]:348        try:349            for stream_resp in self._create_chat_stream(messages, stop, **kwargs):350                if stream_resp:351                    chunk = _chat_stream_response_to_chat_generation_chunk(stream_resp)352                    if run_manager:353                        run_manager.on_llm_new_token(354                            chunk.text,355                            chunk=chunk,356                            verbose=self.verbose,357                        )358                    yield chunk359        except OllamaEndpointNotFoundError:360            yield from self._legacy_stream(messages, stop, **kwargs)361 362    async def _astream(363        self,364        messages: List[BaseMessage],365        stop: Optional[List[str]] = None,366        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,367        **kwargs: Any,368    ) -> AsyncIterator[ChatGenerationChunk]:369        async for stream_resp in self._acreate_chat_stream(messages, stop, **kwargs):370            if stream_resp:371                chunk = _chat_stream_response_to_chat_generation_chunk(stream_resp)372                if run_manager:373                    await run_manager.on_llm_new_token(374                        chunk.text,375                        chunk=chunk,376                        verbose=self.verbose,377                    )378                yield chunk379 380    @deprecated("0.0.3", alternative="_stream")381    def _legacy_stream(382        self,383        messages: List[BaseMessage],384        stop: Optional[List[str]] = None,385        run_manager: Optional[CallbackManagerForLLMRun] = None,386        **kwargs: Any,387    ) -> Iterator[ChatGenerationChunk]:388        prompt = self._format_messages_as_text(messages)389        for stream_resp in self._create_generate_stream(prompt, stop, **kwargs):390            if stream_resp:391                chunk = _stream_response_to_chat_generation_chunk(stream_resp)392                if run_manager:393                    run_manager.on_llm_new_token(394                        chunk.text,395                        chunk=chunk,396                        verbose=self.verbose,397                    )398                yield chunk399