diff --git a/langchain/llms/cohere.py b/langchain/llms/cohere.py index d9a3f51a..fd7b9106 100644 --- a/langchain/llms/cohere.py +++ b/langchain/llms/cohere.py @@ -1,5 +1,6 @@ """Wrapper around Cohere APIs.""" -from typing import Any, Dict, List, Mapping, Optional +import logging +from typing import Any, Dict, List, Optional from pydantic import BaseModel, Extra, root_validator @@ -7,6 +8,8 @@ from langchain.llms.base import LLM from langchain.llms.utils import enforce_stop_tokens from langchain.utils import get_from_dict_or_env +logger = logging.getLogger(__name__) + class Cohere(LLM, BaseModel): """Wrapper around Cohere large language models. @@ -46,6 +49,8 @@ class Cohere(LLM, BaseModel): cohere_api_key: Optional[str] = None + stop: Optional[List[str]] = None + class Config: """Configuration for this pydantic object.""" @@ -69,7 +74,7 @@ class Cohere(LLM, BaseModel): return values @property - def _default_params(self) -> Mapping[str, Any]: + def _default_params(self) -> Dict[str, Any]: """Get the default parameters for calling Cohere API.""" return { "max_tokens": self.max_tokens, @@ -81,7 +86,7 @@ class Cohere(LLM, BaseModel): } @property - def _identifying_params(self) -> Mapping[str, Any]: + def _identifying_params(self) -> Dict[str, Any]: """Get the identifying parameters.""" return {**{"model": self.model}, **self._default_params} @@ -105,9 +110,15 @@ class Cohere(LLM, BaseModel): response = cohere("Tell me a joke.") """ - response = self.client.generate( - model=self.model, prompt=prompt, stop_sequences=stop, **self._default_params - ) + params = self._default_params + if self.stop is not None and stop is not None: + raise ValueError("`stop` found in both the input and default params.") + elif self.stop is not None: + params["stop_sequences"] = self.stop + else: + params["stop_sequences"] = stop + + response = self.client.generate(model=self.model, prompt=prompt, **params) text = response.generations[0].text # If stop tokens are provided, Cohere's endpoint returns them. # In order to make this consistent with other endpoints, we strip them.