From 5b1b7d6ab85f8a2248fd8889ccc198d809c27d32 Mon Sep 17 00:00:00 2001
From: Aaron <29749331+aarnphm@users.noreply.github.com>
Date: Tue, 30 May 2023 14:37:00 -0700
Subject: [PATCH] chore(llm): expose langchain API to the runner

Signed-off-by: Aaron <29749331+aarnphm@users.noreply.github.com>
---
 src/openllm/_llm.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/src/openllm/_llm.py b/src/openllm/_llm.py
index 11b60321..dec9f9f8 100644
--- a/src/openllm/_llm.py
+++ b/src/openllm/_llm.py
@@ -44,6 +44,8 @@ else:
 
 logger = logging.getLogger(__name__)
 
+_object_setattr = object.__setattr__
+
 # NOTE: `1-2` -> text-generation and text2text-generation
 FRAMEWORK_TO_AUTOCLASS_MAPPING = {
     "pt": ("AutoModelForCausalLM", "AutoModelForSeq2SeqLM"),
@@ -742,7 +744,7 @@ class LLM(LLMInterface, metaclass=LLMMetaclass):
         else:
             _supported_resources = ("nvidia.com/gpu", "cpu")
 
-        return bentoml.Runner(
+        runner = bentoml.Runner(
             type(
                 inflection.camelize(self.config.__openllm_model_name__) + "Runnable",
                 (_Runnable,),
@@ -757,6 +759,12 @@ class LLM(LLMInterface, metaclass=LLMMetaclass):
             scheduling_strategy=scheduling_strategy,
         )
 
+        # NOTE: returning the two langchain API's to the runner
+        _object_setattr(runner, "llm_type", self.llm_type)
+        _object_setattr(runner, "identifying_params", self.identifying_params)
+
+        return runner
+
 
 def Runner(start_name: str, **attrs: t.Any) -> bentoml.Runner:
     """Create a Runner for given LLM. For a list of currently supported LLM, check out 'openllm models'