o
    -i/                     @   s   d dl Z d dlZd dlmZ d dlmZ d dlZd dlm	Z	 d dl
mZ d dlmZ d dlmZmZmZmZmZmZ d dlmZ d dlmZ eeZG d	d
 d
eZG dd deZde jdefddZde j defddZ!de"de#e"e"f dB fddZ$de"dede%de"dedB f
ddZ&dS )    N)Sequence)Any)PreTrainedTokenizerBase)ChatCompletionRequest)DeltaFunctionCallDeltaMessageDeltaToolCallExtractedToolCallInformationFunctionCallToolCall)init_logger)
ToolParserc                   @   s   e Zd ZdS )_UnexpectedAstErrorN)__name__
__module____qualname__ r   r   c/home/ubuntu/veenaModal/venv/lib/python3.10/site-packages/vllm/tool_parsers/pythonic_tool_parser.pyr       s    r   c                       s   e Zd ZdZedejZdef fddZ	e
defddZejd	edd
fddZdededefddZdedededee dee dee deded
B fddZ  ZS )PythonicToolParserz
    Tool call parser for models that produce tool calls in a pythonic style,
    such as Llama 3.2 and Llama 4 models.

    Used when --enable-auto-tool-choice --tool-call-parser pythonic are all set
    z\[([a-zA-Z]+\w*\(([a-zA-Z]+\w*=.*,\s*)*([a-zA-Z]+\w*=.*\s)?\),\s*)*([a-zA-Z]+\w*\(([a-zA-Z]+\w*=.*,\s*)*([a-zA-Z]+\w*=.*\s*)?\)\s*)+\]	tokenizerc                    s   t  | d S N)super__init__)selfr   	__class__r   r   r   8   s   zPythonicToolParser.__init__returnc                 C   s   | j S r   current_tool_id)r   r   r   r   current_tool_index<   s   z%PythonicToolParser.current_tool_indexvalueNc                 C   s
   || _ d S r   r   )r   r    r   r   r   r   @   s   
model_outputrequestc                 C   s   d}z| j j|tjddu}W n ty#   td td| Y nw |s-tdg |dS z0t	
|}t|jd dd}t|t	jrZtd	d
 |jD rZtddd |jD ddW S td tyr   td tdg |d Y S w )zH
        Extract the tool calls from a complete model response.
        F)timeoutNz7Regex timeout occurred when matching tool call pattern.z3Regex timeout occurred when matching user input: %s)tools_called
tool_callscontentr   r    c                 s       | ]	}t |tjV  qd S r   
isinstanceastCall.0er   r   r   	<genexpr>`       
z8PythonicToolParser.extract_tool_calls.<locals>.<genexpr>Tc                 S      g | ]}t |qS r   _handle_single_toolr,   r   r   r   
<listcomp>e       z9PythonicToolParser.extract_tool_calls.<locals>.<listcomp>,Tool output must be a list of function callsz,Error in extracting tool call from response.)TOOL_CALL_REGEXmatchenvs%VLLM_TOOL_PARSE_REGEX_TIMEOUT_SECONDSTimeoutErrorloggerwarningdebugr	   r*   parsegetattrbodyr)   Listalleltsr   	Exception	exception)r   r!   r"   is_tool_call_patternmoduleparsedr   r   r   extract_tool_callsD   sP   

	

z%PythonicToolParser.extract_tool_callsprevious_textcurrent_text
delta_textprevious_token_idscurrent_token_idsdelta_token_idsc                 C   s  | ds
t|dS zt|}|d u rW d S |\}	}
t|	}t|jd dd }t|tjr8t	dd |j
D s<tddd	 |j
D }g }t|D ]t\}}|| jk rTqJ|| _t| j|krd| jd
 |t|d k pod|
v}|ry|  jd7  _|s|
d d nd
}|s|
d dkr|d }|dd}t| j| |||}|d ur|| |jd ur|jjd ur| j|  |jj7  < qJ|r| jsdi ig| _|rt|dW S |
s| jdkrtd
dW S W d S  ty   td td Y d S w )N[)r&   r   r    c                 s   r'   r   r(   r,   r   r   r   r/      r0   zBPythonicToolParser.extract_tool_calls_streaming.<locals>.<genexpr>r6   c                 S   r1   r   r2   r,   r   r   r   r4      r5   zCPythonicToolParser.extract_tool_calls_streaming.<locals>.<listcomp>    z)])}'"	arguments)r%   z+Error trying to handle streaming tool call.z=Skipping chunk as a result of tool streaming extraction error)
startswithr   _make_valid_pythonr*   r?   r@   rA   r)   rB   rC   rD   r   	enumerater   lenstreamed_args_for_toolappendreplace_compute_tool_deltafunctionrY   prev_tool_call_arrr   rE   r<   rF   r>   )r   rK   rL   rM   rN   rO   rP   r"   valid_and_added_text
valid_text
added_textrH   rI   r%   tool_deltasindexnew_callnew_call_completewithheld_suffixdeltar   r   r   extract_tool_calls_streamingv   sp   








z/PythonicToolParser.extract_tool_calls_streaming)r   r   r   __doc__recompileDOTALLr7   r   r   propertyintr   setterstrr   r	   rJ   r   r   rm   __classcell__r   r   r   r   r   $   sF    
2	r   valr   c                 C   sv   t | tjr	| jS t | tjr)tdd | jD stddd t| j| j	D S t | tj
r7dd | jD S td)	Nc                 s   r'   r   )r)   r*   Constant)r-   kr   r   r   r/      s    z'_get_parameter_value.<locals>.<genexpr>z/Dict tool call arguments must have literal keysc                 S   s   i | ]
\}}|j t|qS r   )r    _get_parameter_value)r-   ry   vr   r   r   
<dictcomp>   s    
z(_get_parameter_value.<locals>.<dictcomp>c                 S   r1   r   )rz   )r-   r{   r   r   r   r4      s    z(_get_parameter_value.<locals>.<listcomp>z$Tool call arguments must be literals)r)   r*   rx   r    DictrC   keysr   zipvaluesrB   rD   )rw   r   r   r   rz      s   rz   callc                 C   s\   t | jtjstd| jj}i }| jD ]
}t|j||j	< qt
dt|tj|ddddS )NzInvalid tool call namerb   F)ensure_asciinamerY   )typerb   )r)   funcr*   Namer   idkeywordsrz   r    argr   r
   jsondumps)r   function_namerY   keywordr   r   r   r3      s   
r3   textc                 C   sF  g }t | D ]m\}}|dv r|| q|dkr%|r | dkr$tdq|dkr6|r1| dkr5tdq|dkrG|rB| d	krFtd
q|dv rs|re|d |kre|dkr`| |d  dkr`q|  q|rn|d dv rnq|| q|  } | ds| drd S |r|d d	kr| d | d	 }|d}|d}||krd S |r|d dkr| d | d }|d}|d}	||	krd S | dr| d d } |r|d dkr| ds| dsd S d}
t|D ]1}|dkr|
d7 }
q|dkr|
d7 }
q|d	kr	|
d7 }
q|dkr|
d7 }
q|dkr|
d7 }
q| |
 |
fS )N>   {(rQ   ]rQ   zMismatched square bracketsrU   r   zMismatched parenthesesrV   r   zMismatched curly braces>   rX   rW   r   rS   \=:,rR   rW   rX   )	r\   r_   popr   rstripendswithrfindcountreversed)r   bracket_stackrh   chartrailing_dict_textnum_keys
num_valuestrailing_params_textnum_full_param_namesnum_full_param_valuesrf   r   r   r   r[      s   













r[   previously_sent_argsri   rh   rk   c                 C   sz   |j j}|r||sJ |d t|  }| s't|jd|t|j j|ddS |t| d  }|r;td |t|ddS d S )Nrb   r   )r   r   rh   rb   )rY   )r   rh   rb   )rb   rY   r   r]   r   r   r   r   )r   ri   rh   rk   new_call_argsarg_diffr   r   r   ra   5  s,   
ra   )'r*   r   collections.abcr   typingr   regexro   transformersr   	vllm.envsr9   0vllm.entrypoints.openai.chat_completion.protocolr   'vllm.entrypoints.openai.engine.protocolr   r   r   r	   r
   r   vllm.loggerr   &vllm.tool_parsers.abstract_tool_parserr   r   r<   rE   r   r   exprrz   r+   r3   ru   tupler[   rs   ra   r   r   r   r   <module>   s:     -F