o
    Zh|                     @  s   d Z ddlmZ ddlmZmZmZmZ ddlm	Z	 ddl
mZ ddl
mZ ddlmZ ddlmZ dddZdddZdddZdddZG dd dejZG dd dejZdS )zSheetsInputs.    )annotations)AnyCallableMappingSequence)parse)gspread_client)	sheets_id)llmfn_inputs_source)llmfn_outputsvaluestrreturn!sheets_id.SheetsIdentifier | Nonec                 C  sR   zt | }W n
 ty   Y dS w |jr'tjt| d}t 	| |S dS )z4Try to open a Sheets document with `value` as a URL.N)url)
r   urlparse
ValueErrorschemer	   SheetsIdentifierZ	SheetsURLr   
get_clientvalidate)r   Zparse_resultsid r   `/var/www/html/lang_env/lib/python3.10/site-packages/google/generativeai/notebook/sheets_utils.py_try_sheet_id_as_url   s   r   c                 C  sX   zt jt | d}W n
 ty   Y dS w z
t | W |S  tjy+   Y dS w )z4Try to open a Sheets document with `value` as a key.)keyN)r	   r   Z	SheetsKeyr   r   r   r   SpreadsheetNotFoundErrorr   r   r   r   r   _try_sheet_id_as_key-   s   r   c                 C  s8   t j| d}z
t | W |S  tjy   Y dS w )z5Try to open a Sheets document with `value` as a name.)nameN)r	   r   r   r   r   r   r   r   r   r   _try_sheet_id_as_name<   s   r    sheets_id.SheetsIdentifierc                 C  s>   t |  }r|S t|  }r|S t|  }r|S td| )Nz-No Sheets found with "{}" as URL, key or name)r   r   r    RuntimeErrorformatr   r   r   r   get_sheets_id_from_strF   s   r$   c                      s.   e Zd ZdZdd fddZdddZ  ZS )SheetsInputsz,Inputs to an LLMFunction from Google Sheets.r   r   r!   worksheet_idintc                   s   t    || _|| _d S N)super__init___sid_worksheet_id)selfr   r&   	__class__r   r   r*   S   s   

zSheetsInputs.__init__r   6tuple[Sequence[Mapping[str, str]], Callable[[], None]]c                 C  s   t  j| j| jdS )N)r   r&   )r   r   Zget_all_recordsr+   r,   )r-   r   r   r   _to_normalized_inputs_implX   s   z'SheetsInputs._to_normalized_inputs_impl)r   )r   r!   r&   r'   )r   r0   )__name__
__module____qualname____doc__r*   r1   __classcell__r   r   r.   r   r%   P   s    r%   c                   @  s$   e Zd ZdZdddZdd
dZdS )SheetsOutputsz4Writes outputs from an LLMFunction to Google Sheets.r   r!   c                 C  s
   || _ d S r(   )r+   )r-   r   r   r   r   r*   c   s   
zSheetsOutputs.__init__outputsllmfn_outputs.LLMFnOutputsBaser   Nonec                 C  sJ   |  }t| g}|dd t|  D  t j| j	|d d S )Nc                 S  s   g | ]}t |qS r   )list).0xr   r   r   
<listcomp>j   s    z/SheetsOutputs.write_outputs.<locals>.<listcomp>)r   rows)
as_dictr;   keysextendzipvaluesr   r   Zwrite_recordsr+   )r-   r8   Zoutputs_dictZoutputs_rowsr   r   r   write_outputsf   s   
zSheetsOutputs.write_outputsN)r   r!   )r8   r9   r   r:   )r2   r3   r4   r5   r*   rE   r   r   r   r   r7   `   s    
r7   N)r   r   r   r   )r   r   r   r!   )r5   
__future__r   typingr   r   r   r   urllibr   Zgoogle.generativeai.notebookr   r	   Z google.generativeai.notebook.libr
   r   r   r   r    r$   ZLLMFnInputsSourcer%   ZLLMFnOutputsSinkr7   r   r   r   r   <module>   s   





