682 lines
29 KiB
Python
682 lines
29 KiB
Python
from collections.abc import AsyncGenerator, Generator, Iterable
|
|
from typing import (
|
|
Any,
|
|
ClassVar,
|
|
Optional,
|
|
cast,
|
|
get_origin,
|
|
get_args,
|
|
Union,
|
|
TYPE_CHECKING,
|
|
)
|
|
import json
|
|
from pydantic import BaseModel, Field, create_model
|
|
from ..mode import Mode
|
|
from ..utils import extract_json_from_stream, extract_json_from_stream_async
|
|
|
|
if TYPE_CHECKING:
|
|
pass
|
|
|
|
|
|
class IterableBase:
|
|
task_type: ClassVar[Optional[type[BaseModel]]] = None
|
|
|
|
@classmethod
|
|
def from_streaming_response(
|
|
cls, completion: Iterable[Any], mode: Mode, **kwargs: Any
|
|
) -> Generator[BaseModel, None, None]: # noqa: ARG003
|
|
json_chunks = cls.extract_json(completion, mode)
|
|
|
|
if mode in {Mode.MD_JSON, Mode.GEMINI_TOOLS}:
|
|
json_chunks = extract_json_from_stream(json_chunks)
|
|
|
|
if mode in {Mode.VERTEXAI_TOOLS, Mode.MISTRAL_TOOLS}:
|
|
response = next(json_chunks)
|
|
if not response:
|
|
return
|
|
|
|
json_response = json.loads(response)
|
|
if not json_response["tasks"]:
|
|
return
|
|
|
|
for item in json_response["tasks"]:
|
|
yield cls.extract_cls_task_type(json.dumps(item), **kwargs)
|
|
|
|
yield from cls.tasks_from_chunks(json_chunks, **kwargs)
|
|
|
|
@classmethod
|
|
async def from_streaming_response_async(
|
|
cls, completion: AsyncGenerator[Any, None], mode: Mode, **kwargs: Any
|
|
) -> AsyncGenerator[BaseModel, None]:
|
|
json_chunks = cls.extract_json_async(completion, mode)
|
|
|
|
if mode in {Mode.MD_JSON, Mode.GEMINI_TOOLS}:
|
|
json_chunks = extract_json_from_stream_async(json_chunks)
|
|
|
|
if mode in {Mode.MISTRAL_TOOLS, Mode.VERTEXAI_TOOLS}:
|
|
async for item in cls.tasks_from_mistral_chunks(json_chunks, **kwargs):
|
|
yield item
|
|
else:
|
|
async for item in cls.tasks_from_chunks_async(json_chunks, **kwargs):
|
|
yield item
|
|
|
|
@classmethod
|
|
async def tasks_from_mistral_chunks(
|
|
cls, json_chunks: AsyncGenerator[str, None], **kwargs: Any
|
|
) -> AsyncGenerator[BaseModel, None]:
|
|
"""Process streaming chunks from Mistral and VertexAI.
|
|
|
|
Handles the specific JSON format used by these providers when streaming."""
|
|
|
|
async for chunk in json_chunks:
|
|
if not chunk:
|
|
continue
|
|
json_response = json.loads(chunk)
|
|
if not json_response["tasks"]:
|
|
continue
|
|
|
|
for item in json_response["tasks"]:
|
|
obj = cls.extract_cls_task_type(json.dumps(item), **kwargs)
|
|
yield obj
|
|
|
|
@classmethod
|
|
def tasks_from_chunks(
|
|
cls, json_chunks: Iterable[str], **kwargs: Any
|
|
) -> Generator[BaseModel, None, None]:
|
|
started = False
|
|
potential_object = ""
|
|
for chunk in json_chunks:
|
|
potential_object += chunk
|
|
if not started:
|
|
if "[" in chunk:
|
|
started = True
|
|
potential_object = chunk[chunk.find("[") + 1 :]
|
|
|
|
while True:
|
|
task_json, potential_object = cls.get_object(potential_object, 0)
|
|
if task_json:
|
|
assert cls.task_type is not None
|
|
obj = cls.extract_cls_task_type(task_json, **kwargs)
|
|
yield obj
|
|
else:
|
|
break
|
|
|
|
@classmethod
|
|
async def tasks_from_chunks_async(
|
|
cls, json_chunks: AsyncGenerator[str, None], **kwargs: Any
|
|
) -> AsyncGenerator[BaseModel, None]:
|
|
started = False
|
|
potential_object = ""
|
|
async for chunk in json_chunks:
|
|
potential_object += chunk
|
|
if not started:
|
|
if "[" in chunk:
|
|
started = True
|
|
potential_object = chunk[chunk.find("[") + 1 :]
|
|
|
|
while True:
|
|
task_json, potential_object = cls.get_object(potential_object, 0)
|
|
if task_json:
|
|
assert cls.task_type is not None
|
|
obj = cls.extract_cls_task_type(task_json, **kwargs)
|
|
yield obj
|
|
else:
|
|
break
|
|
|
|
@classmethod
|
|
def extract_cls_task_type(
|
|
cls,
|
|
task_json: str,
|
|
**kwargs: Any,
|
|
):
|
|
assert cls.task_type is not None
|
|
if get_origin(cls.task_type) is Union:
|
|
union_members = get_args(cls.task_type)
|
|
for member in union_members:
|
|
try:
|
|
obj = member.model_validate_json(task_json, **kwargs)
|
|
return obj
|
|
except Exception:
|
|
pass
|
|
else:
|
|
return cls.task_type.model_validate_json(task_json, **kwargs)
|
|
raise ValueError(
|
|
f"Failed to extract task type with {task_json} for {cls.task_type}"
|
|
)
|
|
|
|
@staticmethod
|
|
def extract_json(
|
|
completion: Iterable[Any], mode: Mode
|
|
) -> Generator[str, None, None]:
|
|
json_started = False
|
|
for chunk in completion:
|
|
try:
|
|
if mode in {Mode.COHERE_TOOLS, Mode.COHERE_JSON_SCHEMA}:
|
|
event_type = getattr(chunk, "event_type", None)
|
|
if event_type == "text-generation":
|
|
if text := getattr(chunk, "text", None):
|
|
if not json_started:
|
|
json_start = min(
|
|
(
|
|
pos
|
|
for pos in (text.find("{"), text.find("["))
|
|
if pos != -1
|
|
),
|
|
default=-1,
|
|
)
|
|
if json_start == -1:
|
|
continue
|
|
json_started = True
|
|
text = text[json_start:]
|
|
yield text
|
|
elif event_type == "tool-calls-chunk":
|
|
delta = getattr(chunk, "tool_call_delta", None)
|
|
args = getattr(delta, "parameters", None) or getattr(
|
|
delta, "text", None
|
|
)
|
|
if args:
|
|
if not json_started:
|
|
json_start = min(
|
|
(
|
|
pos
|
|
for pos in (args.find("{"), args.find("["))
|
|
if pos != -1
|
|
),
|
|
default=-1,
|
|
)
|
|
if json_start == -1:
|
|
continue
|
|
json_started = True
|
|
args = args[json_start:]
|
|
yield args
|
|
elif text := getattr(chunk, "text", None):
|
|
if not json_started:
|
|
json_start = min(
|
|
(
|
|
pos
|
|
for pos in (text.find("{"), text.find("["))
|
|
if pos != -1
|
|
),
|
|
default=-1,
|
|
)
|
|
if json_start == -1:
|
|
continue
|
|
json_started = True
|
|
text = text[json_start:]
|
|
yield text
|
|
elif event_type == "tool-calls-generation":
|
|
tool_calls = getattr(chunk, "tool_calls", None)
|
|
if tool_calls:
|
|
args = json.dumps(tool_calls[0].parameters)
|
|
if not json_started:
|
|
json_start = min(
|
|
(
|
|
pos
|
|
for pos in (args.find("{"), args.find("["))
|
|
if pos != -1
|
|
),
|
|
default=-1,
|
|
)
|
|
if json_start == -1:
|
|
continue
|
|
json_started = True
|
|
args = args[json_start:]
|
|
yield args
|
|
elif text := getattr(chunk, "text", None):
|
|
if not json_started:
|
|
json_start = min(
|
|
(
|
|
pos
|
|
for pos in (text.find("{"), text.find("["))
|
|
if pos != -1
|
|
),
|
|
default=-1,
|
|
)
|
|
if json_start == -1:
|
|
continue
|
|
json_started = True
|
|
text = text[json_start:]
|
|
yield text
|
|
else:
|
|
chunk_type = getattr(chunk, "type", None)
|
|
if chunk_type == "content-delta":
|
|
delta = getattr(chunk, "delta", None)
|
|
message = getattr(delta, "message", None)
|
|
content = getattr(message, "content", None)
|
|
if text := getattr(content, "text", None):
|
|
if not json_started:
|
|
json_start = min(
|
|
(
|
|
pos
|
|
for pos in (
|
|
text.find("{"),
|
|
text.find("["),
|
|
)
|
|
if pos != -1
|
|
),
|
|
default=-1,
|
|
)
|
|
if json_start == -1:
|
|
continue
|
|
json_started = True
|
|
text = text[json_start:]
|
|
yield text
|
|
elif chunk_type == "tool-call-delta":
|
|
delta = getattr(chunk, "delta", None)
|
|
message = getattr(delta, "message", None)
|
|
tool_calls = getattr(message, "tool_calls", None)
|
|
function = getattr(tool_calls, "function", None)
|
|
if args := getattr(function, "arguments", None):
|
|
if not json_started:
|
|
json_start = min(
|
|
(
|
|
pos
|
|
for pos in (
|
|
args.find("{"),
|
|
args.find("["),
|
|
)
|
|
if pos != -1
|
|
),
|
|
default=-1,
|
|
)
|
|
if json_start == -1:
|
|
continue
|
|
json_started = True
|
|
args = args[json_start:]
|
|
yield args
|
|
if mode == Mode.ANTHROPIC_JSON:
|
|
if json_chunk := chunk.delta.text:
|
|
yield json_chunk
|
|
if mode == Mode.ANTHROPIC_TOOLS:
|
|
yield chunk.delta.partial_json
|
|
if mode == Mode.GEMINI_JSON:
|
|
yield chunk.text
|
|
if mode == Mode.VERTEXAI_JSON:
|
|
yield chunk.candidates[0].content.parts[0].text
|
|
if mode == Mode.VERTEXAI_TOOLS:
|
|
yield json.dumps(
|
|
chunk.candidates[0].content.parts[0].function_call.args
|
|
)
|
|
if mode == Mode.MISTRAL_STRUCTURED_OUTPUTS:
|
|
yield chunk.data.choices[0].delta.content
|
|
if mode == Mode.MISTRAL_TOOLS:
|
|
if not chunk.data.choices[0].delta.tool_calls:
|
|
continue
|
|
yield chunk.data.choices[0].delta.tool_calls[0].function.arguments
|
|
|
|
if mode in {Mode.GENAI_TOOLS}:
|
|
yield json.dumps(
|
|
chunk.candidates[0].content.parts[0].function_call.args
|
|
)
|
|
if mode in {Mode.GENAI_STRUCTURED_OUTPUTS}:
|
|
yield chunk.candidates[0].content.parts[0].text
|
|
|
|
if mode in {Mode.GEMINI_TOOLS}:
|
|
resp = chunk.candidates[0].content.parts[0].function_call
|
|
resp_dict = type(resp).to_dict(resp) # type:ignore
|
|
|
|
if "args" in resp_dict:
|
|
yield json.dumps(resp_dict["args"])
|
|
|
|
if mode in {
|
|
Mode.RESPONSES_TOOLS,
|
|
Mode.RESPONSES_TOOLS_WITH_INBUILT_TOOLS,
|
|
}:
|
|
from openai.types.responses import (
|
|
ResponseFunctionCallArgumentsDeltaEvent,
|
|
)
|
|
|
|
if isinstance(chunk, ResponseFunctionCallArgumentsDeltaEvent):
|
|
yield chunk.delta
|
|
elif chunk.choices:
|
|
if mode == Mode.FUNCTIONS:
|
|
Mode.warn_mode_functions_deprecation()
|
|
if json_chunk := chunk.choices[0].delta.function_call.arguments:
|
|
yield json_chunk
|
|
elif mode in {
|
|
Mode.JSON,
|
|
Mode.MD_JSON,
|
|
Mode.JSON_SCHEMA,
|
|
Mode.CEREBRAS_JSON,
|
|
Mode.FIREWORKS_JSON,
|
|
Mode.PERPLEXITY_JSON,
|
|
Mode.WRITER_JSON,
|
|
}:
|
|
if json_chunk := chunk.choices[0].delta.content:
|
|
yield json_chunk
|
|
elif mode in {
|
|
Mode.TOOLS,
|
|
Mode.TOOLS_STRICT,
|
|
Mode.FIREWORKS_TOOLS,
|
|
Mode.WRITER_TOOLS,
|
|
}:
|
|
if json_chunk := chunk.choices[0].delta.tool_calls:
|
|
if json_chunk[0].function.arguments is not None:
|
|
yield json_chunk[0].function.arguments
|
|
else:
|
|
raise NotImplementedError(
|
|
f"Mode {mode} is not supported for MultiTask streaming"
|
|
)
|
|
except AttributeError:
|
|
pass
|
|
|
|
@staticmethod
|
|
async def extract_json_async(
|
|
completion: AsyncGenerator[Any, None], mode: Mode
|
|
) -> AsyncGenerator[str, None]:
|
|
json_started = False
|
|
async for chunk in completion:
|
|
try:
|
|
if mode in {Mode.COHERE_TOOLS, Mode.COHERE_JSON_SCHEMA}:
|
|
event_type = getattr(chunk, "event_type", None)
|
|
if event_type == "text-generation":
|
|
if text := getattr(chunk, "text", None):
|
|
if not json_started:
|
|
json_start = min(
|
|
(
|
|
pos
|
|
for pos in (text.find("{"), text.find("["))
|
|
if pos != -1
|
|
),
|
|
default=-1,
|
|
)
|
|
if json_start == -1:
|
|
continue
|
|
json_started = True
|
|
text = text[json_start:]
|
|
yield text
|
|
elif event_type == "tool-calls-chunk":
|
|
delta = getattr(chunk, "tool_call_delta", None)
|
|
args = getattr(delta, "parameters", None) or getattr(
|
|
delta, "text", None
|
|
)
|
|
if args:
|
|
if not json_started:
|
|
json_start = min(
|
|
(
|
|
pos
|
|
for pos in (args.find("{"), args.find("["))
|
|
if pos != -1
|
|
),
|
|
default=-1,
|
|
)
|
|
if json_start == -1:
|
|
continue
|
|
json_started = True
|
|
args = args[json_start:]
|
|
yield args
|
|
elif text := getattr(chunk, "text", None):
|
|
if not json_started:
|
|
json_start = min(
|
|
(
|
|
pos
|
|
for pos in (text.find("{"), text.find("["))
|
|
if pos != -1
|
|
),
|
|
default=-1,
|
|
)
|
|
if json_start == -1:
|
|
continue
|
|
json_started = True
|
|
text = text[json_start:]
|
|
yield text
|
|
elif event_type == "tool-calls-generation":
|
|
tool_calls = getattr(chunk, "tool_calls", None)
|
|
if tool_calls:
|
|
args = json.dumps(tool_calls[0].parameters)
|
|
if not json_started:
|
|
json_start = min(
|
|
(
|
|
pos
|
|
for pos in (args.find("{"), args.find("["))
|
|
if pos != -1
|
|
),
|
|
default=-1,
|
|
)
|
|
if json_start == -1:
|
|
continue
|
|
json_started = True
|
|
args = args[json_start:]
|
|
yield args
|
|
elif text := getattr(chunk, "text", None):
|
|
if not json_started:
|
|
json_start = min(
|
|
(
|
|
pos
|
|
for pos in (text.find("{"), text.find("["))
|
|
if pos != -1
|
|
),
|
|
default=-1,
|
|
)
|
|
if json_start == -1:
|
|
continue
|
|
json_started = True
|
|
text = text[json_start:]
|
|
yield text
|
|
else:
|
|
chunk_type = getattr(chunk, "type", None)
|
|
if chunk_type == "content-delta":
|
|
delta = getattr(chunk, "delta", None)
|
|
message = getattr(delta, "message", None)
|
|
content = getattr(message, "content", None)
|
|
if text := getattr(content, "text", None):
|
|
if not json_started:
|
|
json_start = min(
|
|
(
|
|
pos
|
|
for pos in (
|
|
text.find("{"),
|
|
text.find("["),
|
|
)
|
|
if pos != -1
|
|
),
|
|
default=-1,
|
|
)
|
|
if json_start == -1:
|
|
continue
|
|
json_started = True
|
|
text = text[json_start:]
|
|
yield text
|
|
elif chunk_type == "tool-call-delta":
|
|
delta = getattr(chunk, "delta", None)
|
|
message = getattr(delta, "message", None)
|
|
tool_calls = getattr(message, "tool_calls", None)
|
|
function = getattr(tool_calls, "function", None)
|
|
if args := getattr(function, "arguments", None):
|
|
if not json_started:
|
|
json_start = min(
|
|
(
|
|
pos
|
|
for pos in (
|
|
args.find("{"),
|
|
args.find("["),
|
|
)
|
|
if pos != -1
|
|
),
|
|
default=-1,
|
|
)
|
|
if json_start == -1:
|
|
continue
|
|
json_started = True
|
|
args = args[json_start:]
|
|
yield args
|
|
if mode == Mode.ANTHROPIC_JSON:
|
|
if json_chunk := chunk.delta.text:
|
|
yield json_chunk
|
|
if mode == Mode.ANTHROPIC_TOOLS:
|
|
yield chunk.delta.partial_json
|
|
if mode == Mode.VERTEXAI_JSON:
|
|
yield chunk.candidates[0].content.parts[0].text
|
|
if mode == Mode.VERTEXAI_TOOLS:
|
|
yield json.dumps(
|
|
chunk.candidates[0].content.parts[0].function_call.args
|
|
)
|
|
if mode == Mode.MISTRAL_STRUCTURED_OUTPUTS:
|
|
yield chunk.data.choices[0].delta.content
|
|
if mode == Mode.MISTRAL_TOOLS:
|
|
if not chunk.data.choices[0].delta.tool_calls:
|
|
continue
|
|
yield chunk.data.choices[0].delta.tool_calls[0].function.arguments
|
|
if mode == Mode.GENAI_STRUCTURED_OUTPUTS:
|
|
yield chunk.text
|
|
if mode in {Mode.GENAI_TOOLS}:
|
|
yield json.dumps(
|
|
chunk.candidates[0].content.parts[0].function_call.args
|
|
)
|
|
if mode in {
|
|
Mode.RESPONSES_TOOLS,
|
|
Mode.RESPONSES_TOOLS_WITH_INBUILT_TOOLS,
|
|
}:
|
|
from openai.types.responses import (
|
|
ResponseFunctionCallArgumentsDeltaEvent,
|
|
)
|
|
|
|
if isinstance(chunk, ResponseFunctionCallArgumentsDeltaEvent):
|
|
yield chunk.delta
|
|
elif chunk.choices:
|
|
if mode == Mode.FUNCTIONS:
|
|
Mode.warn_mode_functions_deprecation()
|
|
if json_chunk := chunk.choices[0].delta.function_call.arguments:
|
|
yield json_chunk
|
|
elif mode in {
|
|
Mode.JSON,
|
|
Mode.MD_JSON,
|
|
Mode.JSON_SCHEMA,
|
|
Mode.CEREBRAS_JSON,
|
|
Mode.FIREWORKS_JSON,
|
|
Mode.PERPLEXITY_JSON,
|
|
Mode.WRITER_JSON,
|
|
}:
|
|
if json_chunk := chunk.choices[0].delta.content:
|
|
yield json_chunk
|
|
elif mode in {
|
|
Mode.TOOLS,
|
|
Mode.TOOLS_STRICT,
|
|
Mode.FIREWORKS_TOOLS,
|
|
Mode.WRITER_TOOLS,
|
|
}:
|
|
if json_chunk := chunk.choices[0].delta.tool_calls:
|
|
if json_chunk[0].function.arguments is not None:
|
|
yield json_chunk[0].function.arguments
|
|
else:
|
|
raise NotImplementedError(
|
|
f"Mode {mode} is not supported for MultiTask streaming"
|
|
)
|
|
except AttributeError:
|
|
pass
|
|
|
|
@staticmethod
|
|
def get_object(s: str, stack: int) -> tuple[Optional[str], str]:
|
|
start_index = s.find("{")
|
|
for i, c in enumerate(s):
|
|
if c == "{":
|
|
stack += 1
|
|
if c == "}":
|
|
stack -= 1
|
|
if stack == 0:
|
|
return s[start_index : i + 1], s[i + 2 :]
|
|
return None, s
|
|
|
|
|
|
def IterableModel(
|
|
subtask_class: type[BaseModel],
|
|
name: Optional[str] = None,
|
|
description: Optional[str] = None,
|
|
) -> type[BaseModel]:
|
|
# Import at runtime to avoid circular import
|
|
from ..processing.function_calls import OpenAISchema
|
|
|
|
"""
|
|
Dynamically create a IterableModel OpenAISchema that can be used to segment multiple
|
|
tasks given a base class. This creates class that can be used to create a toolkit
|
|
for a specific task, names and descriptions are automatically generated. However
|
|
they can be overridden.
|
|
|
|
## Usage
|
|
|
|
```python
|
|
from pydantic import BaseModel, Field
|
|
from instructor import IterableModel
|
|
|
|
class User(BaseModel):
|
|
name: str = Field(description="The name of the person")
|
|
age: int = Field(description="The age of the person")
|
|
role: str = Field(description="The role of the person")
|
|
|
|
MultiUser = IterableModel(User)
|
|
```
|
|
|
|
## Result
|
|
|
|
```python
|
|
class MultiUser(OpenAISchema, MultiTaskBase):
|
|
tasks: List[User] = Field(
|
|
default_factory=list,
|
|
repr=False,
|
|
description="Correctly segmented list of `User` tasks",
|
|
)
|
|
|
|
@classmethod
|
|
def from_streaming_response(cls, completion) -> Generator[User]:
|
|
'''
|
|
Parse the streaming response from OpenAI and yield a `User` object
|
|
for each task in the response
|
|
'''
|
|
json_chunks = cls.extract_json(completion)
|
|
yield from cls.tasks_from_chunks(json_chunks)
|
|
```
|
|
|
|
Parameters:
|
|
subtask_class (Type[OpenAISchema]): The base class to use for the MultiTask
|
|
name (Optional[str]): The name of the MultiTask class, if None then the name
|
|
of the subtask class is used as `Multi{subtask_class.__name__}`
|
|
description (Optional[str]): The description of the MultiTask class, if None
|
|
then the description is set to `Correct segmentation of `{subtask_class.__name__}` tasks`
|
|
|
|
Returns:
|
|
schema (OpenAISchema): A new class that can be used to segment multiple tasks
|
|
"""
|
|
if name is not None:
|
|
task_name = name
|
|
else:
|
|
# Handle `Union[A, B]` / `A | B` task types.
|
|
# `types.UnionType` does not have `__name__`, so fall back to a stable name.
|
|
task_name = getattr(subtask_class, "__name__", None)
|
|
if task_name is None and get_origin(subtask_class) is Union:
|
|
members = get_args(subtask_class)
|
|
task_name = "Or".join(getattr(m, "__name__", str(m)) for m in members)
|
|
if task_name is None:
|
|
task_name = str(subtask_class)
|
|
|
|
name = f"Iterable{task_name}"
|
|
|
|
list_tasks = (
|
|
list[subtask_class], # type: ignore
|
|
Field(
|
|
default_factory=list,
|
|
repr=False,
|
|
description=f"Correctly segmented list of `{task_name}` tasks",
|
|
),
|
|
)
|
|
|
|
base_models = cast(tuple[type[BaseModel], ...], (OpenAISchema, IterableBase))
|
|
new_cls = create_model(
|
|
name,
|
|
tasks=list_tasks,
|
|
__base__=base_models,
|
|
)
|
|
new_cls = cast(type[IterableBase], new_cls)
|
|
|
|
# set the class constructor BaseModel
|
|
new_cls.task_type = subtask_class
|
|
|
|
new_cls.__doc__ = (
|
|
f"Correct segmentation of `{task_name}` tasks"
|
|
if description is None
|
|
else description
|
|
)
|
|
assert issubclass(new_cls, OpenAISchema), (
|
|
"The new class should be a subclass of OpenAISchema"
|
|
)
|
|
return new_cls
|