1  # This is an example for logging a Langchain model from code using the
 2  # mlflow.langchain.log_model API. When a path to a valid Python script is submitted to the
 3  # lc_model argument, the model code itself is serialized instead of the model object.
 4  # Within the targeted script, the model implementation must be defined and set by
 5  # using the mlflow.models.set_model API.
 6  
 7  import mlflow
 8  
 9  input_example = {
10      "messages": [
11          {
12              "role": "user",
13              "content": "What is Retrieval-augmented Generation?",
14          }
15      ]
16  }
17  
18  # Specify the path to the chain notebook
19  chain_path = "chain_as_code.py"
20  
21  print(f"Chain path: {chain_path}")
22  
23  print("Logging model as code using Langchain log model API")
24  with mlflow.start_run():
25      logged_chain_info = mlflow.langchain.log_model(
26          lc_model=chain_path,
27          name="chain",
28          input_example=input_example,
29      )
30  
31  print("Loading model using Langchain load model API")
32  model = mlflow.langchain.load_model(logged_chain_info.model_uri)
33  output = model.invoke(input_example)
34  print(f"Output: {output}")
35  
36  print("Loading model using Pyfunc load model API")
37  pyfunc_model = mlflow.pyfunc.load_model(logged_chain_info.model_uri)
38  output = pyfunc_model.predict([input_example])
39  print(f"Output: {output}")