
    Ng8                        d Z ddlZddlZddlmZ ddlmZmZmZ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  ej        e          Z ej        ej                  Z ej        d
          Ze                    e           e                     e           defdZ! G d d          Z" G d de          Z#dS )aH  
UpTrain Callback Handler

UpTrain is an open-source platform to evaluate and improve LLM applications. It provides
grades for 20+ preconfigured checks (covering language, code, embedding use cases),
performs root cause analyses on instances of failure cases and provides guidance for
resolving them.

This module contains a callback handler for integrating UpTrain seamlessly into your
pipeline and facilitating diverse evaluations. The callback handler automates various
evaluations to assess the performance and effectiveness of the components within the
pipeline.

The evaluations conducted include:

1. RAG:
   - Context Relevance: Determines the relevance of the context extracted from the query
   to the response.
   - Factual Accuracy: Assesses if the Language Model (LLM) is providing accurate
   information or hallucinating.
   - Response Completeness: Checks if the response contains all the information
   requested by the query.

2. Multi Query Generation:
   MultiQueryRetriever generates multiple variants of a question with similar meanings
   to the original question. This evaluation includes previous assessments and adds:
   - Multi Query Accuracy: Ensures that the multi-queries generated convey the same
   meaning as the original query.

3. Context Compression and Reranking:
   Re-ranking involves reordering nodes based on relevance to the query and selecting
   top n nodes.
   Due to the potential reduction in the number of nodes after re-ranking, the following
   evaluations
   are performed in addition to the RAG evaluations:
   - Context Reranking: Determines if the order of re-ranked nodes is more relevant to
   the query than the original order.
   - Context Conciseness: Examines whether the reduced number of nodes still provides
   all the required information.

These evaluations collectively ensure the robustness and effectiveness of the RAG query
engine, MultiQueryRetriever, and the re-ranking process within the pipeline.

Useful links:
Github: https://github.com/uptrain-ai/uptrain
Website: https://uptrain.ai/
Docs: https://docs.uptrain.ai/getting-started/introduction

    N)defaultdict)AnyDefaultDictDictListOptionalSequenceSet)UUID)BaseCallbackHandler)Document)	LLMResultguard_importz%(message)sreturnc                       t          d          S )zImport the `uptrain` package.uptrainr        j/var/www/html/ai-engine/env/lib/python3.11/site-packages/langchain_community/callbacks/uptrain_callback.pyimport_uptrainr   M   s    	"""r   c                   "    e Zd ZdZdeddfdZdS )UpTrainDataSchemaa-  The UpTrain data schema for tracking evaluation results.

    Args:
        project_name (str): The project name to be shown in UpTrain dashboard.

    Attributes:
        project_name (str): The project name to be shown in UpTrain dashboard.
        uptrain_results (DefaultDict[str, Any]): Dictionary to store evaluation results.
        eval_types (Set[str]): Set to store the types of evaluations.
        query (str): Query for the RAG evaluation.
        context (str): Context for the RAG evaluation.
        response (str): Response for the RAG evaluation.
        old_context (List[str]): Old context nodes for Context Conciseness evaluation.
        new_context (List[str]): New context nodes for Context Conciseness evaluation.
        context_conciseness_run_id (str): Run ID for Context Conciseness evaluation.
        multi_queries (List[str]): List of multi queries for Multi Query evaluation.
        multi_query_run_id (str): Run ID for Multi Query evaluation.
        multi_query_daugher_run_id (str): Run ID for Multi Query daughter evaluation.

    project_namer   Nc                 >   || _         t          t                    | _        t	                      | _        d| _        d| _        d| _        g | _	        g | _
        t          d          | _        g | _        t          d          | _        t          d          | _        dS )z#Initialize the UpTrain data schema. r   )intN)r   r   listuptrain_resultsset
eval_typesquerycontextresponseold_contextnew_contextr   context_conciseness_run_idmulti_queriesmulti_query_run_idmulti_query_daugher_run_id)selfr   s     r   __init__zUpTrainDataSchema.__init__h   s     ".6A$6G6G %(EE 
 ')&(04' )+(,04'''r   )__name__
__module____qualname____doc__strr,   r   r   r   r   r   R   s@         *<S <T < < < < < <r   r   c                   *    e Zd ZdZddddddded	ed
edededdf fdZdedeeee	f                  dee         ddfdZ
dddededee         de	ddf
dZdddddddeee	f         deee	f         dedeee                  dee         deeee	f                  dee         d ee         de	ddfd!Zdddd"deee	f         d#ededee         deee                  deeee	f                  de	ddfd$Zddd%ee         dedee         de	de	f
d&Z xZS )'UpTrainCallbackHandlera  Callback Handler that logs evaluation results to uptrain and the console.

    Args:
        project_name (str): The project name to be shown in UpTrain dashboard.
        key_type (str): Type of key to use. Must be 'uptrain' or 'openai'.
        api_key (str): API key for the UpTrain or OpenAI API.
        (This key is required to perform evaluations using GPT.)

    Raises:
        ValueError: If the key type is invalid.
        ImportError: If the `uptrain` package is not installed.

    	langchainopenaizsk-****************zgpt-3.5-turboT)r   key_typeapi_keymodellog_resultsr   r6   r7   r8   r9   r   Nc                   t                                                       t                      }|| _        t	          |          | _        d| _        |dk    r4|                    ||          }|                    |          | _	        d
S |dk    r5|                    |d|          }|
                    |          | _	        d
S t          d	          )z)Initializes the `UpTrainCallbackHandler`.)r   Fr   )uptrain_access_tokenr8   )settingsr5   T)openai_api_keyevaluate_locallyr8   z/Invalid key type: Must be 'uptrain' or 'openai'N)superr,   r   r9   r   schemafirst_score_printed_flagSettings	APIClientuptrain_clientEvalLLM
ValueError)	r+   r   r6   r7   r8   r9   r   r<   	__class__s	           r   r,   zUpTrainCallbackHandler.__init__   s     	 ""& (\BBB(-%y  ''WE'RRH")"3"3X"3"F"FD!!''&U (  H #*//8/"D"DDNOOOr   evaluation_namedatachecksc                    | j         j        j        dk    r)| j                             | j        j        |||          }n(| j                             | j        j        |||          }| j        j        | j        j                                     |           ddddddd	d
}| j	        r$t                              t          j                   |D ]p}t          |                                          }|D ]H}|dk    r+t                              d||                     d| _        4|dk    r+t                              d||                     d| _        e|dk    rJt                              d           ||         D ]}	t                              d|	             d| _        |                    d          r~| j        s!t                              d           d| _        ||v r-t                              ||          d||                     #t                              | d||                     Jr| j	        r&t                              t          j                   dS dS )z=Run an evaluation on the UpTrain server using UpTrain client.rC   )r   rH   rI   rJ   zContext Relevance ScorezFactual Accuracy ScorezResponse Completeness ScorezSub Query Completeness ScorezContext Reranking ScorezContext Conciseness ScorezMulti Query Accuracy Score)score_context_relevancescore_factual_accuracyscore_response_completenessscore_sub_query_completenessscore_context_rerankingscore_context_concisenessscore_multi_query_accuracyquestionz
Question: Fr$   z
Response: variantszMulti Queries:z  - scorer   Tz: N)rD   rG   r-   log_and_evaluater@   r   evaluater   appendr9   loggersetLevelloggingINFOr   keysinforA   
startswithWARNING)
r+   rH   rI   rJ   uptrain_resultscore_name_maprowcolumnscolumnvariants
             r   uptrain_evaluatez'UpTrainCallbackHandler.uptrain_evaluate   s    (1[@@!0AA![5 /	 B  NN "099![5 /	 :  N 	#DK$<=DD^TTT (A&>+H,J'@)D*F
 
  	*OOGL)))! 	@ 	@C388::&&G! @ @Z''KK <s6{ < <===49D11z))KK :S[ : :;;;49D11z))KK 0111#&v; 6 6$47$4$4555549D11&&w// @8 =B8<5//~f'=$N$NV$N$NOOOOv$>$>V$>$>???'@*  	- OOGO,,,,,	- 	-r   )parent_run_idr$   run_idrh   kwargsc                f   t                      }|j        d         d         j        | j        _        d| j        j        v rp|| j        j        k    rb| j        j        | j        j        | j        j        dg}| 	                    d||j
        j        |j
        j        |j
        j        g           dS dS dS )z(Log records to uptrain when an LLM ends.r   qa_rag)rS   r#   r$   ragrH   rI   rJ   N)r   generationstextr@   r$   r!   r*   r"   r#   rg   EvalsCONTEXT_RELEVANCEFACTUAL_ACCURACYRESPONSE_COMPLETENESS)r+   r$   ri   rh   rj   r   rI   s          r   
on_llm_endz!UpTrainCallbackHandler.on_llm_end   s     !""'3A6q9>...!GGG !% 1#{2 $ 4 D !! %M3M2M7 "      /.GGr   )tagsrh   metadatarun_typename
serializedinputsrv   rw   rx   ry   c                4   || j         j        k    r|| j         _        t          |t                    rt          |                                          ddhk    r| j         j                            d           d}
t          |d         t                    r|d         j
        }
nWt          |d         t                    r|d         D ]}|
|j
        dz   z  }
n#t          |d         t                    r|d         }
|
| j         _        |d         | j         _        dS )zDo nothing when chain startsr#   rS   rl   r   
N)r@   r)   r*   
isinstancedictr    r]   r!   addr   page_contentr   r1   r#   r"   )r+   rz   r{   ri   rv   rh   rw   rx   ry   rj   r#   docs               r   on_chain_startz%UpTrainCallbackHandler.on_chain_start  s    DK:::5;DK2fd## 	3FKKMM(:(:y*>U(U(UK"&&x000G&+X66 , +8F9-t44 ,!), 7 7Cs/$66GG7F9-s33 , +")DK &z 2DKr   )rh   rv   rw   r"   c                l   d|d         v r7| j         j                            d           || j         _        || j         _        d|d         v r9| j         j                            d           || j         _        || j         _        d S d| j         j        v r!| j         j                            |           d S d S )Ncontextual_compressionidmulti_query)r@   r!   r   r"   r'   r)   r(   rX   )r+   rz   r"   ri   rh   rv   rw   rj   s           r   on_retriever_startz)UpTrainCallbackHandler.on_retriever_start0  s     $z$'777K"&&'?@@@ %DK5;DK2Jt,,,K"&&}555-3DK* %DKdk444K%,,U33333 54r   	documentsc                *   t                      }|| j        j        k    r=| j        j        | j        j        dg}|                     d||j        j        g           d| j        j        v r&|| j        j	        k    r+|D ]&}| j        j
                            |j                   'dS || j        j	        k    r|D ]&}| j        j                            |j                   'd                    d t          | j        j
        d          D                       }d                    d	 t          | j        j        d          D                       }	| j        j        ||	|	d
g}|                     d||j        j        |j        j        g           dS dS dS )z Run when Retriever ends running.)rS   rT   r   rn   r   r}   c                 "    g | ]\  }}| d | S z. r   .0indexstrings      r   
<listcomp>z;UpTrainCallbackHandler.on_retriever_end.<locals>.<listcomp>f  :       )E6 !,,F,,  r      )startc                 "    g | ]\  }}| d | S r   r   r   s      r   r   z;UpTrainCallbackHandler.on_retriever_end.<locals>.<listcomp>l  r   r   )rS   r#   concise_contextreranked_contextcontext_rerankingN)r   r@   r)   r"   r(   rg   rq   MULTI_QUERY_ACCURACYr!   r'   r%   rX   r   r&   join	enumerateCONTEXT_CONCISENESSCONTEXT_RERANKING)
r+   r   ri   rh   rj   r   rI   r   r#   r   s
             r   on_retriever_endz'UpTrainCallbackHandler.on_retriever_endG  s    !""T[333 !% 1 $ 9 D !! -:; "   
 $t{'=== FFF$ E ECK+2233CDDDDE E4;AAA$ E ECK+2233CDDDD)) -6t{7NVW-X-X-X    $(99 -6t{7NVW-X-X-X  $ $  %)K$5#*+;,<	  %%$797 &     7 >= BAr   )r-   r.   r/   r0   r1   boolr,   r   r   r   rg   r   r   r   ru   r   r   r	   r   r   __classcell__)rG   s   @r   r3   r3      s	        " ( ,$ P P P P 	P
 P P P 
P P P P P P>?-?- 4S>"?- S		?-
 
?- ?- ?- ?-L )-   	
  ~  
   N %)(,-1"&"  cN S#X
  tCy!  ~ 4S>* 3- sm  
   J )-$(-14 4 4cN4 4
 4  ~4 tCy!4 4S>*4 4 
4 4 4 48 )-9 9 9H%9 	9
  ~9 9 
9 9 9 9 9 9 9 9r   r3   )$r0   r[   syscollectionsr   typingr   r   r   r   r   r	   r
   uuidr   langchain_core.callbacks.baser   langchain_core.documentsr   langchain_core.outputsr   langchain_core.utilsr   	getLoggerr-   rY   StreamHandlerstdouthandler	Formatter	formattersetFormatter
addHandlerr   r   r3   r   r   r   <module>r      s  0 0d  



 # # # # # #                        = = = = = = - - - - - - , , , , , , - - - - - -		8	$	$
'


+
+Gm,,	   Y      '   # # # # #
,< ,< ,< ,< ,< ,< ,< ,<^    0     r   