o
    Zh/                     @  s  d dl mZ 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 d d	lmZ d d
lmZ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# d dl$m%Z%m&Z&m'Z' d dl(m)Z)m*Z*m+Z+m,Z,m-Z-m.Z.m/Z/ d dl0m1Z1 d dl2m3Z3m4Z4m5Z5 d dl6m7Z7m8Z8m9Z9m:Z: ej;rd dl<m%Z= d dl>m%Z? e@eAZBejCejDeEf ZFe/de.de-diZGeG dd dZHeG dd dZIdS )    )annotationsN)	dataclass)choices)Dataset)
Embeddings)BaseLanguageModel)
ChatOpenAI)OpenAIEmbeddings)TestsetGenerationEventtrack)BaseRagasEmbeddingsLangchainEmbeddingsWrapper)ExceptionInRunner)Executor)BaseRagasLLMLangchainLLMWrapper)	RunConfig)DocumentDocumentStoreInMemoryDocumentStore)ComplexEvolutionCurrentNodesDataRow	Evolutionmulti_context	reasoningsimple)KeyphraseExtractor)EvolutionFilter
NodeFilterQuestionFilter)check_if_sum_is_close
deprecatedget_feature_languageis_nan)r   g      ?g      ?c                   @  s8   e Zd ZU dZded< dddZdd	d
ZdddZdS )TestDatasetz
    TestDataset class
    zt.List[DataRow]	test_datareturnt.List[t.Dict]c                 C  s.   g }| j D ]}t|}d|d< || q|S )NTZepisode_done)r&   dictappend)selfZdata_samplesdataZ	data_dict r-   N/var/www/html/lang_env/lib/python3.10/site-packages/ragas/testset/generator.py_to_records5   s   
zTestDataset._to_recordspd.DataFramec                 C  s   t j|  S N)pdZ	DataFrameZfrom_recordsr/   r+   r-   r-   r.   	to_pandas=   s   zTestDataset.to_pandasr   c                 C  s   t |  S r1   )r   	from_listr/   r3   r-   r-   r.   
to_dataset@   s   zTestDataset.to_datasetN)r'   r(   )r'   r0   )r'   r   )__name__
__module____qualname____doc____annotations__r/   r4   r6   r-   r-   r-   r.   r%   -   s   
 

r%   c                   @  s   e Zd ZU ded< ded< ded< ded< e				d>d?ddZeedddd						d@dAddZ			 	 	dBdCd)d*Z			 	 	dBdDd,d-Z	dEd1d2Z
			 	 	dBdFd3d4Z	dGdHd:d;Z	dGdId<d=ZdS )JTestsetGeneratorr   generator_llm
critic_llmr   
embeddingsr   docstoreN   r   r   t.Optional[DocumentStore]
run_configt.Optional[RunConfig]
chunk_sizeintr'   'TestsetGenerator'c                 C  sr   t |}t |}t|}	t|d}
|d u r1ddlm} ||dd}t||	|
|d}| |||	|dS | |||	|dS )NZllmr   )TokenTextSplitter)rE   Zchunk_overlap)splitterr?   	extractorrC   )r=   r>   r?   r@   )r   r   r   Zlangchain.text_splitterrI   r   )clsr=   r>   r?   r@   rC   rE   generator_llm_modelcritic_llm_modelembeddings_modelZkeyphrase_extractorrI   rJ   r-   r-   r.   from_langchainK   s2   

zTestsetGenerator.from_langchainz0.1.4z0.2.0rP   )Zremovalalternativegpt-3.5-turbo-16kgpt-4text-embedding-ada-002strc           	      C  s2   t |d}t |d}t|d}| j|||||dS )N)model)r=   r>   r?   r@   rE   )r   r	   rP   )	rL   r=   r>   r?   r@   rE   rM   rN   rO   r-   r-   r.   with_openair   s   



zTestsetGenerator.with_openaiFT	documentst.Sequence[LlamaindexDocument]	test_sizedistributionst.Optional[Distributions]is_asyncboolraise_exceptionsc                 C  s4   |pi }| j dd |D  | j||||||dS )Nc                 S     g | ]}t |qS r-   )r   Zfrom_llamaindex_document.0docr-   r-   r.   
<listcomp>       zBTestsetGenerator.generate_with_llamaindex_docs.<locals>.<listcomp>)rZ   r[   with_debugging_logsr]   rC   r_   r@   Zadd_documentsgenerater+   rX   rZ   r[   rf   r]   r_   rC   r-   r-   r.   generate_with_llamaindex_docs      
z.TestsetGenerator.generate_with_llamaindex_docst.Sequence[LCDocument]c                 C  s4   |pi }| j dd |D  | j||||||dS )Nc                 S  r`   r-   )r   Zfrom_langchain_documentra   r-   r-   r.   rd      re   zATestsetGenerator.generate_with_langchain_docs.<locals>.<listcomp>)rZ   r[   rf   r]   r_   rC   rg   ri   r-   r-   r.   generate_with_langchain_docs   rk   z-TestsetGenerator.generate_with_langchain_docs	evolutionr   Nonec                 C  s   |j d u r=| j |_ |jd u r| j|_|jd u rt| jd|_|jd u r*t| jd|_t|tr?|j	d u rAt
| jd|_	d S d S d S d S )NrH   )r=   r@   question_filterr    r>   node_filterr   
isinstancer   evolution_filterr   )r+   rn   r-   r-   r.   init_evolution   s   






zTestsetGenerator.init_evolutionc              
     sN   pt  tt  ddstdtt   d|d u r&tddd}| j|  D ]}| 	| |j
||d q.|rcd	d
lm} |dtj |dtj |dtj |dtj |dtj tdd||d}	dd | jj|dD }
d	}  D ]%\}}tt|| D ]}|	j|j|
| |jj d| d |d7 }qq}||krtt || d}|D ]}|	j|j|
| |jj d| d |d7 }qz|	 }|st W n ty } z|d }~ww dd |D }t|d}dd  D }dd |D }ttddd  D  fdd D t|jt|d	kr |d	 nd d! |S )"Ng      ?   z,distributions passed do not sum to 1.0 [got z"]. Please check the distributions.   Z   )max_retriesZmax_wait)r]   rC   r   )patch_loggerzragas.testset.evolutionszragas.testset.extractorzragas.testset.filterszragas.testset.docstorezragas.llms.promptZ
GeneratingT)ZdescZkeep_progress_barr_   rC   c                 S  s   g | ]	}t ||gd qS ))Z	root_nodenodes)r   )rb   nr-   r-   r.   rd      s    z-TestsetGenerator.generate.<locals>.<listcomp>)k-)name   c                 S  s   g | ]}t |s|qS r-   )r$   )rb   rr-   r-   r.   rd         )r&   c                 S  s   g | ]}t |qS r-   )r#   rb   er-   r-   r.   rd         c                 S  s   g | ]}|d ur|qS r1   r-   r   r-   r-   r.   rd     r   Ztestset_generationc                 S  s   g | ]}|j j qS r-   )	__class__r7   lowerr   r-   r-   r.   rd     r   c                   s   g | ]} | qS r-   r-   r   r[   r-   r.   rd      r    )Z
event_typeZevolution_namesZevolution_percentagesZnum_rowslanguage) DEFAULT_DISTRIBUTIONr!   listvalues
ValueErrorsumr   r@   Zset_run_configrt   initragas.utilsry   loggingDEBUGr   Zget_random_nodesitemsrangeroundZsubmitevolver   r7   r   resultsr   r%   r   r
   lenr&   )r+   rZ   r[   rf   r]   r_   rC   rn   ry   execZcurrent_nodesZtotal_evolutionsZprobabilityiZfiller_evolutionsZtest_data_rowsr   Ztest_datasetZ	evol_langr-   r   r.   rh      s   	




zTestsetGenerator.generater   
evolutionst.List[Evolution]	cache_dirt.Optional[str]c                 C  sh   t | jts
J d| jjd usJ d| jjj||d |D ]}| | |  |j||d qd S )N)Must be an instance of in-memory docstoreExtractor is not setr   )rr   r@   r   rK   adaptrt   r   )r+   r   r   r   rn   r-   r-   r.   r   (  s   
zTestsetGenerator.adaptc                 C  s   t | jts
J d| jjdusJ d| jj| |D ](}|jdus(J d|jdus1J dt |tr?|jdus?J d|j|d qdS )z6
        Save the docstore prompts to a path.
        r   Nr   zNodeFilter is not setzQuestionFilter is not setzEvolutionFilter is not setr   )	rr   r@   r   rK   saverq   rp   r   rs   )r+   r   r   rn   r-   r-   r.   r   9  s    
zTestsetGenerator.save)NNrA   )r=   r   r>   r   r?   r   r@   rB   rC   rD   rE   rF   r'   rG   )rR   rS   rT   NrA   )r=   rU   r>   rU   r?   rU   r@   rB   rE   rF   r'   rG   )NFTTN)rX   rY   rZ   rF   r[   r\   r]   r^   r_   r^   rC   rD   )rX   rl   rZ   rF   r[   r\   r]   r^   r_   r^   rC   rD   )rn   r   r'   ro   )
rZ   rF   r[   r\   r]   r^   r_   r^   rC   rD   r1   )r   rU   r   r   r   r   r'   ro   )r   r   r   r   r'   ro   )r7   r8   r9   r;   classmethodrP   r"   rW   rj   rm   rt   rh   r   r   r-   r-   r-   r.   r<   D   sR   
 &
ar<   )J
__future__r   r   typingtdataclassesr   randomr   Zpandasr2   Zdatasetsr   Zlangchain_core.embeddingsr   Zlangchain_core.language_modelsr   Zlangchain_openai.chat_modelsr   Zlangchain_openai.embeddingsr	   Zragas._analyticsr
   r   Zragas.embeddings.baser   r   Zragas.exceptionsr   Zragas.executorr   Z
ragas.llmsr   r   Zragas.run_configr   Zragas.testset.docstorer   r   r   Zragas.testset.evolutionsr   r   r   r   r   r   r   Zragas.testset.extractorr   Zragas.testset.filtersr   r   r    r   r!   r"   r#   r$   TYPE_CHECKINGZlangchain_core.documentsZ
LCDocumentZllama_index.core.schemaZLlamaindexDocument	getLoggerr7   loggerDictAnyfloatZDistributionsr   r%   r<   r-   r-   r-   r.   <module>   s@    $	
