Jasonozuzu/pts 7374 langchain chat v2 (stream added)

libs/cohere/langchain_cohere/cohere_agent.py

@@ -3,6 +3,8 @@
 
 from cohere.types import (
     Tool,
+    ToolV2,
+    ToolV2Function,
     ToolCall,
     ToolParameterDefinitionsValue,
     ToolResult,
@@ -69,6 +71,12 @@ def _format_to_cohere_tools(
     return [_convert_to_cohere_tool(tool) for tool in tools]
 
 
+def _format_to_cohere_tools_v2(
+    tools: Sequence[Union[Dict[str, Any], Type[BaseModel], Callable, BaseTool]],
+) -> List[Dict[str, Any]]:
+    return [_convert_to_cohere_tool_v2(tool) for tool in tools]
+
+
 def _format_to_cohere_tools_messages(
     intermediate_steps: Sequence[Tuple[AgentAction, str]],
 ) -> List[Dict[str, Any]]:
@@ -176,6 +184,97 @@ def _convert_to_cohere_tool(
         )
 
 
+def _convert_to_cohere_tool_v2(
+    tool: Union[Union[Dict[str, Any], Type[BaseModel], Callable, BaseTool]],
+) -> Dict[str, Any]:
+    """
+    Convert a BaseTool instance, JSON schema dict, or BaseModel type to a V2 Cohere tool.
+    """
+    if isinstance(tool, dict):
+        if not all(k in tool for k in ("title", "description", "properties")):
+            raise ValueError(
+                "Unsupported dict type. Tool must be passed in as a BaseTool instance, JSON schema dict, or BaseModel type."  # noqa: E501
+            )
+        return ToolV2(
+            type="function",
+            function=ToolV2Function(
+                name=tool.get("title"),
+                description=tool.get("description"),
+                parameters={
+                    "type": "object",
+                    "properties": {
+                        param_name: {
+                            "description": param_definition.get("description"),
+                            "type": JSON_TO_PYTHON_TYPES.get(
+                                param_definition.get("type"), param_definition.get("type")
+                            ),
+                        }
+                        for param_name, param_definition in tool.get("properties", {}).items()
+                    },
+                    "required": [param_name 
+                                 for param_name, param_definition 
+                                 in tool.get("properties", {}).items() 
+                                 if "default" not in param_definition],
+                },
+            )
+        ).dict()
+    elif (
+        (isinstance(tool, type) and issubclass(tool, BaseModel))
+        or callable(tool)
+        or isinstance(tool, BaseTool)
+    ):
+        as_json_schema_function = convert_to_openai_function(tool)
+        parameters = as_json_schema_function.get("parameters", {})
+        properties = parameters.get("properties", {})
+        parameter_definitions = {}
+        required_params = []
+        for param_name, param_definition in properties.items():
+            if "type" in param_definition:
+                _type_str = param_definition.get("type")
+                _type = JSON_TO_PYTHON_TYPES.get(_type_str)
+            elif "anyOf" in param_definition:
+                _type_str = next(
+                    (
+                        t.get("type")
+                        for t in param_definition.get("anyOf", [])
+                        if t.get("type") != "null"
+                    ),
+                    param_definition.get("type"),
+                )
+                _type = JSON_TO_PYTHON_TYPES.get(_type_str)
+            else:
+                _type = None
+            tool_definition = {
+                "type": _type,
+                "description": param_definition.get("description"),
+            }
+            parameter_definitions[param_name] = tool_definition
+            if param_name in parameters.get("required", []):
+                required_params.append(param_name)
+        return ToolV2(
+            type="function",
+            function=ToolV2Function(
+                name=as_json_schema_function.get("name"),
+                description=as_json_schema_function.get(
+                    # The Cohere API requires the description field.
+                    "description",
+                    as_json_schema_function.get("name"),
+                ),
+                parameters={
+                    "type": "object",
+                    "properties": {
+                        **parameter_definitions,
+                    },
+                    "required": required_params,
+                },
+            )
+        ).dict()
+    else:
+        raise ValueError(
+            f"Unsupported tool type {type(tool)}. Tool must be passed in as a BaseTool instance, JSON schema dict, or BaseModel type."  # noqa: E501
+        )
+
+
 class _CohereToolsAgentOutputParser(
     BaseOutputParser[Union[List[AgentAction], AgentFinish]]
 ):

libs/cohere/langchain_cohere/llms.py

@@ -2,7 +2,7 @@
 
 import logging
 import re
-from typing import Any, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 import cohere
 from langchain_core.callbacks import (
@@ -30,6 +30,9 @@ def enforce_stop_tokens(text: str, stop: List[str]) -> str:
 
 logger = logging.getLogger(__name__)
 
+if TYPE_CHECKING:
+    from cohere.types import ListModelsResponse  # noqa: F401
+
 
 def completion_with_retry(llm: Cohere, **kwargs: Any) -> Any:
     """Use tenacity to retry the completion call."""
@@ -58,6 +61,19 @@ class BaseCohere(Serializable):
 
     client: Any = None  #: :meta private:
     async_client: Any = None  #: :meta private:
+
+    chat_v2: Optional[Any] = None
+    "Cohere chat v2."
+
+    async_chat_v2: Optional[Any] = None
+    "Cohere async chat v2."
+
+    chat_stream_v2: Optional[Any] = None
+    "Cohere chat stream v2."
+
+    async_chat_stream_v2: Optional[Any] = None
+    "Cohere async chat stream v2."
+
     model: Optional[str] = Field(default=None)
     """Model name to use."""
 
@@ -83,6 +99,15 @@ class BaseCohere(Serializable):
     base_url: Optional[str] = None
     """Override the default Cohere API URL."""
 
+    def _get_default_model(self) -> str:
+        """Fetches the current default model name."""
+        response = self.client.models.list(default_only=True, endpoint="chat")  # type: "ListModelsResponse"
+        if not response.models:
+            raise Exception("invalid cohere list models response")
+        if not response.models[0].name:
+            raise Exception("invalid cohere list models response")
+        return response.models[0].name
+
     @model_validator(mode="after")
     def validate_environment(self) -> Self:  # type: ignore[valid-type]
         """Validate that api key and python package exists in environment."""
@@ -98,12 +123,21 @@ def validate_environment(self) -> Self:  # type: ignore[valid-type]
             client_name=client_name,
             base_url=self.base_url,
         )
+        self.chat_v2 = self.client.v2.chat
+        self.chat_stream_v2 = self.client.v2.chat_stream
+
         self.async_client = cohere.AsyncClient(
             api_key=cohere_api_key,
             client_name=client_name,
             timeout=timeout_seconds,
             base_url=self.base_url,
         )
+        self.async_chat_v2 = self.async_client.v2.chat
+        self.async_chat_stream_v2 = self.async_client.v2.chat_stream
+
+        if not self.model:
+            self.model = self._get_default_model()
+
         return self
 
 

libs/cohere/langchain_cohere/rag_retrievers.py

@@ -33,8 +33,6 @@ def _get_docs(response: Any) -> List[Document]:
             metadata={
                 "type": "model_response",
                 "citations": response.generation_info["citations"],
-                "search_results": response.generation_info["search_results"],
-                "search_queries": response.generation_info["search_queries"],
                 "token_count": response.generation_info["token_count"],
             },
         )

libs/cohere/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langchain-cohere"
-version = "0.3.1"
+version = "1.0.0-rc1"
 description = "An integration package connecting Cohere and LangChain"
 authors = []
 readme = "README.md"

libs/cohere/tests/clear_cassettes.py

@@ -0,0 +1,18 @@
+import os
+import shutil
+
+def delete_cassettes_directories(root_dir):
+    for dirpath, dirnames, filenames in os.walk(root_dir):
+        for dirname in dirnames:
+            if dirname == "cassettes":
+                dir_to_delete = os.path.join(dirpath, dirname)
+                print(f"Deleting directory: {dir_to_delete}")
+                shutil.rmtree(dir_to_delete)
+
+if __name__ == "__main__":
+    # Clear all cassettes directories in the integration_tests directory
+    # run using: python clear_cassettes.py
+    directory_to_clear = os.path.join(os.getcwd(), "integration_tests")
+    if not os.path.isdir(directory_to_clear):
+        raise Exception("integration_tests directory not found in current working directory")
+    delete_cassettes_directories(directory_to_clear)

libs/cohere/tests/conftest.py

@@ -1,6 +1,8 @@
-from typing import Dict
+from typing import Dict, Generator, Optional
 
 import pytest
+from unittest.mock import MagicMock, patch
+from langchain_cohere.llms import BaseCohere
 
 
 @pytest.fixture(scope="module")
@@ -10,3 +12,10 @@ def vcr_config() -> Dict:
         "filter_headers": [("Authorization", None)],
         "ignore_hosts": ["storage.googleapis.com"],
     }
+
+@pytest.fixture
+def patch_base_cohere_get_default_model() -> Generator[Optional[MagicMock], None, None]:
+    with patch.object(
+        BaseCohere, "_get_default_model", return_value="command-r-plus"
+    ) as mock_get_default_model:
+        yield mock_get_default_model