Spaces:
Sleeping
Sleeping
File size: 1,133 Bytes
f0235ce |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 |
import os
from langchain import hub
from langchain.agents import load_tools, initialize_agent, AgentType
from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
from langchain_huggingface import HuggingFaceEndpoint
from langchain_huggingface import ChatHuggingFace
class LLMResource:
llm = HuggingFaceEndpoint(
repo_id="mistralai/Mistral-7B-Instruct-v0.2",
task="text-generation",
huggingfacehub_api_token=os.environ["API_KEY"],
streaming=True,
temperature=0.5
)
chat_model = ChatHuggingFace(
llm = llm,
)
tools = load_tools(["wikipedia"], llm=llm)
agent = initialize_agent(
tools,
llm,
agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
handle_parsing_errors=True,
)
@staticmethod
def invoke(message:str):
response = LLMResource.agent.invoke(message)
print(response)
return response["output"][:-4]
@staticmethod
def test():
print(LLMResource.agent.invoke("When was the movie \'Oppenheimer\' released?"))
if __name__ == "__main__":
LLMResource.test() |