Wrapped LLM into LangChain Chat model

app.py

@@ -2,7 +2,8 @@
 import uuid
 import streamlit as st
 import datasets
-from langchain_huggingface import HuggingFaceEndpointEmbeddings, HuggingFaceEndpoint
+from langchain_huggingface import HuggingFaceEndpointEmbeddings, ChatHuggingFace
+from langchain_community.llms import HuggingFaceEndpoint
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.runnables import RunnablePassthrough
 from langchain.agents import create_react_agent, AgentExecutor
@@ -49,8 +50,8 @@ def setup_chroma_embedding_function():
     return chroma_embedding_function
 
 # Set up HuggingFaceEndpoint model
-def setup_huggingface_endpoint():
-    model = HuggingFaceEndpoint(
+def setup_huggingface_endpoint(model_id):
+    llm = HuggingFaceEndpoint(
         endpoint_url="http://{host}:{port}".format(
             host=os.getenv("TGI_HOST", "localhost"), port=os.getenv("TGI_PORT", "8080")
         ),
@@ -61,6 +62,10 @@ def setup_huggingface_endpoint():
             "{your_token}".format(your_token=os.getenv("STOP_TOKEN", "<|end_of_text|>")),
         ],
     )
+
+    model = ChatHuggingFace(llm=llm,
+                            model_id=model_id)
+
     return model
 
 def setup_portkey_integrated_model():
@@ -220,6 +225,7 @@ def setup_tools(_model, _client, _chroma_embedding_function, _embedder):
     #    chroma_embedding_function=_chroma_embedding_function,
     #    embedder=_embedder,
     #)
+
     if os.getenv("USE_RERANKER", "False") == "True":
         retriever = create_reranker_retriever(
             name="slack_conversations_retriever",
@@ -245,7 +251,7 @@ def setup_tools(_model, _client, _chroma_embedding_function, _embedder):
 @st.cache_resource
 def setup_agent(_model, _prompt, _client, _chroma_embedding_function, _embedder):
     tools = setup_tools(_model, _client, _chroma_embedding_function, _embedder)
-    agent = create_react_agent(llm=_model, prompt=_prompt, tools=tools,)
+    agent = create_react_agent(llm=_model, prompt=_prompt, tools=tools, )
     agent_executor = AgentExecutor(
         agent=agent, verbose=True, tools=tools, handle_parsing_errors=True
     )
@@ -258,7 +264,7 @@ def main():
     if os.getenv("ENABLE_PORTKEY", "False") == "True":
         model = setup_portkey_integrated_model()
     else:
-        model = setup_huggingface_endpoint()
+        model = setup_huggingface_endpoint(model_id="qwen/Qwen2-7B-Instruct")
     embedder = setup_huggingface_embeddings()
 
     agent_executor = setup_agent(