MlflowEmbeddings: input_type argument is missing, required by Cohere embeddings models
Labels: name: area: embeddings, name: auto:bug

### System Info

python = "3.11"
langchain = "0.0.352"
cohere = "4.39"
mlflow = {extras = ["genai"], version = "2.9.2"}

### Who can help?

@harupy
@hwchase17
@agola11

### Information

- [ ] The official example notebooks/scripts
- [X] My own modified scripts

### Related Components

- [ ] LLMs/Chat Models
- [X] Embedding Models
- [ ] Prompts / Prompt Templates / Prompt Selectors
- [ ] Output Parsers
- [ ] Document Loaders
- [ ] Vector Stores / Retrievers
- [ ] Memory
- [ ] Agents / Agent Executors
- [ ] Tools / Toolkits
- [ ] Chains
- [ ] Callbacks/Tracing
- [ ] Async

### Reproduction

I followed the official example for embeddings here, except that I am using Cohere instead of OpenAI: https://python.langchain.com/docs/integrations/providers/mlflow

<details>
<summary>Click for specific steps</summary>

More specifically, I installed mlflow genai and set my `COHERE_API_KEY` environment variable:
```bash
pip install 'mlflow[genai]'
export COHERE_API_KEY=...
```

I created `config.yaml` like so:
```yaml
endpoints:
  - name: completions
    endpoint_type: llm/v1/completions
    model:
      provider: cohere
      name: command
      config:
        cohere_api_key: $COHERE_API_KEY


  - name: embeddings
    endpoint_type: llm/v1/embeddings
    model:
      provider: cohere
      name: embed-english-light-v3.0
      config:
        cohere_api_key: $COHERE_API_KEY
```

I started the mlflow deployments server:
```bash
mlflow deployments start-server --config-path config.yaml
```

<details>
<summary>The server started as expected</summary>

```
xxx/python3.11/site-packages/pydantic/_internal/_config.py:321: UserWarning: Valid config keys have changed in V2:
* 'schema_extra' has been renamed to 'json_schema_extra'
  warnings.warn(message, UserWarning)
[2023-12-27 13:53:18 -0800] [22480] [INFO] Starting gunicorn 21.2.0
[2023-12-27 13:53:18 -0800] [22480] [INFO] Listening at: http://127.0.0.1:5000 (22480)
[2023-12-27 13:53:18 -0800] [22480] [INFO] Using worker: uvicorn.workers.UvicornWorker
[2023-12-27 13:53:18 -0800] [22481] [INFO] Booting worker with pid: 22481
[2023-12-27 13:53:18 -0800] [22482] [INFO] Booting worker with pid: 22482
xxx/python3.11/site-packages/pydantic/_internal/_config.py:321: UserWarning: Valid config keys have changed in V2:
* 'schema_extra' has been renamed to 'json_schema_extra'
  warnings.warn(message, UserWarning)
xxx/python3.11/site-packages/pydantic/_internal/_config.py:321: UserWarning: Valid config keys have changed in V2:
* 'schema_extra' has been renamed to 'json_schema_extra'
  warnings.warn(message, UserWarning)
[2023-12-27 13:53:20 -0800] [22481] [INFO] Started server process [22481]
[2023-12-27 13:53:20 -0800] [22481] [INFO] Waiting for application startup.
[2023-12-27 13:53:20 -0800] [22481] [INFO] Application startup complete.
[2023-12-27 13:53:20 -0800] [22482] [INFO] Started server process [22482]
[2023-12-27 13:53:20 -0800] [22482] [INFO] Waiting for application startup.
[2023-12-27 13:53:20 -0800] [22482] [INFO] Application startup complete.
```
</details>

In `test.py`, I added the embeddings example:

```python
from langchain.embeddings import MlflowEmbeddings

embeddings = MlflowEmbeddings(
    target_uri="http://127.0.0.1:5000",
    endpoint="embeddings",
)

print(embeddings.embed_query("hello"))
print(embeddings.embed_documents(["hello"]))
```

And I ran it with `python test.py`.
</details>

Here is the error I got:
```
raise HTTPError(
requests.exceptions.HTTPError: 400 Client Error: Bad Request for url: http://127.0.0.1:5000/endpoints/embeddings/invocations. Response text: {"detail":{"message":"invalid request: valid input_type must be provided with the provided model"}}
```

<details>
<summary>Full trace</summary>

```
xxx/python3.11/site-packages/pydantic/_internal/_config.py:321: UserWarning: Valid config keys have changed in V2:
* 'schema_extra' has been renamed to 'json_schema_extra'
  warnings.warn(message, UserWarning)
Traceback (most recent call last):
  File "xxx/python3.11/site-packages/mlflow/utils/request_utils.py", line 52, in augmented_raise_for_status
    response.raise_for_status()
  File "xxx/python3.11/site-packages/requests/models.py", line 1021, in raise_for_status
    raise HTTPError(http_error_msg, response=self)
requests.exceptions.HTTPError: 400 Client Error: Bad Request for url: http://127.0.0.1:5000/endpoints/embeddings/invocations

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "yyy/test.py", line 8, in <module>
    print(embeddings.embed_query("hello"))
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "xxx/python3.11/site-packages/langchain_community/embeddings/mlflow.py", line 74, in embed_query
    return self.embed_documents([text])[0]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "xxx/python3.11/site-packages/langchain_community/embeddings/mlflow.py", line 69, in embed_documents
    resp = self._client.predict(endpoint=self.endpoint, inputs={"input": txt})
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "xxx/python3.11/site-packages/mlflow/deployments/mlflow/__init__.py", line 294, in predict
    return self._call_endpoint(
           ^^^^^^^^^^^^^^^^^^^^
  File "xxx/python3.11/site-packages/mlflow/deployments/mlflow/__init__.py", line 139, in _call_endpoint
    augmented_raise_for_status(response)
  File "xxx/python3.11/site-packages/mlflow/utils/request_utils.py", line 55, in augmented_raise_for_status
    raise HTTPError(
requests.exceptions.HTTPError: 400 Client Error: Bad Request for url: http://127.0.0.1:5000/endpoints/embeddings/invocations. Response text: {"detail":{"message":"invalid request: valid input_type must be provided with the provided model"}}
```
</details>

The issue is that `embed_query`/`embed_documents` don't allow passing in the input_type argument, which is needed by the Cohere API -- see https://github.com/langchain-ai/langchain/blob/master/libs/community/langchain_community/embeddings/cohere.py#L81

### Proposed solution

My quick solution was to modify the two methods in `MlflowEmbeddings` to allow for kwargs:

```python
class MlflowEmbeddings:
    ....

    def embed_documents(self, texts: List[str], **kwargs) -> List[List[float]]:
        embeddings: List[List[float]] = []
        for txt in _chunk(texts, 20):
            resp = self._client.predict(endpoint=self.endpoint, inputs={"input": txt, **kwargs})
            embeddings.extend(r["embedding"] for r in resp["data"])
        return embeddings

    def embed_query(self, text: str, **kwargs) -> List[float]:
        return self.embed_documents([text], **kwargs)[0]
```

So `test.py` changes to:
```python
print(embeddings.embed_query("hello", input_type="search_query"))
print(embeddings.embed_documents(["hello"], input_type="search_document"))
```

This might not be the best solution since it kind of defeats the purpose of separating `embed_query` and `embed_documents` for Cohere. Another solution is to subclass MlflowEmbeddings for Cohere (and others?).

I intend to open a PR with this change, so any guidance on the best approach is much appreciated!

### Expected behavior

The code should generate embeddings for the given words