60 lines
1.7 KiB
Python
60 lines
1.7 KiB
Python
import os
|
|
from langchain_openai import ChatOpenAI
|
|
from langchain.agents import AgentExecutor, create_react_agent
|
|
from langchain.tools import tool
|
|
from tools import (
|
|
dalle,
|
|
tcp_dump,
|
|
tshark,
|
|
get_wireless_interface,
|
|
ip_loc
|
|
)
|
|
from langchain import hub
|
|
from langchain_community.tools import ShellTool
|
|
from langsmith import Client
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
|
|
"""
|
|
Main rag tool-chain application. Simply load the tools and prompt the user.
|
|
"""
|
|
|
|
os.environ["LANGCHAIN_TRACING_V2"] = "true"
|
|
os.environ["LANGCHAIN_PROJECT"] = f"LangSmith Introduction"
|
|
os.environ["LANGCHAIN_ENDPOINT"] = "https://api.smith.langchain.com"
|
|
|
|
client = Client()
|
|
shell_tool = ShellTool()
|
|
llm = ChatOpenAI(model_name="gpt-4o", temperature=0)
|
|
tools = []
|
|
tools.extend([tcp_dump,ip_loc, dalle,get_wireless_interface])
|
|
base_prompt = hub.pull("langchain-ai/react-agent-template")
|
|
prompt = base_prompt.partial(
|
|
instructions="""
|
|
You are a packet analysis assistant. Use any combination of the tools provided to best serve the users request.
|
|
If the request cannot be served with the tools provided, state why and offer advice on how the user could solve the problem.
|
|
"""
|
|
)
|
|
agent = create_react_agent(llm, tools, prompt)
|
|
agent_executor = AgentExecutor(
|
|
agent=agent, tools=tools, verbose=True, handle_parsing_errors=True
|
|
)
|
|
|
|
|
|
print(
|
|
"I am a packet analysis an assistant. I can perform various tasks related to packet capture files."
|
|
)
|
|
print(f"I am configured with the following tools")
|
|
|
|
for tool in tools:
|
|
print(f" Tool: {tool.name} = {tool.description}")
|
|
|
|
while True:
|
|
line = input("llm>> ")
|
|
if line:
|
|
result = agent_executor.invoke({"input": line})
|
|
print(result["output"])
|
|
else:
|
|
break
|