wrigleyDan commited on
Commit
5c5fbc2
·
verified ·
1 Parent(s): cc122da

Update agent.py

Browse files
Files changed (1) hide show
  1. agent.py +250 -3
agent.py CHANGED
@@ -1,13 +1,16 @@
1
  import os
2
  import base64
 
3
  from langchain_core.messages import HumanMessage, SystemMessage
4
  from langchain_openai import ChatOpenAI
5
  from langchain_community.tools import DuckDuckGoSearchResults
6
  from langchain_community.utilities import DuckDuckGoSearchAPIWrapper
7
  import wikipediaapi
8
  import json
9
- import asyncio
10
- import aiohttp
 
 
11
  from langchain_core.tools import tool
12
  from langgraph.graph import START, StateGraph, MessagesState
13
  from langgraph.prebuilt import tools_condition
@@ -61,9 +64,253 @@ tavily_search_tool = TavilySearch(
61
  topic="general",
62
  )
63
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
64
  tools = [
65
  tavily_search_tool,
66
- search_wiki
 
 
 
 
 
 
 
67
  ]
68
 
69
  def build_graph():
 
1
  import os
2
  import base64
3
+ import tempfile
4
  from langchain_core.messages import HumanMessage, SystemMessage
5
  from langchain_openai import ChatOpenAI
6
  from langchain_community.tools import DuckDuckGoSearchResults
7
  from langchain_community.utilities import DuckDuckGoSearchAPIWrapper
8
  import wikipediaapi
9
  import json
10
+ from urllib.parse import urlparse
11
+ import pytesseract
12
+ from PIL import Image, ImageDraw, ImageFont, ImageEnhance, ImageFilter
13
+ import cmath
14
  from langchain_core.tools import tool
15
  from langgraph.graph import START, StateGraph, MessagesState
16
  from langgraph.prebuilt import tools_condition
 
64
  topic="general",
65
  )
66
 
67
+ @tool
68
+ def save_and_read_file(content: str, filename: Optional[str] = None) -> str:
69
+ """
70
+ Save content to a file and return the path.
71
+ Args:
72
+ content (str): the content to save to the file
73
+ filename (str, optional): the name of the file. If not provided, a random name file will be created.
74
+ """
75
+ temp_dir = tempfile.gettempdir()
76
+ if filename is None:
77
+ temp_file = tempfile.NamedTemporaryFile(delete=False, dir=temp_dir)
78
+ filepath = temp_file.name
79
+ else:
80
+ filepath = os.path.join(temp_dir, filename)
81
+
82
+ with open(filepath, "w") as f:
83
+ f.write(content)
84
+
85
+ return f"File saved to {filepath}. You can read this file to process its contents."
86
+
87
+ @tool
88
+ def download_file_from_url(url: str, filename: Optional[str] = None) -> str:
89
+ """
90
+ Download a file from a URL and save it to a temporary location.
91
+ Args:
92
+ url (str): the URL of the file to download.
93
+ filename (str, optional): the name of the file. If not provided, a random name file will be created.
94
+ """
95
+ try:
96
+ # Parse URL to get filename if not provided
97
+ if not filename:
98
+ path = urlparse(url).path
99
+ filename = os.path.basename(path)
100
+ if not filename:
101
+ filename = f"downloaded_{uuid.uuid4().hex[:8]}"
102
+
103
+ # Create temporary file
104
+ temp_dir = tempfile.gettempdir()
105
+ filepath = os.path.join(temp_dir, filename)
106
+
107
+ # Download the file
108
+ response = requests.get(url, stream=True)
109
+ response.raise_for_status()
110
+
111
+ # Save the file
112
+ with open(filepath, "wb") as f:
113
+ for chunk in response.iter_content(chunk_size=8192):
114
+ f.write(chunk)
115
+
116
+ return f"File downloaded to {filepath}. You can read this file to process its contents."
117
+ except Exception as e:
118
+ return f"Error downloading file: {str(e)}"
119
+
120
+ @tool
121
+ def sum(a: int, b:int) -> int:
122
+ """Sum up two numbers.
123
+ Args:
124
+ a: first int
125
+ b: second int
126
+ """
127
+ return a + b
128
+
129
+ @tool
130
+ def extract_text_from_image(image_path: str) -> str:
131
+ """
132
+ Extract text from an image using OCR library pytesseract (if available).
133
+ Args:
134
+ image_path (str): the path to the image file.
135
+ """
136
+ try:
137
+ # Open the image
138
+ image = Image.open(image_path)
139
+
140
+ # Extract text from the image
141
+ text = pytesseract.image_to_string(image)
142
+
143
+ return f"Extracted text from image:\n\n{text}"
144
+ except Exception as e:
145
+ return f"Error extracting text from image: {str(e)}"
146
+
147
+
148
+ @tool
149
+ def analyze_csv_file(file_path: str, query: str) -> str:
150
+ """
151
+ Analyze a CSV file using pandas and answer a question about it.
152
+ Args:
153
+ file_path (str): the path to the CSV file.
154
+ query (str): Question about the data
155
+ """
156
+ try:
157
+ # Read the CSV file
158
+ df = pd.read_csv(file_path)
159
+
160
+ # Run various analyses based on the query
161
+ result = f"CSV file loaded with {len(df)} rows and {len(df.columns)} columns.\n"
162
+ result += f"Columns: {', '.join(df.columns)}\n\n"
163
+
164
+ # Add summary statistics
165
+ result += "Summary statistics:\n"
166
+ result += str(df.describe())
167
+
168
+ return result
169
+
170
+ except Exception as e:
171
+ return f"Error analyzing CSV file: {str(e)}"
172
+
173
+
174
+ @tool
175
+ def analyze_excel_file(file_path: str, query: str) -> str:
176
+ """
177
+ Analyze an Excel file using pandas and answer a question about it.
178
+ Args:
179
+ file_path (str): the path to the Excel file.
180
+ query (str): Question about the data
181
+ """
182
+ try:
183
+ # Read the Excel file
184
+ df = pd.read_excel(file_path)
185
+
186
+ # Run various analyses based on the query
187
+ result = (
188
+ f"Excel file loaded with {len(df)} rows and {len(df.columns)} columns.\n"
189
+ )
190
+ result += f"Columns: {', '.join(df.columns)}\n\n"
191
+
192
+ # Add summary statistics
193
+ result += "Summary statistics:\n"
194
+ result += str(df.describe())
195
+
196
+ return result
197
+
198
+ except Exception as e:
199
+ return f"Error analyzing Excel file: {str(e)}"
200
+
201
+
202
+ @tool
203
+ def analyze_image(image_base64: str) -> Dict[str, Any]:
204
+ """
205
+ Analyze basic properties of an image (size, mode, color analysis, thumbnail preview).
206
+ Args:
207
+ image_base64 (str): Base64 encoded image string
208
+ Returns:
209
+ Dictionary with analysis result
210
+ """
211
+ try:
212
+ img = decode_image(image_base64)
213
+ width, height = img.size
214
+ mode = img.mode
215
+
216
+ if mode in ("RGB", "RGBA"):
217
+ arr = np.array(img)
218
+ avg_colors = arr.mean(axis=(0, 1))
219
+ dominant = ["Red", "Green", "Blue"][np.argmax(avg_colors[:3])]
220
+ brightness = avg_colors.mean()
221
+ color_analysis = {
222
+ "average_rgb": avg_colors.tolist(),
223
+ "brightness": brightness,
224
+ "dominant_color": dominant,
225
+ }
226
+ else:
227
+ color_analysis = {"note": f"No color analysis for mode {mode}"}
228
+
229
+ thumbnail = img.copy()
230
+ thumbnail.thumbnail((100, 100))
231
+ thumb_path = save_image(thumbnail, "thumbnails")
232
+ thumbnail_base64 = encode_image(thumb_path)
233
+
234
+ return {
235
+ "dimensions": (width, height),
236
+ "mode": mode,
237
+ "color_analysis": color_analysis,
238
+ "thumbnail": thumbnail_base64,
239
+ }
240
+ except Exception as e:
241
+ return {"error": str(e)}
242
+
243
+ @tool
244
+ def transform_image(
245
+ image_base64: str, operation: str, params: Optional[Dict[str, Any]] = None
246
+ ) -> Dict[str, Any]:
247
+ """
248
+ Apply transformations: resize, rotate, crop, flip, brightness, contrast, blur, sharpen, grayscale.
249
+ Args:
250
+ image_base64 (str): Base64 encoded input image
251
+ operation (str): Transformation operation
252
+ params (Dict[str, Any], optional): Parameters for the operation
253
+ Returns:
254
+ Dictionary with transformed image (base64)
255
+ """
256
+ try:
257
+ img = decode_image(image_base64)
258
+ params = params or {}
259
+
260
+ if operation == "resize":
261
+ img = img.resize(
262
+ (
263
+ params.get("width", img.width // 2),
264
+ params.get("height", img.height // 2),
265
+ )
266
+ )
267
+ elif operation == "rotate":
268
+ img = img.rotate(params.get("angle", 90), expand=True)
269
+ elif operation == "crop":
270
+ img = img.crop(
271
+ (
272
+ params.get("left", 0),
273
+ params.get("top", 0),
274
+ params.get("right", img.width),
275
+ params.get("bottom", img.height),
276
+ )
277
+ )
278
+ elif operation == "flip":
279
+ if params.get("direction", "horizontal") == "horizontal":
280
+ img = img.transpose(Image.FLIP_LEFT_RIGHT)
281
+ else:
282
+ img = img.transpose(Image.FLIP_TOP_BOTTOM)
283
+ elif operation == "adjust_brightness":
284
+ img = ImageEnhance.Brightness(img).enhance(params.get("factor", 1.5))
285
+ elif operation == "adjust_contrast":
286
+ img = ImageEnhance.Contrast(img).enhance(params.get("factor", 1.5))
287
+ elif operation == "blur":
288
+ img = img.filter(ImageFilter.GaussianBlur(params.get("radius", 2)))
289
+ elif operation == "sharpen":
290
+ img = img.filter(ImageFilter.SHARPEN)
291
+ elif operation == "grayscale":
292
+ img = img.convert("L")
293
+ else:
294
+ return {"error": f"Unknown operation: {operation}"}
295
+
296
+ result_path = save_image(img)
297
+ result_base64 = encode_image(result_path)
298
+ return {"transformed_image": result_base64}
299
+
300
+ except Exception as e:
301
+ return {"error": str(e)}
302
+
303
+
304
  tools = [
305
  tavily_search_tool,
306
+ search_wiki,
307
+ save_and_read_file,
308
+ transform_image,
309
+ analyze_image,
310
+ analyze_excel_file,
311
+ analyze_csv_file,
312
+ extract_text_from_image,
313
+ download_file_from_url
314
  ]
315
 
316
  def build_graph():