from model.chain_spec import ChainSpec, LLMChainSpec, SequentialChainSpec, CaseChainSpec, APIChainSpec, ReformatChainSpec, TransformChainSpec, VectorSearchChainSpec
from model.chain_revision import ChainRevision
from model.lang_chain_context import LangChainContext
from langchain.llms.fake import FakeListLLM
from model.tests.factory import SpecFactoryContext, llm_factory, sequential_factory, case_factory

def test_llm_chain_spec_serialization():
    llm_chain_spec = LLMChainSpec(
        chain_id=1,
        input_keys=["input1", "input2"],
        output_key="output1",
        prompt="prompt",
        llm_key="llm_key",
        chain_type="llm_chain_spec",
    )
    serialized = llm_chain_spec.json()
    assert serialized == '{"chain_id": 1, "input_keys": ["input1", "input2"], "output_key": "output1", "chain_type": "llm_chain_spec", "prompt": "prompt", "llm_key": "llm_key"}'

    revision = ChainRevision(chain=llm_chain_spec, llms={})
    serialized_revision = revision.json()
    deserialized = ChainRevision.parse_raw(serialized_revision).chain
    assert deserialized == llm_chain_spec


def test_llm_chain_spec_to_lang_chain_creates_valid_chain():
    prompt_template = "the prompt {input1} {input2}"
    llm_chain_spec = LLMChainSpec(
        chain_id=1,
        input_keys=["input1", "input2"],
        output_key="output1",
        prompt=prompt_template,
        llm_key="test",
        chain_type="llm_chain_spec",
    )
    llms = llms={"test": FakeListLLM(responses=["response1"])}
    ctx = LangChainContext(llms=llms)
    llm_chain = llm_chain_spec.to_lang_chain(ctx)

    assert llm_chain.input_keys == ["input1", "input2"]
    assert llm_chain.output_key == "output1"
    assert llm_chain.prompt.template == prompt_template
    assert llm_chain.llm == llms["test"]

    output = llm_chain._call({"input1": "input1", "input2": "input2"})
    assert output == {"output1": "response1"}


def test_llm_chain_spec_to_lang_chain_creates_recording_chain():
    prompt_template = "the prompt {input1} {input2}"
    llm_chain_spec = LLMChainSpec(
        chain_id=1,
        input_keys=["input1", "input2"],
        output_key="output1",
        prompt=prompt_template,
        llm_key="test",
        chain_type="llm_chain_spec",
    )
    llms = llms={"test": FakeListLLM(responses=["response1"])}
    ctx = LangChainContext(llms=llms, recording=True)
    llm_chain = llm_chain_spec.to_lang_chain(ctx)

    assert len(ctx.prompts) == 1

    output = llm_chain._call({"input1": "input1", "input2": "input2"})

    assert output == {"output1": "response1"}
    assert ctx.prompts[0].calls == [
        ({"input1": "input1", "input2": "input2"}, {"output1": "response1"})
    ]


def test_sequential_chain_spec_serialization():
    sequential_chain_spec = SequentialChainSpec(
        chain_id=1,
        input_keys=["input1", "input2"],
        output_keys=["output1", "output2"],
        chain_type="sequential_chain_spec",
        chains=[
            LLMChainSpec(
                chain_id=1,
                input_keys=["input1", "input2"],
                output_key="output1",
                prompt="prompt",
                llm_key="llm_key",
                chain_type="llm_chain_spec"
            )
        ],
    )
    serialized = sequential_chain_spec.json()
    deserialized = SequentialChainSpec.parse_raw(serialized)
    assert deserialized == sequential_chain_spec 


def test_sequential_chain_spec_to_lang_chain_creates_valid_chain():
    llm_chain_spec1 = LLMChainSpec(
                chain_id=1,
                input_keys=["input1", "input2"],
                output_key="llm1_output",
                prompt="the prompt {input1} {input2}",
                llm_key="test",
                chain_type="llm_chain_spec",
            )
    llm_chain_spec2 = LLMChainSpec(
                chain_id=2,
                input_keys=["llm1_output", "input3"],
                output_key="llm2_output",
                prompt="the prompt {llm1_output} {input3}",
                llm_key="test",
                chain_type="llm_chain_spec",
            )

    sequential_chain_spec = SequentialChainSpec(
        chain_id=3,
        input_keys=["input1", "input2", "input3"],
        output_keys=["llm2_output"],
        chain_type="sequential_chain_spec",
        chains=[llm_chain_spec1, llm_chain_spec2],
    )

    llms = llms={"test": FakeListLLM(responses=["fake_response1", "fake_response2"])}
    ctx = LangChainContext(llms=llms)
    sequential_chain = sequential_chain_spec.to_lang_chain(ctx)

    assert sequential_chain.input_keys == ["input1", "input2", "input3"]
    assert sequential_chain.output_keys == ["llm2_output"]
    assert sequential_chain.chains[0].input_keys == ["input1", "input2"]
    assert sequential_chain.chains[0].output_keys == ["llm1_output"]
    assert sequential_chain.chains[0].prompt.template == llm_chain_spec1.prompt
    assert sequential_chain.chains[0].llm == llms["test"]

    assert sequential_chain.chains[1].input_keys == ["llm1_output", "input3"]
    assert sequential_chain.chains[1].output_keys == ["llm2_output"]
    assert sequential_chain.chains[1].prompt.template == llm_chain_spec2.prompt
    assert sequential_chain.chains[1].llm == llms["test"]

    output = sequential_chain.run({"input1": "input1", "input2": "input2", "input3": "input3"})
    assert output == "fake_response2"


def test_case_chain_spec_serialization():
    case_chain_spec = CaseChainSpec(
        chain_id=3,
        chain_type="case_chain_spec",
        categorization_key="categorization_input",
        cases={
            "response1": LLMChainSpec(
                chain_id=1,
                input_keys=["input1", "input2"],
                output_key="output",
                prompt="prompt",
                llm_key="test",
                chain_type="llm_chain_spec"
            ),
            "response2": LLMChainSpec(
                chain_id=2,
                input_keys=["input3", "input4"],
                output_key="output",
                prompt="prompt",
                llm_key="test",
                chain_type="llm_chain_spec"
            )
        },
        default_case=LLMChainSpec(
            chain_id=4,
            input_keys=["input1", "input2"],
            output_key="output",
            prompt="prompt",
            llm_key="test",
            chain_type="llm_chain_spec"
        ),
    )
    serialized = case_chain_spec.json()
    deserialized = CaseChainSpec.parse_raw(serialized)
    assert deserialized == case_chain_spec


def test_case_chain_spec_to_lang_chain_creates_valid_chain():
    llms = {"test1": FakeListLLM(responses=["fake_response1"]), "test2": FakeListLLM(responses=["fake_response2"])}
    llm_chain_spec1 = LLMChainSpec(
                chain_id=1,
                input_keys=["input1", "input2"],
                output_key="output",
                prompt="ll1 prompt {input1} {input2}",
                llm_key="test1",
                chain_type="llm_chain_spec",
            )
    llm_chain_spec2 = LLMChainSpec(
                chain_id=2,
                input_keys=["input3", "input4"],
                output_key="output",
                prompt="llm2 prompt {input3} {input4}",
                llm_key="test2",
                chain_type="llm_chain_spec",
            )
    llm_chain_spec3 = LLMChainSpec(
                chain_id=4,
                input_keys=["input1", "input2"],
                output_key="output",
                prompt="default prompt {input1} {input2}",
                llm_key="test1",
                chain_type="llm_chain_spec"
            )

    case_chain_spec = CaseChainSpec(
        chain_id=3,
        chain_type="case_chain_spec",  
        categorization_key="categorization_input",      
        cases={
            "response1": llm_chain_spec1,
            "response2": llm_chain_spec2,
        },
        default_case=llm_chain_spec3,
    )

    ctx = LangChainContext(llms=llms)
    case_chain = case_chain_spec.to_lang_chain(ctx)

    assert case_chain.input_keys == ["categorization_input", "input1", "input2", "input3", "input4"]
    assert case_chain.output_keys == ["output"]
    assert case_chain.subchains["response1"].input_keys == ["input1", "input2"]
    assert case_chain.subchains["response1"].output_keys == ["output"]
    assert case_chain.subchains["response1"].prompt.template == llm_chain_spec1.prompt
    assert case_chain.subchains["response1"].llm == llms["test1"]

    assert case_chain.subchains["response2"].input_keys == ["input3", "input4"]
    assert case_chain.subchains["response2"].output_keys == ["output"]
    assert case_chain.subchains["response2"].prompt.template == llm_chain_spec2.prompt
    assert case_chain.subchains["response2"].llm == llms["test2"]

    assert case_chain.default_chain.input_keys == ["input1", "input2"]
    assert case_chain.default_chain.output_keys == ["output"]
    assert case_chain.default_chain.prompt.template == llm_chain_spec3.prompt
    assert case_chain.default_chain.llm == llms["test1"]

    output = case_chain.run({
        "categorization_input": "response2",
        "input1": "input1", 
        "input2": "input2",
        "input3": "input3",
        "input4": "input4"
    })
    assert output == "fake_response2"


def test_api_chain_spec_serialization():
    api_chain_spec = APIChainSpec(
        chain_id=3,
        chain_type="api_chain_spec",
        input_keys=["input1", "input2"],
        output_key="output",
        url="http://test.com",
        method="POST",
        body="body {input1} {input2}",
        headers={"header1": "header1"},
    )

    revision = ChainRevision(chain=api_chain_spec, llms={})
    serialized_revision = revision.json()
    deserialized = ChainRevision.parse_raw(serialized_revision).chain
    assert deserialized == api_chain_spec


def test_api_chain_spec_to_lang_chain_creates_valid_chain():
    api_chain_spec = APIChainSpec(
        chain_id=3,
        chain_type="api_chain_spec",
        input_keys=["input1", "input2"],
        output_key="output",
        url="http://test.com",
        method="POST",
        body="body {input1} {input2}",
        headers={"header1": "header1"},
    )
    ctx = LangChainContext(llms={})
    api_chain = api_chain_spec.to_lang_chain(ctx)

    assert api_chain.input_keys == ["input1", "input2"]
    assert api_chain.output_keys == ["output"]
    assert api_chain.url == "http://test.com"
    assert api_chain.method == "POST"
    assert api_chain.body == "body {input1} {input2}"
    assert api_chain.headers == {"header1": "header1"}


def test_reformat_chain_spec_serialization():
    reformat_chain_spec = ReformatChainSpec(
        chain_id=3,
        chain_type="reformat_chain_spec",
        input_keys=["input1", "input2"],
        formatters={"output1": "formatter1", "output2": "formatter2"},
    )
    serialized = reformat_chain_spec.json()
    deserialized = ReformatChainSpec.parse_raw(serialized)
    assert deserialized == reformat_chain_spec


def test_reformat_chain_spec_to_lang_chain_creates_valid_chain():
    reformat_chain_spec = ReformatChainSpec(
        chain_id=3,
        chain_type="reformat_chain_spec",
        input_keys=["input1", "input2"],
        formatters={"output1": "formatter1", "output2": "formatter2"},
    )
    ctx = LangChainContext(llms={})
    reformat_chain = reformat_chain_spec.to_lang_chain(ctx)

    assert reformat_chain.input_keys == ["input1", "input2"]
    assert reformat_chain.output_keys == ["output1", "output2"]
    assert reformat_chain.formatters == {"output1": "formatter1", "output2": "formatter2"}


def test_transform_chain_spec_to_lang_chain_creates_valid_chain():
    transform_chain_spec = TransformChainSpec(
        chain_id=4,
        chain_type="transform_chain_spec",
        input_keys=["x", "y"],
        output_keys=["z", "w"],
        transform_func="""
z = int(inputs['x']) + int(inputs['y'])
w = int(inputs['x']) - int(inputs['y'])
return {'z': z, 'w': w}"""
    )
    ctx = LangChainContext(llms={})
    transform_chain = transform_chain_spec.to_lang_chain(ctx)

    assert transform_chain.input_keys == ["x", "y"]
    assert transform_chain.output_keys == ["z", "w"]
    assert transform_chain.transform({'x': 4, 'y': 3}) == {'z': 7, 'w': 1}


def test_vector_search_chain_spec_serialization():
    vector_search_chain_spec = VectorSearchChainSpec(
        chain_id=3,
        chain_type="vector_search_chain_spec",
        input_keys=["input1", "input2"],
        output_key="output",
        embedding_engine="openai",
        embedding_params={"openai_api_key": "test"},
        database="pinecone",
        database_params={"index": "test", "text_key": "text"},
        num_results=10,
        query="How does {input}?"
    )
    serialized = vector_search_chain_spec.json()
    deserialized = VectorSearchChainSpec.parse_raw(serialized)
    assert deserialized == vector_search_chain_spec


class ChainDict:
    def __init__(self):
        self.chains = {}
    
    def add_chain(self, chain):
        self.chains[chain.chain_id] = chain


def test_chain_spec_copy_replace():
    ctx = SpecFactoryContext()
    chain = sequential_factory(ctx, chains=[
        llm_factory(ctx),
        case_factory(ctx, cases={
            "case1": llm_factory(ctx),
            "case2": sequential_factory(ctx, chains=[llm_factory(ctx), llm_factory(ctx)]),
        }, default_case=llm_factory(ctx)),
    ])

    original_specs = ChainDict()
    chain.traverse(original_specs.add_chain)
    copied_specs = {}
    
    copied_chain = chain.copy_replace(lambda spec: spec)
    copied_specs = ChainDict()
    copied_chain.traverse(copied_specs.add_chain)

    assert len(original_specs.chains) == len(copied_specs.chains)

    for chain_id, spec in original_specs.chains.items():
        copied_spec = copied_specs.chains.get(chain_id)
        assert copied_spec is not None
        assert copied_spec == spec
        assert copied_spec is not spec

    replacement = copied_specs.chains[3]
    replacement.prompt = "replaced"
    replace_chain = chain.copy_replace(lambda spec: spec if spec.chain_id != 3 else replacement)
    replace_specs = ChainDict()
    replace_chain.traverse(replace_specs.add_chain)

    assert len(original_specs.chains) == len(replace_specs.chains)
    assert replace_specs.chains[3] == replacement
    assert replace_specs.chains[3].prompt == "replaced"


def test_base_spec_find_by_chain_id():
    ctx = SpecFactoryContext()
    deep_llm = llm_factory(ctx)
    chain = sequential_factory(ctx, chains=[
        llm_factory(ctx),
        case_factory(ctx, cases={
            "case1": llm_factory(ctx),
            "case2": sequential_factory(ctx, chains=[llm_factory(ctx), deep_llm]),
        }, default_case=llm_factory(ctx)),
    ])

    assert chain.find_by_chain_id(deep_llm.chain_id) == deep_llm
langchain.llms.fake.FakeListLLM
[((471, 620), 'model.chain_spec.LLMChainSpec', 'LLMChainSpec', ([], {'chain_id': '(1)', 'input_keys': "['input1', 'input2']", 'output_key': '"""output1"""', 'prompt': '"""prompt"""', 'llm_key': '"""llm_key"""', 'chain_type': '"""llm_chain_spec"""'}), "(chain_id=1, input_keys=['input1', 'input2'], output_key=\n    'output1', prompt='prompt', llm_key='llm_key', chain_type='llm_chain_spec')\n", (483, 620), False, 'from model.chain_spec import ChainSpec, LLMChainSpec, SequentialChainSpec, CaseChainSpec, APIChainSpec, ReformatChainSpec, TransformChainSpec, VectorSearchChainSpec\n'), ((904, 948), 'model.chain_revision.ChainRevision', 'ChainRevision', ([], {'chain': 'llm_chain_spec', 'llms': '{}'}), '(chain=llm_chain_spec, llms={})\n', (917, 948), False, 'from model.chain_revision import ChainRevision\n'), ((1240, 1398), 'model.chain_spec.LLMChainSpec', 'LLMChainSpec', ([], {'chain_id': '(1)', 'input_keys': "['input1', 'input2']", 'output_key': '"""output1"""', 'prompt': 'prompt_template', 'llm_key': '"""test"""', 'chain_type': '"""llm_chain_spec"""'}), "(chain_id=1, input_keys=['input1', 'input2'], output_key=\n    'output1', prompt=prompt_template, llm_key='test', chain_type=\n    'llm_chain_spec')\n", (1252, 1398), False, 'from model.chain_spec import ChainSpec, LLMChainSpec, SequentialChainSpec, CaseChainSpec, APIChainSpec, ReformatChainSpec, TransformChainSpec, VectorSearchChainSpec\n'), ((1517, 1544), 'model.lang_chain_context.LangChainContext', 'LangChainContext', ([], {'llms': 'llms'}), '(llms=llms)\n', (1533, 1544), False, 'from model.lang_chain_context import LangChainContext\n'), ((2053, 2211), 'model.chain_spec.LLMChainSpec', 'LLMChainSpec', ([], {'chain_id': '(1)', 'input_keys': "['input1', 'input2']", 'output_key': '"""output1"""', 'prompt': 'prompt_template', 'llm_key': '"""test"""', 'chain_type': '"""llm_chain_spec"""'}), "(chain_id=1, input_keys=['input1', 'input2'], output_key=\n    'output1', prompt=prompt_template, llm_key='test', chain_type=\n    'llm_chain_spec')\n", (2065, 2211), False, 'from model.chain_spec import ChainSpec, LLMChainSpec, SequentialChainSpec, CaseChainSpec, APIChainSpec, ReformatChainSpec, TransformChainSpec, VectorSearchChainSpec\n'), ((2330, 2373), 'model.lang_chain_context.LangChainContext', 'LangChainContext', ([], {'llms': 'llms', 'recording': '(True)'}), '(llms=llms, recording=True)\n', (2346, 2373), False, 'from model.lang_chain_context import LangChainContext\n'), ((3311, 3352), 'model.chain_spec.SequentialChainSpec.parse_raw', 'SequentialChainSpec.parse_raw', (['serialized'], {}), '(serialized)\n', (3340, 3352), False, 'from model.chain_spec import ChainSpec, LLMChainSpec, SequentialChainSpec, CaseChainSpec, APIChainSpec, ReformatChainSpec, TransformChainSpec, VectorSearchChainSpec\n'), ((3495, 3671), 'model.chain_spec.LLMChainSpec', 'LLMChainSpec', ([], {'chain_id': '(1)', 'input_keys': "['input1', 'input2']", 'output_key': '"""llm1_output"""', 'prompt': '"""the prompt {input1} {input2}"""', 'llm_key': '"""test"""', 'chain_type': '"""llm_chain_spec"""'}), "(chain_id=1, input_keys=['input1', 'input2'], output_key=\n    'llm1_output', prompt='the prompt {input1} {input2}', llm_key='test',\n    chain_type='llm_chain_spec')\n", (3507, 3671), False, 'from model.chain_spec import ChainSpec, LLMChainSpec, SequentialChainSpec, CaseChainSpec, APIChainSpec, ReformatChainSpec, TransformChainSpec, VectorSearchChainSpec\n'), ((3796, 3983), 'model.chain_spec.LLMChainSpec', 'LLMChainSpec', ([], {'chain_id': '(2)', 'input_keys': "['llm1_output', 'input3']", 'output_key': '"""llm2_output"""', 'prompt': '"""the prompt {llm1_output} {input3}"""', 'llm_key': '"""test"""', 'chain_type': '"""llm_chain_spec"""'}), "(chain_id=2, input_keys=['llm1_output', 'input3'], output_key=\n    'llm2_output', prompt='the prompt {llm1_output} {input3}', llm_key=\n    'test', chain_type='llm_chain_spec')\n", (3808, 3983), False, 'from model.chain_spec import ChainSpec, LLMChainSpec, SequentialChainSpec, CaseChainSpec, APIChainSpec, ReformatChainSpec, TransformChainSpec, VectorSearchChainSpec\n'), ((4114, 4305), 'model.chain_spec.SequentialChainSpec', 'SequentialChainSpec', ([], {'chain_id': '(3)', 'input_keys': "['input1', 'input2', 'input3']", 'output_keys': "['llm2_output']", 'chain_type': '"""sequential_chain_spec"""', 'chains': '[llm_chain_spec1, llm_chain_spec2]'}), "(chain_id=3, input_keys=['input1', 'input2', 'input3'],\n    output_keys=['llm2_output'], chain_type='sequential_chain_spec', chains\n    =[llm_chain_spec1, llm_chain_spec2])\n", (4133, 4305), False, 'from model.chain_spec import ChainSpec, LLMChainSpec, SequentialChainSpec, CaseChainSpec, APIChainSpec, ReformatChainSpec, TransformChainSpec, VectorSearchChainSpec\n'), ((4441, 4468), 'model.lang_chain_context.LangChainContext', 'LangChainContext', ([], {'llms': 'llms'}), '(llms=llms)\n', (4457, 4468), False, 'from model.lang_chain_context import LangChainContext\n'), ((6447, 6482), 'model.chain_spec.CaseChainSpec.parse_raw', 'CaseChainSpec.parse_raw', (['serialized'], {}), '(serialized)\n', (6470, 6482), False, 'from model.chain_spec import ChainSpec, LLMChainSpec, SequentialChainSpec, CaseChainSpec, APIChainSpec, ReformatChainSpec, TransformChainSpec, VectorSearchChainSpec\n'), ((6728, 6900), 'model.chain_spec.LLMChainSpec', 'LLMChainSpec', ([], {'chain_id': '(1)', 'input_keys': "['input1', 'input2']", 'output_key': '"""output"""', 'prompt': '"""ll1 prompt {input1} {input2}"""', 'llm_key': '"""test1"""', 'chain_type': '"""llm_chain_spec"""'}), "(chain_id=1, input_keys=['input1', 'input2'], output_key=\n    'output', prompt='ll1 prompt {input1} {input2}', llm_key='test1',\n    chain_type='llm_chain_spec')\n", (6740, 6900), False, 'from model.chain_spec import ChainSpec, LLMChainSpec, SequentialChainSpec, CaseChainSpec, APIChainSpec, ReformatChainSpec, TransformChainSpec, VectorSearchChainSpec\n'), ((7025, 7198), 'model.chain_spec.LLMChainSpec', 'LLMChainSpec', ([], {'chain_id': '(2)', 'input_keys': "['input3', 'input4']", 'output_key': '"""output"""', 'prompt': '"""llm2 prompt {input3} {input4}"""', 'llm_key': '"""test2"""', 'chain_type': '"""llm_chain_spec"""'}), "(chain_id=2, input_keys=['input3', 'input4'], output_key=\n    'output', prompt='llm2 prompt {input3} {input4}', llm_key='test2',\n    chain_type='llm_chain_spec')\n", (7037, 7198), False, 'from model.chain_spec import ChainSpec, LLMChainSpec, SequentialChainSpec, CaseChainSpec, APIChainSpec, ReformatChainSpec, TransformChainSpec, VectorSearchChainSpec\n'), ((7323, 7499), 'model.chain_spec.LLMChainSpec', 'LLMChainSpec', ([], {'chain_id': '(4)', 'input_keys': "['input1', 'input2']", 'output_key': '"""output"""', 'prompt': '"""default prompt {input1} {input2}"""', 'llm_key': '"""test1"""', 'chain_type': '"""llm_chain_spec"""'}), "(chain_id=4, input_keys=['input1', 'input2'], output_key=\n    'output', prompt='default prompt {input1} {input2}', llm_key='test1',\n    chain_type='llm_chain_spec')\n", (7335, 7499), False, 'from model.chain_spec import ChainSpec, LLMChainSpec, SequentialChainSpec, CaseChainSpec, APIChainSpec, ReformatChainSpec, TransformChainSpec, VectorSearchChainSpec\n'), ((7624, 7829), 'model.chain_spec.CaseChainSpec', 'CaseChainSpec', ([], {'chain_id': '(3)', 'chain_type': '"""case_chain_spec"""', 'categorization_key': '"""categorization_input"""', 'cases': "{'response1': llm_chain_spec1, 'response2': llm_chain_spec2}", 'default_case': 'llm_chain_spec3'}), "(chain_id=3, chain_type='case_chain_spec', categorization_key=\n    'categorization_input', cases={'response1': llm_chain_spec1,\n    'response2': llm_chain_spec2}, default_case=llm_chain_spec3)\n", (7637, 7829), False, 'from model.chain_spec import ChainSpec, LLMChainSpec, SequentialChainSpec, CaseChainSpec, APIChainSpec, ReformatChainSpec, TransformChainSpec, VectorSearchChainSpec\n'), ((7922, 7949), 'model.lang_chain_context.LangChainContext', 'LangChainContext', ([], {'llms': 'llms'}), '(llms=llms)\n', (7938, 7949), False, 'from model.lang_chain_context import LangChainContext\n'), ((9327, 9543), 'model.chain_spec.APIChainSpec', 'APIChainSpec', ([], {'chain_id': '(3)', 'chain_type': '"""api_chain_spec"""', 'input_keys': "['input1', 'input2']", 'output_key': '"""output"""', 'url': '"""http://test.com"""', 'method': '"""POST"""', 'body': '"""body {input1} {input2}"""', 'headers': "{'header1': 'header1'}"}), "(chain_id=3, chain_type='api_chain_spec', input_keys=['input1',\n    'input2'], output_key='output', url='http://test.com', method='POST',\n    body='body {input1} {input2}', headers={'header1': 'header1'})\n", (9339, 9543), False, 'from model.chain_spec import ChainSpec, LLMChainSpec, SequentialChainSpec, CaseChainSpec, APIChainSpec, ReformatChainSpec, TransformChainSpec, VectorSearchChainSpec\n'), ((9623, 9667), 'model.chain_revision.ChainRevision', 'ChainRevision', ([], {'chain': 'api_chain_spec', 'llms': '{}'}), '(chain=api_chain_spec, llms={})\n', (9636, 9667), False, 'from model.chain_revision import ChainRevision\n'), ((9906, 10122), 'model.chain_spec.APIChainSpec', 'APIChainSpec', ([], {'chain_id': '(3)', 'chain_type': '"""api_chain_spec"""', 'input_keys': "['input1', 'input2']", 'output_key': '"""output"""', 'url': '"""http://test.com"""', 'method': '"""POST"""', 'body': '"""body {input1} {input2}"""', 'headers': "{'header1': 'header1'}"}), "(chain_id=3, chain_type='api_chain_spec', input_keys=['input1',\n    'input2'], output_key='output', url='http://test.com', method='POST',\n    body='body {input1} {input2}', headers={'header1': 'header1'})\n", (9918, 10122), False, 'from model.chain_spec import ChainSpec, LLMChainSpec, SequentialChainSpec, CaseChainSpec, APIChainSpec, ReformatChainSpec, TransformChainSpec, VectorSearchChainSpec\n'), ((10196, 10221), 'model.lang_chain_context.LangChainContext', 'LangChainContext', ([], {'llms': '{}'}), '(llms={})\n', (10212, 10221), False, 'from model.lang_chain_context import LangChainContext\n'), ((10643, 10811), 'model.chain_spec.ReformatChainSpec', 'ReformatChainSpec', ([], {'chain_id': '(3)', 'chain_type': '"""reformat_chain_spec"""', 'input_keys': "['input1', 'input2']", 'formatters': "{'output1': 'formatter1', 'output2': 'formatter2'}"}), "(chain_id=3, chain_type='reformat_chain_spec', input_keys=\n    ['input1', 'input2'], formatters={'output1': 'formatter1', 'output2':\n    'formatter2'})\n", (10660, 10811), False, 'from model.chain_spec import ChainSpec, LLMChainSpec, SequentialChainSpec, CaseChainSpec, APIChainSpec, ReformatChainSpec, TransformChainSpec, VectorSearchChainSpec\n'), ((10905, 10944), 'model.chain_spec.ReformatChainSpec.parse_raw', 'ReformatChainSpec.parse_raw', (['serialized'], {}), '(serialized)\n', (10932, 10944), False, 'from model.chain_spec import ChainSpec, LLMChainSpec, SequentialChainSpec, CaseChainSpec, APIChainSpec, ReformatChainSpec, TransformChainSpec, VectorSearchChainSpec\n'), ((11086, 11254), 'model.chain_spec.ReformatChainSpec', 'ReformatChainSpec', ([], {'chain_id': '(3)', 'chain_type': '"""reformat_chain_spec"""', 'input_keys': "['input1', 'input2']", 'formatters': "{'output1': 'formatter1', 'output2': 'formatter2'}"}), "(chain_id=3, chain_type='reformat_chain_spec', input_keys=\n    ['input1', 'input2'], formatters={'output1': 'formatter1', 'output2':\n    'formatter2'})\n", (11103, 11254), False, 'from model.chain_spec import ChainSpec, LLMChainSpec, SequentialChainSpec, CaseChainSpec, APIChainSpec, ReformatChainSpec, TransformChainSpec, VectorSearchChainSpec\n'), ((11295, 11320), 'model.lang_chain_context.LangChainContext', 'LangChainContext', ([], {'llms': '{}'}), '(llms={})\n', (11311, 11320), False, 'from model.lang_chain_context import LangChainContext\n'), ((11694, 11947), 'model.chain_spec.TransformChainSpec', 'TransformChainSpec', ([], {'chain_id': '(4)', 'chain_type': '"""transform_chain_spec"""', 'input_keys': "['x', 'y']", 'output_keys': "['z', 'w']", 'transform_func': '"""\nz = int(inputs[\'x\']) + int(inputs[\'y\'])\nw = int(inputs[\'x\']) - int(inputs[\'y\'])\nreturn {\'z\': z, \'w\': w}"""'}), '(chain_id=4, chain_type=\'transform_chain_spec\',\n    input_keys=[\'x\', \'y\'], output_keys=[\'z\', \'w\'], transform_func=\n    """\nz = int(inputs[\'x\']) + int(inputs[\'y\'])\nw = int(inputs[\'x\']) - int(inputs[\'y\'])\nreturn {\'z\': z, \'w\': w}"""\n    )\n', (11712, 11947), False, 'from model.chain_spec import ChainSpec, LLMChainSpec, SequentialChainSpec, CaseChainSpec, APIChainSpec, ReformatChainSpec, TransformChainSpec, VectorSearchChainSpec\n'), ((11990, 12015), 'model.lang_chain_context.LangChainContext', 'LangChainContext', ([], {'llms': '{}'}), '(llms={})\n', (12006, 12015), False, 'from model.lang_chain_context import LangChainContext\n'), ((12343, 12678), 'model.chain_spec.VectorSearchChainSpec', 'VectorSearchChainSpec', ([], {'chain_id': '(3)', 'chain_type': '"""vector_search_chain_spec"""', 'input_keys': "['input1', 'input2']", 'output_key': '"""output"""', 'embedding_engine': '"""openai"""', 'embedding_params': "{'openai_api_key': 'test'}", 'database': '"""pinecone"""', 'database_params': "{'index': 'test', 'text_key': 'text'}", 'num_results': '(10)', 'query': '"""How does {input}?"""'}), "(chain_id=3, chain_type='vector_search_chain_spec',\n    input_keys=['input1', 'input2'], output_key='output', embedding_engine=\n    'openai', embedding_params={'openai_api_key': 'test'}, database=\n    'pinecone', database_params={'index': 'test', 'text_key': 'text'},\n    num_results=10, query='How does {input}?')\n", (12364, 12678), False, 'from model.chain_spec import ChainSpec, LLMChainSpec, SequentialChainSpec, CaseChainSpec, APIChainSpec, ReformatChainSpec, TransformChainSpec, VectorSearchChainSpec\n'), ((12815, 12858), 'model.chain_spec.VectorSearchChainSpec.parse_raw', 'VectorSearchChainSpec.parse_raw', (['serialized'], {}), '(serialized)\n', (12846, 12858), False, 'from model.chain_spec import ChainSpec, LLMChainSpec, SequentialChainSpec, CaseChainSpec, APIChainSpec, ReformatChainSpec, TransformChainSpec, VectorSearchChainSpec\n'), ((13108, 13128), 'model.tests.factory.SpecFactoryContext', 'SpecFactoryContext', ([], {}), '()\n', (13126, 13128), False, 'from model.tests.factory import SpecFactoryContext, llm_factory, sequential_factory, case_factory\n'), ((14436, 14456), 'model.tests.factory.SpecFactoryContext', 'SpecFactoryContext', ([], {}), '()\n', (14454, 14456), False, 'from model.tests.factory import SpecFactoryContext, llm_factory, sequential_factory, case_factory\n'), ((14472, 14488), 'model.tests.factory.llm_factory', 'llm_factory', (['ctx'], {}), '(ctx)\n', (14483, 14488), False, 'from model.tests.factory import SpecFactoryContext, llm_factory, sequential_factory, case_factory\n'), ((1010, 1054), 'model.chain_revision.ChainRevision.parse_raw', 'ChainRevision.parse_raw', (['serialized_revision'], {}), '(serialized_revision)\n', (1033, 1054), False, 'from model.chain_revision import ChainRevision\n'), ((1469, 1505), 'langchain.llms.fake.FakeListLLM', 'FakeListLLM', ([], {'responses': "['response1']"}), "(responses=['response1'])\n", (1480, 1505), False, 'from langchain.llms.fake import FakeListLLM\n'), ((2282, 2318), 'langchain.llms.fake.FakeListLLM', 'FakeListLLM', ([], {'responses': "['response1']"}), "(responses=['response1'])\n", (2293, 2318), False, 'from langchain.llms.fake import FakeListLLM\n'), ((4370, 4429), 'langchain.llms.fake.FakeListLLM', 'FakeListLLM', ([], {'responses': "['fake_response1', 'fake_response2']"}), "(responses=['fake_response1', 'fake_response2'])\n", (4381, 4429), False, 'from langchain.llms.fake import FakeListLLM\n'), ((6611, 6652), 'langchain.llms.fake.FakeListLLM', 'FakeListLLM', ([], {'responses': "['fake_response1']"}), "(responses=['fake_response1'])\n", (6622, 6652), False, 'from langchain.llms.fake import FakeListLLM\n'), ((6663, 6704), 'langchain.llms.fake.FakeListLLM', 'FakeListLLM', ([], {'responses': "['fake_response2']"}), "(responses=['fake_response2'])\n", (6674, 6704), False, 'from langchain.llms.fake import FakeListLLM\n'), ((9729, 9773), 'model.chain_revision.ChainRevision.parse_raw', 'ChainRevision.parse_raw', (['serialized_revision'], {}), '(serialized_revision)\n', (9752, 9773), False, 'from model.chain_revision import ChainRevision\n'), ((6158, 6303), 'model.chain_spec.LLMChainSpec', 'LLMChainSpec', ([], {'chain_id': '(4)', 'input_keys': "['input1', 'input2']", 'output_key': '"""output"""', 'prompt': '"""prompt"""', 'llm_key': '"""test"""', 'chain_type': '"""llm_chain_spec"""'}), "(chain_id=4, input_keys=['input1', 'input2'], output_key=\n    'output', prompt='prompt', llm_key='test', chain_type='llm_chain_spec')\n", (6170, 6303), False, 'from model.chain_spec import ChainSpec, LLMChainSpec, SequentialChainSpec, CaseChainSpec, APIChainSpec, ReformatChainSpec, TransformChainSpec, VectorSearchChainSpec\n'), ((2974, 3123), 'model.chain_spec.LLMChainSpec', 'LLMChainSpec', ([], {'chain_id': '(1)', 'input_keys': "['input1', 'input2']", 'output_key': '"""output1"""', 'prompt': '"""prompt"""', 'llm_key': '"""llm_key"""', 'chain_type': '"""llm_chain_spec"""'}), "(chain_id=1, input_keys=['input1', 'input2'], output_key=\n    'output1', prompt='prompt', llm_key='llm_key', chain_type='llm_chain_spec')\n", (2986, 3123), False, 'from model.chain_spec import ChainSpec, LLMChainSpec, SequentialChainSpec, CaseChainSpec, APIChainSpec, ReformatChainSpec, TransformChainSpec, VectorSearchChainSpec\n'), ((5598, 5743), 'model.chain_spec.LLMChainSpec', 'LLMChainSpec', ([], {'chain_id': '(1)', 'input_keys': "['input1', 'input2']", 'output_key': '"""output"""', 'prompt': '"""prompt"""', 'llm_key': '"""test"""', 'chain_type': '"""llm_chain_spec"""'}), "(chain_id=1, input_keys=['input1', 'input2'], output_key=\n    'output', prompt='prompt', llm_key='test', chain_type='llm_chain_spec')\n", (5610, 5743), False, 'from model.chain_spec import ChainSpec, LLMChainSpec, SequentialChainSpec, CaseChainSpec, APIChainSpec, ReformatChainSpec, TransformChainSpec, VectorSearchChainSpec\n'), ((5875, 6020), 'model.chain_spec.LLMChainSpec', 'LLMChainSpec', ([], {'chain_id': '(2)', 'input_keys': "['input3', 'input4']", 'output_key': '"""output"""', 'prompt': '"""prompt"""', 'llm_key': '"""test"""', 'chain_type': '"""llm_chain_spec"""'}), "(chain_id=2, input_keys=['input3', 'input4'], output_key=\n    'output', prompt='prompt', llm_key='test', chain_type='llm_chain_spec')\n", (5887, 6020), False, 'from model.chain_spec import ChainSpec, LLMChainSpec, SequentialChainSpec, CaseChainSpec, APIChainSpec, ReformatChainSpec, TransformChainSpec, VectorSearchChainSpec\n'), ((13182, 13198), 'model.tests.factory.llm_factory', 'llm_factory', (['ctx'], {}), '(ctx)\n', (13193, 13198), False, 'from model.tests.factory import SpecFactoryContext, llm_factory, sequential_factory, case_factory\n'), ((14542, 14558), 'model.tests.factory.llm_factory', 'llm_factory', (['ctx'], {}), '(ctx)\n', (14553, 14558), False, 'from model.tests.factory import SpecFactoryContext, llm_factory, sequential_factory, case_factory\n'), ((13388, 13404), 'model.tests.factory.llm_factory', 'llm_factory', (['ctx'], {}), '(ctx)\n', (13399, 13404), False, 'from model.tests.factory import SpecFactoryContext, llm_factory, sequential_factory, case_factory\n'), ((14740, 14756), 'model.tests.factory.llm_factory', 'llm_factory', (['ctx'], {}), '(ctx)\n', (14751, 14756), False, 'from model.tests.factory import SpecFactoryContext, llm_factory, sequential_factory, case_factory\n'), ((13255, 13271), 'model.tests.factory.llm_factory', 'llm_factory', (['ctx'], {}), '(ctx)\n', (13266, 13271), False, 'from model.tests.factory import SpecFactoryContext, llm_factory, sequential_factory, case_factory\n'), ((14615, 14631), 'model.tests.factory.llm_factory', 'llm_factory', (['ctx'], {}), '(ctx)\n', (14626, 14631), False, 'from model.tests.factory import SpecFactoryContext, llm_factory, sequential_factory, case_factory\n'), ((13326, 13342), 'model.tests.factory.llm_factory', 'llm_factory', (['ctx'], {}), '(ctx)\n', (13337, 13342), False, 'from model.tests.factory import SpecFactoryContext, llm_factory, sequential_factory, case_factory\n'), ((13344, 13360), 'model.tests.factory.llm_factory', 'llm_factory', (['ctx'], {}), '(ctx)\n', (13355, 13360), False, 'from model.tests.factory import SpecFactoryContext, llm_factory, sequential_factory, case_factory\n'), ((14686, 14702), 'model.tests.factory.llm_factory', 'llm_factory', (['ctx'], {}), '(ctx)\n', (14697, 14702), False, 'from model.tests.factory import SpecFactoryContext, llm_factory, sequential_factory, case_factory\n')]