Cohere Client (microsoft#3004)

* initial setup for cohere client * client update * changes: ClintType added to the utils * Revert "changes: ClintType added to the utils" This reverts commit 80d6155. * Message conversion to Cohere, Parameter handling, cost calculation, streaming, tool calling * Changed Groq references. * minor fix * tests added * ref fix * added in the workflows * Fixed bug on non-streaming text generation * fix: formatting * Support Cohere rule for last message not USER when tool_results exist * Added Cohere to documentation * fixed client.py merge, removed unnecessary comments in groq.py, updated Cohere documentation, added Groq documentation * log: ignored params * update: custom exception added --------- Co-authored-by: Mark Sze <[email protected]> Co-authored-by: Mark Sze <[email protected]>
pyscht · Jul 3, 2024 · b4a3f26 · b4a3f26
1 parent 8133b7d
commit b4a3f26
Show file tree

Hide file tree

Showing 11 changed files with 1,661 additions and 10 deletions.
diff --git a/.github/workflows/contrib-tests.yml b/.github/workflows/contrib-tests.yml
@@ -638,3 +638,39 @@ jobs:
         with:
           file: ./coverage.xml
           flags: unittests
+
+  CohereTest:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ubuntu-latest, macos-latest, windows-latest]
+        python-version: ["3.9", "3.10", "3.11", "3.12"]
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          lfs: true
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install packages and dependencies for all tests
+        run: |
+          python -m pip install --upgrade pip wheel
+          pip install pytest-cov>=5
+      - name: Install packages and dependencies for Cohere
+        run: |
+          pip install -e .[cohere,test]
+      - name: Set AUTOGEN_USE_DOCKER based on OS
+        shell: bash
+        run: |
+          if [[ ${{ matrix.os }} != ubuntu-latest ]]; then
+            echo "AUTOGEN_USE_DOCKER=False" >> $GITHUB_ENV
+          fi
+      - name: Coverage
+        run: |
+          pytest test/oai/test_cohere.py --skip-openai
+      - name: Upload coverage to Codecov
+        uses: codecov/codecov-action@v3
+        with:
+          file: ./coverage.xml
+          flags: unittests
diff --git a/autogen/logger/file_logger.py b/autogen/logger/file_logger.py
@@ -18,6 +18,7 @@
 if TYPE_CHECKING:
     from autogen import Agent, ConversableAgent, OpenAIWrapper
     from autogen.oai.anthropic import AnthropicClient
+    from autogen.oai.cohere import CohereClient
     from autogen.oai.gemini import GeminiClient
     from autogen.oai.groq import GroqClient
     from autogen.oai.mistral import MistralAIClient
@@ -205,7 +206,16 @@ def log_new_wrapper(
 
     def log_new_client(
         self,
-        client: AzureOpenAI | OpenAI | GeminiClient | AnthropicClient | MistralAIClient | TogetherClient | GroqClient,
+        client: (
+            AzureOpenAI
+            | OpenAI
+            | GeminiClient
+            | AnthropicClient
+            | MistralAIClient
+            | TogetherClient
+            | GroqClient
+            | CohereClient
+        ),
         wrapper: OpenAIWrapper,
         init_args: Dict[str, Any],
     ) -> None:

diff --git a/autogen/logger/sqlite_logger.py b/autogen/logger/sqlite_logger.py
@@ -19,6 +19,7 @@
 if TYPE_CHECKING:
     from autogen import Agent, ConversableAgent, OpenAIWrapper
     from autogen.oai.anthropic import AnthropicClient
+    from autogen.oai.cohere import CohereClient
     from autogen.oai.gemini import GeminiClient
     from autogen.oai.groq import GroqClient
     from autogen.oai.mistral import MistralAIClient
@@ -392,7 +393,16 @@ def log_function_use(self, source: Union[str, Agent], function: F, args: Dict[st
 
     def log_new_client(
         self,
-        client: Union[AzureOpenAI, OpenAI, GeminiClient, AnthropicClient, MistralAIClient, TogetherClient, GroqClient],
+        client: Union[
+            AzureOpenAI,
+            OpenAI,
+            GeminiClient,
+            AnthropicClient,
+            MistralAIClient,
+            TogetherClient,
+            GroqClient,
+            CohereClient,
+        ],
         wrapper: OpenAIWrapper,
         init_args: Dict[str, Any],
     ) -> None:

diff --git a/autogen/oai/client.py b/autogen/oai/client.py
@@ -77,6 +77,13 @@
 except ImportError as e:
     groq_import_exception = e
 
+try:
+    from autogen.oai.cohere import CohereClient
+
+    cohere_import_exception: Optional[ImportError] = None
+except ImportError as e:
+    cohere_import_exception = e
+
 logger = logging.getLogger(__name__)
 if not logger.handlers:
     # Add the console handler.
@@ -497,6 +504,11 @@ def _register_default_client(self, config: Dict[str, Any], openai_config: Dict[s
                     raise ImportError("Please install `groq` to use the Groq API.")
                 client = GroqClient(**openai_config)
                 self._clients.append(client)
+            elif api_type is not None and api_type.startswith("cohere"):
+                if cohere_import_exception:
+                    raise ImportError("Please install `cohere` to use the Groq API.")
+                client = CohereClient(**openai_config)
+                self._clients.append(client)
             else:
                 client = OpenAI(**openai_config)
                 self._clients.append(OpenAIClient(client))