small code fixes

final/README.md

@@ -19,8 +19,10 @@ python3 app.py
 ```
 
 ## Example tests
-*Show a summary of network traffic on enp7s0*
+*Give a summary of network traffic on my ethernet adapter*
 
-*Show a visual representation of the network traffic on enp7s0*
+*Based on ethernet traffic, what applications might I be running*
 
-*Show a list of IP addresses that communicate on enp7s0 and identify the countries of those addresses*
+*What countries or regions am I sending traffic to via my ethernet adapter*
+
+*Show a visual representation of the network traffic on my ethernet adapter*

final/app.py

@@ -2,7 +2,7 @@ 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, get_wireless_interface, ip_loc, pcap_summary
+from tools import dalle, tcp_dump, get_adapter_interface, ip_loc, pcap_summary
 from langchain import hub
 from langchain_community.tools import ShellTool
 from langsmith import Client
@@ -22,12 +22,13 @@ client = Client()
 shell_tool = ShellTool()
 llm = ChatOpenAI(model_name="gpt-4o", temperature=0)
 tools = []
-tools.extend([tcp_dump, ip_loc, dalle, pcap_summary, get_wireless_interface])
+tools.extend([tcp_dump, ip_loc, dalle, pcap_summary, get_adapter_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.
+                Never make assumptions about which network adapter to use, always check the adapter name with the correct tool.
                 """
 )
 agent = create_react_agent(llm, tools, prompt)

final/tools.py

@@ -1,6 +1,7 @@
 from datetime import datetime
 from io import StringIO
 import sys
+from time import sleep
 from langchain_core.pydantic_v1 import BaseModel, Field
 from langchain.tools import tool
 from langchain_community.tools import ShellTool
@@ -50,7 +51,7 @@ class Image(BaseModel):
 
 
 @tool(
-    "Create an image of a network. This tool must take as input a summary of packet information, and not a file name",
+    "Create an image of a network. This tool must take as input a summary of packet information, and not a file name. If necessary, condense input to 900 characters or less",
     args_schema=Image,
     return_direct=False,
 )
@@ -96,9 +97,9 @@ class Iwconfig(BaseModel):
     )
 
 
-@tool("Get interface information", args_schema=Iwconfig, return_direct=False)
+@tool("Get interface adapter name", args_schema=Iwconfig, return_direct=False)
-def get_wireless_interface(params: str) -> str:
+def get_adapter_interface(params: str) -> str:
-    """Return wireless interface information via iwconfig"""
+    """Network interface adapters via iwconfig"""
     params = params.replace("`", "").replace("\n", "")  # fix buggy input from LLM
     print("params ", params)
     res = shell_tool.run({"commands": [f"iwconfig {params}"]})
@@ -109,6 +110,7 @@ def get_wireless_interface(params: str) -> str:
 def ip_loc(address):
     """Get information from an ip address, including geolocation. Takes as a paramater an ip address. Do not use this tool with IP adresses in a reserve range or on LAN"""
     url = f"http://ipwho.is/{address}"
+    sleep(2)
     response = requests.get(url)
     if response.status_code == 200:
         return response.json()