@ -1,7 +1,7 @@
from typing import Any , Dict
import pytest
from test_utils import MockEncoder
from test_utils import MockEncoder , MockEncoderReturnsList
import langchain . chains . rl_chain . base as rl_chain
import langchain . chains . rl_chain . pick_best_chain as pick_best_chain
@ -26,7 +26,9 @@ def test_multiple_ToSelectFrom_throws() -> None:
chain = pick_best_chain . PickBest . from_llm (
llm = llm ,
prompt = PROMPT ,
feature_embedder = pick_best_chain . PickBestFeatureEmbedder ( auto_embed = False , model = MockEncoder ( ) ) ,
feature_embedder = pick_best_chain . PickBestFeatureEmbedder (
auto_embed = False , model = MockEncoder ( )
) ,
)
actions = [ " 0 " , " 1 " , " 2 " ]
with pytest . raises ( ValueError ) :
@ -43,7 +45,9 @@ def test_missing_basedOn_from_throws() -> None:
chain = pick_best_chain . PickBest . from_llm (
llm = llm ,
prompt = PROMPT ,
feature_embedder = pick_best_chain . PickBestFeatureEmbedder ( auto_embed = False , model = MockEncoder ( ) ) ,
feature_embedder = pick_best_chain . PickBestFeatureEmbedder (
auto_embed = False , model = MockEncoder ( )
) ,
)
actions = [ " 0 " , " 1 " , " 2 " ]
with pytest . raises ( ValueError ) :
@ -56,7 +60,9 @@ def test_ToSelectFrom_not_a_list_throws() -> None:
chain = pick_best_chain . PickBest . from_llm (
llm = llm ,
prompt = PROMPT ,
feature_embedder = pick_best_chain . PickBestFeatureEmbedder ( auto_embed = False , model = MockEncoder ( ) ) ,
feature_embedder = pick_best_chain . PickBestFeatureEmbedder (
auto_embed = False , model = MockEncoder ( )
) ,
)
actions = { " actions " : [ " 0 " , " 1 " , " 2 " ] }
with pytest . raises ( ValueError ) :
@ -75,7 +81,9 @@ def test_update_with_delayed_score_with_auto_validator_throws() -> None:
llm = llm ,
prompt = PROMPT ,
selection_scorer = rl_chain . AutoSelectionScorer ( llm = auto_val_llm ) ,
feature_embedder = pick_best_chain . PickBestFeatureEmbedder ( auto_embed = False , model = MockEncoder ( ) ) ,
feature_embedder = pick_best_chain . PickBestFeatureEmbedder (
auto_embed = False , model = MockEncoder ( )
) ,
)
actions = [ " 0 " , " 1 " , " 2 " ]
response = chain . run (
@ -98,7 +106,9 @@ def test_update_with_delayed_score_force() -> None:
llm = llm ,
prompt = PROMPT ,
selection_scorer = rl_chain . AutoSelectionScorer ( llm = auto_val_llm ) ,
feature_embedder = pick_best_chain . PickBestFeatureEmbedder ( auto_embed = False , model = MockEncoder ( ) ) ,
feature_embedder = pick_best_chain . PickBestFeatureEmbedder (
auto_embed = False , model = MockEncoder ( )
) ,
)
actions = [ " 0 " , " 1 " , " 2 " ]
response = chain . run (
@ -121,7 +131,9 @@ def test_update_with_delayed_score() -> None:
llm = llm ,
prompt = PROMPT ,
selection_scorer = None ,
feature_embedder = pick_best_chain . PickBestFeatureEmbedder ( auto_embed = False , model = MockEncoder ( ) ) ,
feature_embedder = pick_best_chain . PickBestFeatureEmbedder (
auto_embed = False , model = MockEncoder ( )
) ,
)
actions = [ " 0 " , " 1 " , " 2 " ]
response = chain . run (
@ -153,7 +165,9 @@ def test_user_defined_scorer() -> None:
llm = llm ,
prompt = PROMPT ,
selection_scorer = CustomSelectionScorer ( ) ,
feature_embedder = pick_best_chain . PickBestFeatureEmbedder ( auto_embed = False , model = MockEncoder ( ) ) ,
feature_embedder = pick_best_chain . PickBestFeatureEmbedder (
auto_embed = False , model = MockEncoder ( )
) ,
)
actions = [ " 0 " , " 1 " , " 2 " ]
response = chain . run (
@ -168,7 +182,9 @@ def test_user_defined_scorer() -> None:
@pytest.mark.requires ( " vowpal_wabbit_next " , " sentence_transformers " )
def test_everything_embedded ( ) - > None :
llm , PROMPT = setup ( )
feature_embedder = pick_best_chain . PickBestFeatureEmbedder ( auto_embed = False , model = MockEncoder ( ) )
feature_embedder = pick_best_chain . PickBestFeatureEmbedder (
auto_embed = False , model = MockEncoder ( )
)
chain = pick_best_chain . PickBest . from_llm (
llm = llm , prompt = PROMPT , feature_embedder = feature_embedder , auto_embed = False
)
@ -200,7 +216,9 @@ def test_everything_embedded() -> None:
@pytest.mark.requires ( " vowpal_wabbit_next " , " sentence_transformers " )
def test_default_auto_embedder_is_off ( ) - > None :
llm , PROMPT = setup ( )
feature_embedder = pick_best_chain . PickBestFeatureEmbedder ( auto_embed = False , model = MockEncoder ( ) )
feature_embedder = pick_best_chain . PickBestFeatureEmbedder (
auto_embed = False , model = MockEncoder ( )
)
chain = pick_best_chain . PickBest . from_llm (
llm = llm , prompt = PROMPT , feature_embedder = feature_embedder
)
@ -224,9 +242,11 @@ def test_default_auto_embedder_is_off() -> None:
@pytest.mark.requires ( " vowpal_wabbit_next " , " sentence_transformers " )
def test_default_ embeddings_off( ) - > None :
def test_default_ w_ embeddings_off( ) - > None :
llm , PROMPT = setup ( )
feature_embedder = pick_best_chain . PickBestFeatureEmbedder ( auto_embed = False , model = MockEncoder ( ) )
feature_embedder = pick_best_chain . PickBestFeatureEmbedder (
auto_embed = False , model = MockEncoder ( )
)
chain = pick_best_chain . PickBest . from_llm (
llm = llm , prompt = PROMPT , feature_embedder = feature_embedder , auto_embed = False
)
@ -250,29 +270,54 @@ def test_default_embeddings_off() -> None:
@pytest.mark.requires ( " vowpal_wabbit_next " , " sentence_transformers " )
def test_default_ embeddings_mixed_ w_explicit_user_e mbeddings( ) - > None :
def test_default_ w_embeddings_on ( ) - > None :
llm , PROMPT = setup ( )
feature_embedder = pick_best_chain . PickBestFeatureEmbedder ( auto_embed = True , model = MockEncoder ( ) )
feature_embedder = pick_best_chain . PickBestFeatureEmbedder (
auto_embed = True , model = MockEncoderReturnsList ( )
)
chain = pick_best_chain . PickBest . from_llm (
llm = llm , prompt = PROMPT , feature_embedder = feature_embedder , auto_embed = True
)
str1 = " 0 "
str2 = " 1 "
str3 = " 2 "
encoded_str1 = rl_chain . stringify_embedding ( list ( encoded_keyword + str1 ) )
encoded_str2 = rl_chain . stringify_embedding ( list ( encoded_keyword + str2 ) )
e ncoded_str3 = rl_chain . stringify_embedding ( list ( encoded_keyword + str3 ) )
ctx_str_1 = " context1 "
dot_prod = " dotprod 0:5.0 " # dot prod of [1.0, 2.0] and [1.0, 2.0]
e xpected = f """ shared |User { ctx_str_1 } |@ User= { ctx_str_1 } \n |action { str1 } |# action= { str1 } | { dot_prod } \n |action { str2 } |# action= { str2 } | { dot_prod } """ # noqa
actions = [ str1 , str2 ]
response = chain . run (
User = rl_chain . BasedOn ( ctx_str_1 ) ,
action = rl_chain . ToSelectFrom ( actions ) ,
)
selection_metadata = response [ " selection_metadata " ]
vw_str = feature_embedder . format ( selection_metadata )
assert vw_str == expected
@pytest.mark.requires ( " vowpal_wabbit_next " , " sentence_transformers " )
def test_default_embeddings_mixed_w_explicit_user_embeddings ( ) - > None :
llm , PROMPT = setup ( )
feature_embedder = pick_best_chain . PickBestFeatureEmbedder (
auto_embed = True , model = MockEncoderReturnsList ( )
)
chain = pick_best_chain . PickBest . from_llm (
llm = llm , prompt = PROMPT , feature_embedder = feature_embedder , auto_embed = True
)
str1 = " 0 "
str2 = " 1 "
encoded_str2 = rl_chain . stringify_embedding ( [ 1.0 , 2.0 ] )
ctx_str_1 = " context1 "
ctx_str_2 = " context2 "
encoded_ctx_str_1 = rl_chain . stringify_embedding ( [ 1.0 , 2.0 ] )
dot_prod = " dotprod 0:5.0 1:5.0 " # dot prod of [1.0, 2.0] and [1.0, 2.0]
encoded_ctx_str_1 = rl_chain . stringify_embedding ( list ( encoded_keyword + ctx_str_1 ) )
encoded_ctx_str_2 = rl_chain . stringify_embedding ( list ( encoded_keyword + ctx_str_2 ) )
expected = f """ shared |User { encoded_ctx_str_1 } |User2 { ctx_str_2 + " " + encoded_ctx_str_2 } \n |action { str1 + " " + encoded_str1 } \n |action { str2 + " " + encoded_str2 } \n |action { encoded_str3 } """ # noqa
expected = f """ shared |User { encoded_ctx_str_1 } |@ User= { encoded_ctx_str_1 } |User2 { ctx_str_2 } |@ User2= { ctx_str_2 } \n |action { str1 } |# action= { str1 } | { dot_prod } \n |action { encoded_str2 } |# action= { encoded_str2 } | { dot_prod } """ # noqa
actions = [ str1 , str2 , rl_chain . Embed ( str3 ) ]
actions = [ str1 , rl_chain . Embed ( str2 ) ]
response = chain . run (
User = rl_chain . BasedOn ( rl_chain . Embed ( ctx_str_1 ) ) ,
@ -291,7 +336,9 @@ def test_default_no_scorer_specified() -> None:
chain = pick_best_chain . PickBest . from_llm (
llm = chain_llm ,
prompt = PROMPT ,
feature_embedder = pick_best_chain . PickBestFeatureEmbedder ( auto_embed = False , model = MockEncoder ( ) ) ,
feature_embedder = pick_best_chain . PickBestFeatureEmbedder (
auto_embed = False , model = MockEncoder ( )
) ,
)
response = chain . run (
User = rl_chain . BasedOn ( " Context " ) ,
@ -310,7 +357,9 @@ def test_explicitly_no_scorer() -> None:
llm = llm ,
prompt = PROMPT ,
selection_scorer = None ,
feature_embedder = pick_best_chain . PickBestFeatureEmbedder ( auto_embed = False , model = MockEncoder ( ) ) ,
feature_embedder = pick_best_chain . PickBestFeatureEmbedder (
auto_embed = False , model = MockEncoder ( )
) ,
)
response = chain . run (
User = rl_chain . BasedOn ( " Context " ) ,
@ -330,7 +379,9 @@ def test_auto_scorer_with_user_defined_llm() -> None:
llm = llm ,
prompt = PROMPT ,
selection_scorer = rl_chain . AutoSelectionScorer ( llm = scorer_llm ) ,
feature_embedder = pick_best_chain . PickBestFeatureEmbedder ( auto_embed = False , model = MockEncoder ( ) ) ,
feature_embedder = pick_best_chain . PickBestFeatureEmbedder (
auto_embed = False , model = MockEncoder ( )
) ,
)
response = chain . run (
User = rl_chain . BasedOn ( " Context " ) ,
@ -348,7 +399,9 @@ def test_calling_chain_w_reserved_inputs_throws() -> None:
chain = pick_best_chain . PickBest . from_llm (
llm = llm ,
prompt = PROMPT ,
feature_embedder = pick_best_chain . PickBestFeatureEmbedder ( auto_embed = False , model = MockEncoder ( ) ) ,
feature_embedder = pick_best_chain . PickBestFeatureEmbedder (
auto_embed = False , model = MockEncoder ( )
) ,
)
with pytest . raises ( ValueError ) :
chain . run (
@ -371,7 +424,9 @@ def test_activate_and_deactivate_scorer() -> None:
llm = llm ,
prompt = PROMPT ,
selection_scorer = pick_best_chain . base . AutoSelectionScorer ( llm = scorer_llm ) ,
feature_embedder = pick_best_chain . PickBestFeatureEmbedder ( auto_embed = False , model = MockEncoder ( ) ) ,
feature_embedder = pick_best_chain . PickBestFeatureEmbedder (
auto_embed = False , model = MockEncoder ( )
) ,
)
response = chain . run (
User = pick_best_chain . base . BasedOn ( " Context " ) ,