243 lines
8.5 KiB
Python
243 lines
8.5 KiB
Python
"""Cohere-specific utilities.
|
|
|
|
This module contains utilities specific to the Cohere provider,
|
|
including reask functions, response handlers, and message formatting.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from ...mode import Mode
|
|
|
|
|
|
def reask_cohere_tools(
|
|
kwargs: dict[str, Any],
|
|
response: Any, # Replace with actual response type for Cohere
|
|
exception: Exception,
|
|
):
|
|
"""
|
|
Handle reask for Cohere tools and JSON schema modes.
|
|
Supports both V1 and V2 formats.
|
|
|
|
V1 kwargs modifications:
|
|
- Adds/Modifies: "chat_history" (appends prior message)
|
|
- Modifies: "message" (user prompt describing validation errors)
|
|
|
|
V2 kwargs modifications:
|
|
- Modifies: "messages" (appends error correction message)
|
|
"""
|
|
# Default to marker stored on kwargs (set during client initialization)
|
|
client_version = kwargs.get("_cohere_client_version")
|
|
|
|
# Detect V1 vs V2 response structure and extract text
|
|
if hasattr(response, "text"):
|
|
client_version = "v1"
|
|
response_text = response.text
|
|
elif hasattr(response, "message") and hasattr(response.message, "content"):
|
|
client_version = "v2"
|
|
content_items = response.message.content
|
|
response_text = ""
|
|
if content_items:
|
|
# Find the text content item (skip thinking/other types)
|
|
for item in content_items:
|
|
if (
|
|
hasattr(item, "type")
|
|
and item.type == "text"
|
|
and hasattr(item, "text")
|
|
):
|
|
response_text = item.text
|
|
break
|
|
if not response_text:
|
|
response_text = str(response)
|
|
else:
|
|
# Fallback to string representation
|
|
response_text = str(response)
|
|
if client_version is None:
|
|
if "messages" in kwargs:
|
|
client_version = "v2"
|
|
elif "chat_history" in kwargs or "message" in kwargs:
|
|
client_version = "v1"
|
|
|
|
# Create the correction message
|
|
correction_msg = (
|
|
"Correct the following JSON response, based on the errors given below:\n\n"
|
|
f"JSON:\n{response_text}\n\nExceptions:\n{exception}"
|
|
)
|
|
|
|
if client_version == "v2":
|
|
# V2 format: append to messages list
|
|
kwargs["messages"].append({"role": "user", "content": correction_msg})
|
|
elif client_version == "v1":
|
|
# V1 format: use chat_history and message
|
|
message = kwargs.get("message", "")
|
|
|
|
# Fetch or initialize chat_history in one operation
|
|
if "chat_history" in kwargs:
|
|
kwargs["chat_history"].append({"role": "user", "message": message})
|
|
else:
|
|
kwargs["chat_history"] = [{"role": "user", "message": message}]
|
|
|
|
kwargs["message"] = correction_msg
|
|
else:
|
|
# Unknown version - raise error for future compatibility
|
|
raise ValueError(
|
|
f"Unsupported Cohere client version: {client_version}. "
|
|
f"Expected 'v1' or 'v2'."
|
|
)
|
|
|
|
return kwargs
|
|
|
|
|
|
def handle_cohere_modes(new_kwargs: dict[str, Any]) -> tuple[None, dict[str, Any]]:
|
|
"""
|
|
Convert OpenAI-style messages to Cohere format.
|
|
Handles both V1 and V2 client formats.
|
|
|
|
V1 format:
|
|
- Removes: "messages"
|
|
- Adds: "message" (last user message)
|
|
- Adds: "chat_history" (prior messages)
|
|
|
|
V2 format:
|
|
- Keeps: "messages" (compatible with OpenAI format)
|
|
|
|
Both versions:
|
|
- Renames: "model_name" -> "model"
|
|
- Removes: "strict"
|
|
- Removes: "_cohere_client_version" (internal marker)
|
|
"""
|
|
new_kwargs = new_kwargs.copy()
|
|
client_version = new_kwargs.pop("_cohere_client_version")
|
|
|
|
if client_version == "v2":
|
|
# V2 uses OpenAI-style messages directly - no conversion needed
|
|
# Just clean up incompatible fields
|
|
if "model_name" in new_kwargs and "model" not in new_kwargs:
|
|
new_kwargs["model"] = new_kwargs.pop("model_name")
|
|
new_kwargs.pop("strict", None)
|
|
elif client_version == "v1":
|
|
# V1 needs conversion from OpenAI format to Cohere V1 format
|
|
messages = new_kwargs.pop("messages", [])
|
|
chat_history = []
|
|
for message in messages[:-1]:
|
|
chat_history.append( # type: ignore[arg-type]
|
|
{
|
|
"role": message["role"],
|
|
"message": message["content"],
|
|
}
|
|
)
|
|
new_kwargs["message"] = messages[-1]["content"]
|
|
new_kwargs["chat_history"] = chat_history
|
|
if "model_name" in new_kwargs and "model" not in new_kwargs:
|
|
new_kwargs["model"] = new_kwargs.pop("model_name")
|
|
new_kwargs.pop("strict", None)
|
|
else:
|
|
# Unknown version - raise error for future compatibility
|
|
raise ValueError(
|
|
f"Unsupported Cohere client version: {client_version}. "
|
|
f"Expected 'v1' or 'v2'."
|
|
)
|
|
|
|
return None, new_kwargs
|
|
|
|
|
|
def handle_cohere_json_schema(
|
|
response_model: type[Any] | None, new_kwargs: dict[str, Any]
|
|
) -> tuple[type[Any] | None, dict[str, Any]]:
|
|
"""
|
|
Handle Cohere JSON schema mode.
|
|
|
|
When response_model is None:
|
|
- Converts messages from OpenAI format to Cohere format (message + chat_history)
|
|
- No schema is added to the request
|
|
|
|
When response_model is provided:
|
|
- Converts messages from OpenAI format to Cohere format
|
|
- Adds the model's JSON schema to response_format
|
|
|
|
Kwargs modifications:
|
|
- Removes: "messages" (converted to message + chat_history)
|
|
- Adds: "message" (last message content)
|
|
- Adds: "chat_history" (all messages except last)
|
|
- Modifies: "model" (if "model_name" exists, renames to "model")
|
|
- Removes: "strict"
|
|
- Adds: "response_format" (with JSON schema) - only when response_model provided
|
|
"""
|
|
if response_model is None:
|
|
# Just handle message conversion
|
|
return handle_cohere_modes(new_kwargs)
|
|
|
|
new_kwargs["response_format"] = {
|
|
"type": "json_object",
|
|
"schema": response_model.model_json_schema(),
|
|
}
|
|
_, new_kwargs = handle_cohere_modes(new_kwargs)
|
|
|
|
return response_model, new_kwargs
|
|
|
|
|
|
def handle_cohere_tools(
|
|
response_model: type[Any] | None, new_kwargs: dict[str, Any]
|
|
) -> tuple[type[Any] | None, dict[str, Any]]:
|
|
"""
|
|
Handle Cohere tools mode.
|
|
|
|
When response_model is None:
|
|
- Converts messages from OpenAI format to Cohere format (message + chat_history for V1, messages for V2)
|
|
- No tools or schema instructions are added
|
|
- Allows for unstructured responses from Cohere
|
|
|
|
When response_model is provided:
|
|
- Converts messages from OpenAI format to Cohere format
|
|
- Prepends extraction instructions to the chat history (V1) or messages (V2)
|
|
- Includes the model's JSON schema in the instructions
|
|
- The model is instructed to extract a valid object matching the schema
|
|
|
|
Kwargs modifications:
|
|
- All modifications from handle_cohere_modes (message format conversion)
|
|
- Modifies: "chat_history" (V1) or "messages" (V2) to prepend extraction instruction - only when response_model provided
|
|
"""
|
|
if response_model is None:
|
|
# Just handle message conversion
|
|
return handle_cohere_modes(new_kwargs)
|
|
|
|
_, new_kwargs = handle_cohere_modes(new_kwargs)
|
|
|
|
instruction = f"""\
|
|
Extract a valid {response_model.__name__} object based on the chat history and the json schema below.
|
|
{response_model.model_json_schema()}
|
|
The JSON schema was obtained by running:
|
|
```python
|
|
schema = {response_model.__name__}.model_json_schema()
|
|
```
|
|
|
|
The output must be a valid JSON object that `{response_model.__name__}.model_validate_json()` can successfully parse.
|
|
Respond with JSON only. Do not include code fences, markdown, or extra text.
|
|
"""
|
|
# Check client version explicitly (marker already removed by handle_cohere_modes)
|
|
# Use presence of messages vs chat_history as indicator since marker is already consumed
|
|
if "messages" in new_kwargs:
|
|
# V2 format: prepend to messages
|
|
new_kwargs["messages"].insert(0, {"role": "user", "content": instruction})
|
|
else:
|
|
# V1 format: prepend to chat_history
|
|
new_kwargs["chat_history"] = [
|
|
{"role": "user", "message": instruction}
|
|
] + new_kwargs["chat_history"]
|
|
|
|
return response_model, new_kwargs
|
|
|
|
|
|
# Handler registry for Cohere
|
|
COHERE_HANDLERS = {
|
|
Mode.COHERE_TOOLS: {
|
|
"reask": reask_cohere_tools,
|
|
"response": handle_cohere_tools,
|
|
},
|
|
Mode.COHERE_JSON_SCHEMA: {
|
|
"reask": reask_cohere_tools,
|
|
"response": handle_cohere_json_schema,
|
|
},
|
|
}
|