From 966611bbfa53e46404e64ac71c8a38716e939415 Mon Sep 17 00:00:00 2001
From: Harrison Chase <hw.chase.17@gmail.com>
Date: Sat, 28 Jan 2023 08:24:55 -0800
Subject: [PATCH] add model kwargs to handle stop token from cohere (#773)

---
 langchain/llms/cohere.py | 23 +++++++++++++++++------
 1 file changed, 17 insertions(+), 6 deletions(-)

diff --git a/langchain/llms/cohere.py b/langchain/llms/cohere.py
index d9a3f51a03..fd7b91067d 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.