Yongkang ZOU commited on
Commit
bfdb8e9
·
1 Parent(s): 4e6c049

update tools

Browse files
Files changed (2) hide show
  1. agent.py +209 -1
  2. requirements.txt +2 -1
agent.py CHANGED
@@ -16,6 +16,19 @@ from langchain_community.vectorstores import SupabaseVectorStore
16
  from langchain_openai import ChatOpenAI
17
  from langchain_core.documents import Document
18
  import json
 
 
 
 
 
 
 
 
 
 
 
 
 
19
 
20
  load_dotenv()
21
 
@@ -47,6 +60,28 @@ def modulus(a: int, b: int) -> int:
47
  """Get remainder of a divided by b."""
48
  return a % b
49
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
50
  @tool
51
  def wiki_search(query: str) -> str:
52
  """Search Wikipedia for a query (max 2 results)."""
@@ -80,10 +115,183 @@ def read_excel_file(path: str) -> str:
80
  return content.strip()
81
  except Exception as e:
82
  return f"Error reading Excel file: {str(e)}"
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
83
 
 
 
 
 
 
 
 
84
 
 
 
 
 
 
 
 
 
 
 
 
85
 
86
- tools = [multiply, add, subtract, divide, modulus, wiki_search, web_search, arvix_search, read_excel_file]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
87
 
88
  # ------------------- SYSTEM PROMPT -------------------
89
  system_prompt_path = "system_prompt.txt"
 
16
  from langchain_openai import ChatOpenAI
17
  from langchain_core.documents import Document
18
  import json
19
+ import pdfplumber
20
+ import pandas as pd
21
+ from transformers import BlipProcessor, BlipForConditionalGeneration
22
+ from PIL import Image
23
+ import torch
24
+ import cmath
25
+ from code_interpreter import CodeInterpreter
26
+ import uuid
27
+ import tempfile
28
+ import requests
29
+ from urllib.parse import urlparse
30
+ from typing import Optional
31
+
32
 
33
  load_dotenv()
34
 
 
60
  """Get remainder of a divided by b."""
61
  return a % b
62
 
63
+ @tool
64
+ def square_root(a: float) -> float | complex:
65
+ """
66
+ Get the square root of a number.
67
+ Args:
68
+ a (float): the number to get the square root of
69
+ """
70
+ if a >= 0:
71
+ return a**0.5
72
+ return cmath.sqrt(a)
73
+
74
+ @tool
75
+ def power(a: float, b: float) -> float:
76
+ """
77
+ Get the power of two numbers.
78
+ Args:
79
+ a (float): the first number
80
+ b (float): the second number
81
+ """
82
+ return a**b
83
+
84
+
85
  @tool
86
  def wiki_search(query: str) -> str:
87
  """Search Wikipedia for a query (max 2 results)."""
 
115
  return content.strip()
116
  except Exception as e:
117
  return f"Error reading Excel file: {str(e)}"
118
+
119
+ @tool
120
+ def extract_text_from_pdf(path: str) -> str:
121
+ """Extract text from a PDF file given its local path."""
122
+ try:
123
+ text = ""
124
+ with pdfplumber.open(path) as pdf:
125
+ for page in pdf.pages[:5]: # 限前5页,避免过大
126
+ page_text = page.extract_text()
127
+ if page_text:
128
+ text += page_text + "\n\n"
129
+ return text.strip() if text else "No text extracted from PDF."
130
+ except Exception as e:
131
+ return f"Error reading PDF: {str(e)}"
132
+
133
+ # 初始化模型(首次加载可能稍慢)
134
+ processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
135
+ model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
136
+
137
+ @tool
138
+ def blip_image_caption(image_path: str) -> str:
139
+ """Generate a description for an image using BLIP."""
140
+ try:
141
+ image = Image.open(image_path).convert("RGB")
142
+ inputs = processor(image, return_tensors="pt")
143
+ with torch.no_grad():
144
+ out = model.generate(**inputs)
145
+ caption = processor.decode(out[0], skip_special_tokens=True)
146
+ return caption
147
+ except Exception as e:
148
+ return f"Failed to process image with BLIP: {str(e)}"
149
+
150
+ @tool
151
+ def execute_code_multilang(code: str, language: str = "python") -> str:
152
+ """Execute code in multiple languages (Python, Bash, SQL, C, Java) and return results.
153
+ Args:
154
+ code (str): The source code to execute.
155
+ language (str): The language of the code. Supported: "python", "bash", "sql", "c", "java".
156
+ Returns:
157
+ A string summarizing the execution results (stdout, stderr, errors, plots, dataframes if any).
158
+ """
159
+ supported_languages = ["python", "bash", "sql", "c", "java"]
160
+ language = language.lower()
161
+ interpreter_instance = CodeInterpreter()
162
+
163
+ if language not in supported_languages:
164
+ return f"❌ Unsupported language: {language}. Supported languages are: {', '.join(supported_languages)}"
165
+
166
+ result = interpreter_instance.execute_code(code, language=language)
167
+
168
+ response = []
169
+
170
+ if result["status"] == "success":
171
+ response.append(f"✅ Code executed successfully in **{language.upper()}**")
172
+
173
+ if result.get("stdout"):
174
+ response.append(
175
+ "\n**Standard Output:**\n```\n" + result["stdout"].strip() + "\n```"
176
+ )
177
+
178
+ if result.get("stderr"):
179
+ response.append(
180
+ "\n**Standard Error (if any):**\n```\n"
181
+ + result["stderr"].strip()
182
+ + "\n```"
183
+ )
184
+
185
+ if result.get("result") is not None:
186
+ response.append(
187
+ "\n**Execution Result:**\n```\n"
188
+ + str(result["result"]).strip()
189
+ + "\n```"
190
+ )
191
 
192
+ if result.get("dataframes"):
193
+ for df_info in result["dataframes"]:
194
+ response.append(
195
+ f"\n**DataFrame `{df_info['name']}` (Shape: {df_info['shape']})**"
196
+ )
197
+ df_preview = pd.DataFrame(df_info["head"])
198
+ response.append("First 5 rows:\n```\n" + str(df_preview) + "\n```")
199
 
200
+ if result.get("plots"):
201
+ response.append(
202
+ f"\n**Generated {len(result['plots'])} plot(s)** (Image data returned separately)"
203
+ )
204
+
205
+ else:
206
+ response.append(f"❌ Code execution failed in **{language.upper()}**")
207
+ if result.get("stderr"):
208
+ response.append(
209
+ "\n**Error Log:**\n```\n" + result["stderr"].strip() + "\n```"
210
+ )
211
 
212
+ return "\n".join(response)
213
+
214
+ @tool
215
+ def save_and_read_file(content: str, filename: Optional[str] = None) -> str:
216
+ """
217
+ Save content to a file and return the path.
218
+ Args:
219
+ content (str): the content to save to the file
220
+ filename (str, optional): the name of the file. If not provided, a random name file will be created.
221
+ """
222
+ temp_dir = tempfile.gettempdir()
223
+ if filename is None:
224
+ temp_file = tempfile.NamedTemporaryFile(delete=False, dir=temp_dir)
225
+ filepath = temp_file.name
226
+ else:
227
+ filepath = os.path.join(temp_dir, filename)
228
+
229
+ with open(filepath, "w") as f:
230
+ f.write(content)
231
+
232
+ return f"File saved to {filepath}. You can read this file to process its contents."
233
+
234
+
235
+ @tool
236
+ def download_file_from_url(url: str, filename: Optional[str] = None) -> str:
237
+ """
238
+ Download a file from a URL and save it to a temporary location.
239
+ Args:
240
+ url (str): the URL of the file to download.
241
+ filename (str, optional): the name of the file. If not provided, a random name file will be created.
242
+ """
243
+ try:
244
+ # Parse URL to get filename if not provided
245
+ if not filename:
246
+ path = urlparse(url).path
247
+ filename = os.path.basename(path)
248
+ if not filename:
249
+ filename = f"downloaded_{uuid.uuid4().hex[:8]}"
250
+
251
+ # Create temporary file
252
+ temp_dir = tempfile.gettempdir()
253
+ filepath = os.path.join(temp_dir, filename)
254
+
255
+ # Download the file
256
+ response = requests.get(url, stream=True)
257
+ response.raise_for_status()
258
+
259
+ # Save the file
260
+ with open(filepath, "wb") as f:
261
+ for chunk in response.iter_content(chunk_size=8192):
262
+ f.write(chunk)
263
+
264
+ return f"File downloaded to {filepath}. You can read this file to process its contents."
265
+ except Exception as e:
266
+ return f"Error downloading file: {str(e)}"
267
+
268
+ @tool
269
+ def analyze_csv_file(file_path: str, query: str) -> str:
270
+ """
271
+ Analyze a CSV file using pandas and answer a question about it.
272
+ Args:
273
+ file_path (str): the path to the CSV file.
274
+ query (str): Question about the data
275
+ """
276
+ try:
277
+ # Read the CSV file
278
+ df = pd.read_csv(file_path)
279
+
280
+ # Run various analyses based on the query
281
+ result = f"CSV file loaded with {len(df)} rows and {len(df.columns)} columns.\n"
282
+ result += f"Columns: {', '.join(df.columns)}\n\n"
283
+
284
+ # Add summary statistics
285
+ result += "Summary statistics:\n"
286
+ result += str(df.describe())
287
+
288
+ return result
289
+
290
+ except Exception as e:
291
+ return f"Error analyzing CSV file: {str(e)}"
292
+ tools = [multiply, add, subtract, divide, modulus,
293
+ wiki_search, web_search, arvix_search, read_excel_file, extract_text_from_pdf,
294
+ blip_image_caption, execute_code_multilang, save_and_read_file, download_file_from_url, analyze_csv_file]
295
 
296
  # ------------------- SYSTEM PROMPT -------------------
297
  system_prompt_path = "system_prompt.txt"
requirements.txt CHANGED
@@ -24,4 +24,5 @@ beautifulsoup4
24
  transformers
25
  torch
26
  torchvision
27
- pillow
 
 
24
  transformers
25
  torch
26
  torchvision
27
+ pillow
28
+ matplotlib