o iM@sdZddlmZddlmZddlmZmZm Z ddl m Z ddl m Z ddlmZmZmZmZmZmZmZmZmZmZmZdd lmZeZ eeBZ eGd d d Zd S) aThis module provides functionality to manage and update parts of a model's streamed response. The manager tracks which parts (in particular, text and tool calls) correspond to which vendor-specific identifiers (e.g., `index`, `tool_call_id`, etc., as appropriate for a given model), and produces Pydantic AI-format events as appropriate for consumers of the streaming APIs. The "vendor-specific identifiers" to use depend on the semantics of the responses of the responses from the vendor, and are tightly coupled to the specific model being used, and the Pydantic AI Model subclass implementation. This `ModelResponsePartsManager` is used in each of the subclasses of `StreamedResponse` as a way to consolidate event-emitting logic. ) annotations)Hashable) dataclassfieldreplace)Any)UnexpectedModelBehavior) BuiltinToolCallPartModelResponsePartModelResponseStreamEventPartDeltaEventPartStartEventTextPart TextPartDelta ThinkingPartThinkingPartDelta ToolCallPartToolCallPartDelta)generate_tool_call_idc@seZdZUdZeeddZded< eeddZ ded< d0d d Z d d dd d1ddZ d d d d dd2d d!Z d d d d"d3d'd(Z d d d)d4d*d+Zd5d.d/Zd S)6ModelResponsePartsManagerzManages a sequence of parts that make up a model's streamed response. Parts are generally added and/or updated by providing deltas, which are tracked by vendor-specific IDs. F)default_factoryinitzlist[ManagedPart]_partszdict[VendorId, int]_vendor_id_to_part_indexreturnlist[ModelResponsePart]cCsdd|jDS)zReturn only model response parts that are complete (i.e., not ToolCallPartDelta's). Returns: A list of ModelResponsePart objects. ToolCallPartDelta objects are excluded. cSsg|] }t|ts|qS) isinstancer).0prr`/var/www/html/karishye-ai-python/venv/lib/python3.10/site-packages/pydantic_ai/_parts_manager.py Dsz7ModelResponsePartsManager.get_parts..)r)selfrrr! get_parts>sz#ModelResponsePartsManager.get_partsN)id thinking_tagsignore_leading_whitespacevendor_part_idVendorId | Nonecontentstrr% str | Noner&tuple[str, str] | Noner'boolModelResponseStreamEvent | NonecCstd}|dur|jrt|jd}|j|}t|tr||f}n<|j|}|dur[|j|} |rJt| trJ||dkrC|j|dS|j||dSt| trT| |f}nt d| |rq||dkrq|j|d|j|ddS|dur|rt|dks| rdSt|j} t||d} |dur| |j|<|j | t | | dS|\} }t |d } | | |j|<t|| d S) aHandle incoming text content, creating or updating a TextPart in the manager as appropriate. When `vendor_part_id` is None, the latest part is updated if it exists and is a TextPart; otherwise, a new TextPart is created. When a non-None ID is specified, the TextPart corresponding to that vendor ID is either created or updated. Args: vendor_part_id: The ID the vendor uses to identify this piece of text. If None, a new part will be created unless the latest part is already a TextPart. content: The text content to append to the appropriate TextPart. id: An optional id for the text part. thinking_tags: If provided, will handle content between the thinking tags as thinking parts. ignore_leading_whitespace: If True, will ignore leading whitespace in the content. Returns: - A `PartStartEvent` if a new part was created. - A `PartDeltaEvent` if an existing part was updated. - `None` if no new event is emitted (e.g., the first text part was all whitespace). Raises: UnexpectedModelBehavior: If attempting to apply text content to a part that is not a TextPart. Nr)r(r*z+Cannot apply a text delta to existing_part=r)r*r%indexpart) content_deltar2delta)rlenrrrgetrpophandle_thinking_deltarisspaceappendr rapplyr )r#r(r*r%r&r'existing_text_part_and_index part_index latest_part existing_partnew_part_indexr3existing_text_part part_deltarrr!handle_text_deltaFsF                z+ModelResponsePartsManager.handle_text_deltar*r% signature provider_nameHashable | NonerGrHr cCs.d}|dur|jrt|jd}|j|}t|tr||f}n|j|}|dur>|j|} t| ts:td| | |f}|durr|dusJ|durnt|j} t|pSd|||d} |durb| |j|<|j| t| | dStd|dusz|dur|\} }t |||d} | | |j|<t || d Std ) aHandle incoming thinking content, creating or updating a ThinkingPart in the manager as appropriate. When `vendor_part_id` is None, the latest part is updated if it exists and is a ThinkingPart; otherwise, a new ThinkingPart is created. When a non-None ID is specified, the ThinkingPart corresponding to that vendor ID is either created or updated. Args: vendor_part_id: The ID the vendor uses to identify this piece of thinking. If None, a new part will be created unless the latest part is already a ThinkingPart. content: The thinking content to append to the appropriate ThinkingPart. id: An optional id for the thinking part. signature: An optional signature for the thinking content. provider_name: An optional provider name for the thinking part. Returns: A `PartStartEvent` if a new part was created, or a `PartDeltaEvent` if an existing part was updated. Raises: UnexpectedModelBehavior: If attempting to apply a thinking delta to a part that is not a ThinkingPart. Nrz/Cannot apply a thinking delta to existing_part=r0rFr1z9Cannot create a ThinkingPart with no content or signature)r4signature_deltarHr5z9Cannot update a ThinkingPart with no content or signature) rr7rrrr8rr<r rr=r )r#r(r*r%rGrH existing_thinking_part_and_indexr?r@rArBr3existing_thinking_partrDrrr!r:s>          z/ModelResponsePartsManager.handle_thinking_delta) tool_nameargs tool_call_idrMrNstr | dict[str, Any] | NonerOc Cszd}|dur'|dur&|jr&t|jd}|j|}t|ttBtBr&||f}n#|j|}|durJ|j|}t|ttBtBsFtd|||f}|durt|||d} | pZ| } |durgt|j|j|<t|j} |j | t| ttBrt | | dSdS|\}}t|||d} | |} | |j|<t| ttBrt|trt || dS| j r| j st| | j d} t|| dSdS)atHandle or update a tool call, creating or updating a `ToolCallPart`, `BuiltinToolCallPart`, or `ToolCallPartDelta`. Managed items remain as `ToolCallPartDelta`s until they have at least a tool_name, at which point they are upgraded to `ToolCallPart`s. If `vendor_part_id` is None, updates the latest matching ToolCallPart (or ToolCallPartDelta) if any. Otherwise, a new part (or delta) may be created. Args: vendor_part_id: The ID the vendor uses for this tool call. If None, the latest matching tool call may be updated. tool_name: The name of the tool. If None, the manager does not enforce a name match when `vendor_part_id` is None. args: Arguments for the tool call, either as a string, a dictionary of key-value pairs, or None. tool_call_id: An optional string representing an identifier for this tool call. Returns: - A `PartStartEvent` if a new ToolCallPart or BuiltinToolCallPart is created. - A `PartDeltaEvent` if an existing part is updated. - `None` if no new event is emitted (e.g., the part is still incomplete). Raises: UnexpectedModelBehavior: If attempting to apply a tool call delta to a part that is not a ToolCallPart, BuiltinToolCallPart, or ToolCallPartDelta. Nrz0Cannot apply a tool call delta to existing_part=)tool_name_delta args_deltarOr1)rOr5)rr7rrr rrr8ras_partr<r r=rOrr ) r#r(rMrNrO existing_matching_part_and_indexr?r@rAr6r3rB updated_partrrr!handle_tool_call_deltasH"             z0ModelResponsePartsManager.handle_tool_call_delta)rOr%c Cst|||pt|d}|durt|j}|j|n*|j|}|dur5t|j|tr5|}||j|<n t|j}|j|||j|<t||dS)a4Immediately create or fully-overwrite a ToolCallPart with the given information. This does not apply a delta; it directly sets the tool call part contents. Args: vendor_part_id: The vendor's ID for this tool call part. If not None and an existing part is found, that part is overwritten. tool_name: The name of the tool being invoked. args: The arguments for the tool call, either as a string, a dictionary, or None. tool_call_id: An optional string identifier for this tool call. id: An optional identifier for this tool call part. Returns: ModelResponseStreamEvent: A `PartStartEvent` indicating that a new tool call part has been added to the manager, or replaced an existing part. )rMrNrOr%Nr1) r_generate_tool_call_idr7rr<rr8rr ) r#r(rMrNrOr%new_partrBmaybe_part_indexrrr!handle_tool_call_part4s"       z/ModelResponsePartsManager.handle_tool_call_partr3r cCs|durt|j}|j|n,|j|}|dur,t|j|t|r,|}||j|<n t|j}|j|||j|<t||dS)aCreate or overwrite a ModelResponsePart. Args: vendor_part_id: The vendor's ID for this tool call part. If not None and an existing part is found, that part is overwritten. part: The ModelResponsePart. Returns: ModelResponseStreamEvent: A `PartStartEvent` indicating that a new part has been added to the manager, or replaced an existing part. Nr1)r7rr<rr8rtyper )r#r(r3rBrYrrr! handle_partcs       z%ModelResponsePartsManager.handle_part)rr) r(r)r*r+r%r,r&r-r'r.rr/) r(rIr*r,r%r,rGr,rHr,rr ) r(rIrMr,rNrPrOr,rr/) r(rIrMr+rNrPrOr,r%r,rr )r(rIr3r rr )__name__ __module__ __qualname____doc__rlistr__annotations__dictrr$rEr:rVrZr\rrrr!r2s0   YKX/rN) r` __future__r _annotationscollections.abcr dataclassesrrrtypingrpydantic_ai.exceptionsrpydantic_ai.messagesr r r r r rrrrrr_utilsrrWVendorId ManagedPartrrrrr!s   4