
    %Gg                         d dl Z d dlZd dlmZ d dlmZmZ d dlmZ d dl	m
Z
 d dlmZmZ d dlmZmZ dZd	Z G d
 de      Zy)    N)cast)TemplateSyntaxErrornodes)	Extension)ValidationError)InvalidPromptErrorLLMError)ChatMessageTool)modelzGlitellm is not installed. Please install it with `pip install litellm`.c                   T    e Zd ZdZdhZd Zd Zd Zd Zde	de
ee   ee   f   fd	Zy
)CompletionExtensiona7  
    `completion` can be used to send to the LLM the content of the block in form of messages.

    The rendered value of the block can be used as is but it's usually more useful to
    assign it to a variable and access it from another section of the prompt.

    Example:
        ```
        {% set response %}
        {% completion model="gpt-3.5-turbo-0125" %}
        {% chat role="user" %}You are a helpful assistant{% endchat %}
        {% endcompletion %}
        {% endset %}

        {# output the response content #}
        {{ response }}
        ```
    
completionc                 .   t        |j                        j                  }g }|j                  j                  j                  dk7  rH|j                  t        |j                               |j                  j                  j                  dk7  rHd| }	 |\  }}}|j                  t        vs|j                  dk7  rt        ||      t        j                  |j                        g}|j                  dd      }	|j                  j                  r7t        j                  | j!                  d|      g g |	      j#                  |      S t        j                  | j!                  d|      g g |	      j#                  |      S # t        $ r t        ||      d w xY w)	N	block_endzInvalid syntax for completion: =)zname:endcompletionT)drop_needle_do_completion_async_do_completion)nextstreamlinenocurrenttypeappend
ValueErrorr   valueSUPPORTED_KWARGSr   Constparse_statementsenvironmentis_async	CallBlockcall_method
set_lineno)
selfparserr   gathered	error_msg	attr_nameattr_assign
attr_valueargsbodys
             x/home/kushmeetdev/apache_webroot/langgraph_flaskproject/venv/lib/python3.12/site-packages/banks/extensions/completion.pyparsezCompletionExtension.parse*   ss   fmm$++ mm##((K7OOD/0 mm##((K7 6hZ@		C19.I{J
 ??"22k6G6G36N%i88 #(++j.>.>"?!@ &&'>D&Q &&??4#3#34JD#QSUWY[_`kklrsst//0@$GRQUVaabhii!  	C%i8dB	Cs   E= =Fc                 :   |D ]r  }|j                   j                  |j                   j                  k(  s1|j                  j                  dd      \  }}t	        j
                  |      }t        ||      c S  d|j                   j                   d}t        |      )N.   )maxsplitz	Function z not found in available tools)functionnameimport_pathrsplit	importlibimport_modulegetattrr   )r&   tools	tool_calltoolmodule_name	func_namemodulemsgs           r/   _get_tool_callablez&CompletionExtension._get_tool_callableM   s     	2D}}!!Y%7%7%<%<<)-)9)9)@)@q)@)Q&Y"00=vy11		2
 ),,1122OPo    c           	      N   	 ddl m} ddlm}m} | j                   |             \  }}|D 	cg c]  }	|	j                          }
}	|D cg c]  }|j                          c}xs d}t        | |||
|            }t        t        |   |j                        }|d   j                  j                  }|s|d   j                  j                  S |
j                  |d   j                  j                                |D ]  }|j                   j"                  sd}t%        |      | j'                  ||      }t)        j*                  |j                   j,                        } |d	i |}|
j                  t/        |j0                  d|j                   j"                  |      j                                 t        | |||
|            }t        t        |   |j                        }|d   j                  j                  S # t
        $ r}t        t              |d}~ww xY wc c}	w c c}w )
"
        Helper callback.
        r   )r   ChoicesModelResponseNr   messagesr<   z*Malformed response: function name is emptyr>   tool_call_idroler6   content )litellmr   litellm.types.utilsrH   rI   ImportErrorLITELLM_INSTALL_MSG_body_to_messages
model_dumpr   listchoicesmessage
tool_callsrO   r   r5   r6   r	   rC   jsonloads	argumentsr
   id)r&   
model_namecallerr   rH   rI   erK   r<   mmessage_dictst
tool_dictsresponserX   rZ   r=   rB   funcfunction_argsfunction_responses                        r/   r   z"CompletionExtension._do_completionV   s   	:*B 00:%19:A::.34alln4<
z
]bl'mntG}h&6&67QZ''22
1:%%---WQZ//::<=# 	I%%**Bsm#**5)<D JJy'9'9'C'CDM $ 5} 5  !*FASASAXAXbs*,	 z
]bl'mntG}h&6&67qz!!)))?  	:129	: ;4s"   G< H
H"<	HHHc                   K   	 ddl m} ddlm}m} | j                   |             \  }}|D 	cg c]  }	|	j                          }
}	|D cg c]  }|j                          c}xs d}t        | |||
|       d{         }t        t        |   |j                        }|d   j                  j                  xs g }|s|d   j                  j                  S |
j                  |d   j                  j                                |D ]  }|j                   j"                  sd}t%        |      | j'                  ||      }|
j                  t)        |j*                  d|j                   j"                   |d	i t-        j.                  |j                   j0                              j                                 t        | |||
|       d{         }t        t        |   |j                        }|d   j                  j                  S # t
        $ r}t        t              |d}~ww xY wc c}	w c c}w 7 7 lw)
rF   r   )acompletionrG   NrJ   zFunction name is emptyr>   rL   rP   )rQ   rk   rR   rH   rI   rS   rT   rU   rV   r   rW   rX   rY   rZ   rO   r   r5   r6   r	   rC   r
   r^   r[   r\   r]   )r&   r_   r`   rk   rH   rI   ra   rK   r<   rb   rc   rd   re   rf   rX   rZ   r=   rB   rg   s                      r/   r   z(CompletionExtension._do_completion_async~   s    	:+B 00:%19:A::.34alln4<
[zTais-t'tutG}h&6&67QZ''228b
1:%%---WQZ//::<=# 	I%%**.sm#**5)<D  !*"++00 L4::i.@.@.J.J#KL	
 *,	  [zTais-t'tutG}h&6&67qz!!)))A  	:129	: ;4't. (us]   H>H H>H/H>H4#H>>H9
?EH>H<
>H>	H,H''H,,H><H>r.   returnc                 N   |j                         }g }g }|j                  d      D ]'  }	 |j                  t        j                  |             ) |sd}t        |      ||fS # t
        $ r8 	 |j                  t        j                  |             n# t
        $ r Y nw xY wY }w xY w)z>Converts each line in the body of a block into a chat message.
z1Completion must contain at least one chat message)stripsplitr   r
   model_validate_jsonr   r   r   )r&   r.   rK   r<   linerB   s         r/   rU   z%CompletionExtension._body_to_messages   s    zz|JJt$ 
	D	 ? ? EF
	 EC$S))%   # LL!9!9$!?@& 	s5   $A##	B$-$BB$	BB$BB$#B$N)__name__
__module____qualname____doc__tagsr0   rC   r   r   strtuplerW   r
   r   rU   rP   rD   r/   r   r      sN    ( >D!jF&*P'*R!c !eD4EtDz4Q.R !rD   r   )r9   r[   typingr   jinja2r   r   
jinja2.extr   pydanticr   banks.errorsr   r	   banks.typesr
   r   r   rT   r   rP   rD   r/   <module>r      s9       -   $ 5 ) _ i!) i!rD   