philincloud commited on
Commit
2ea78f5
·
verified ·
1 Parent(s): 1f94fdd

Update langgraph_agent.py

Browse files
Files changed (1) hide show
  1. langgraph_agent.py +57 -27
langgraph_agent.py CHANGED
@@ -1,8 +1,12 @@
1
  import os
2
  import io
3
  import contextlib
4
- import pandas as pd # Added for Excel file handling
5
- from typing import Dict, List, Union # Added for type hinting
 
 
 
 
6
 
7
  from langgraph.graph import START, StateGraph, MessagesState
8
  from langgraph.prebuilt import tools_condition, ToolNode
@@ -45,23 +49,22 @@ def modulus(a: int, b: int) -> int:
45
  def wiki_search(query: str) -> dict:
46
  """Search Wikipedia for a query and return up to 2 documents."""
47
  try:
48
- docs = WikipediaLoader(query=query, load_max_docs=2, lang="en").load() # Added lang="en" for clarity
49
  if not docs:
50
  return {"wiki_results": f"No documents found on Wikipedia for '{query}'."}
51
  formatted = "\n\n---\n\n".join(
52
- f'<Document source="{d.metadata.get("source", "N/A")}"/>\n{d.page_content}' # Added .get for safety
53
  for d in docs
54
  )
55
  return {"wiki_results": formatted}
56
  except Exception as e:
57
- # Log the full error for debugging if possible
58
  print(f"Error in wiki_search tool: {e}")
59
  return {"wiki_results": f"Error occurred while searching Wikipedia for '{query}'. Details: {str(e)}"}
60
 
61
  @tool
62
  def web_search(query: str) -> dict:
63
  """Perform a web search (via Tavily) and return up to 3 results."""
64
- try: # Added try-except block for robustness
65
  docs = TavilySearchResults(max_results=3).invoke(query=query)
66
  formatted = "\n\n---\n\n".join(
67
  f'<Document source="{d.metadata["source"]}"/>\n{d.page_content}'
@@ -82,30 +85,40 @@ def arvix_search(query: str) -> dict:
82
  )
83
  return {"arvix_results": formatted}
84
 
 
 
 
 
 
 
 
 
85
  @tool
86
  def read_file_content(file_path: str) -> Dict[str, str]:
87
  """
88
- Reads the content of a file and returns it.
89
- Supports text (.txt), Python (.py), and Excel (.xlsx) files.
90
- For other file types, returns a message indicating limited support.
91
  """
92
  try:
93
  _, file_extension = os.path.splitext(file_path)
94
- content = ""
95
- if file_extension.lower() in (".txt", ".py"):
 
96
  with open(file_path, "r", encoding="utf-8") as f:
97
  content = f.read()
98
- elif file_extension.lower() == ".xlsx":
99
- # Ensure pandas is installed for this.
100
  df = pd.read_excel(file_path)
101
- content = df.to_string() # Convert Excel to string representation
102
- elif file_extension.lower() == ".mp3":
103
- content = "Audio file provided. Unable to directly process audio. Consider using a transcription service if available."
104
- elif file_extension.lower() == ".png":
105
- content = "Image file provided. Unable to directly process images. Consider using an OCR or image analysis service if available."
 
 
 
106
  else:
107
- content = f"Unsupported file type: {file_extension}. Only .txt, .py, and .xlsx files are fully supported for reading content."
108
- return {"file_content": content, "file_name": file_path}
109
  except FileNotFoundError:
110
  return {"file_error": f"File not found: {file_path}. Please ensure the file exists in the environment."}
111
  except Exception as e:
@@ -118,10 +131,8 @@ def python_interpreter(code: str) -> Dict[str, str]:
118
  If there's an error during execution, it returns the error message.
119
  """
120
  old_stdout = io.StringIO()
121
- # Redirect stdout to capture print statements
122
  with contextlib.redirect_stdout(old_stdout):
123
  try:
124
- # Create a dictionary to hold the execution scope for exec
125
  exec_globals = {}
126
  exec_locals = {}
127
  exec(code, exec_globals, exec_locals)
@@ -130,6 +141,24 @@ def python_interpreter(code: str) -> Dict[str, str]:
130
  except Exception as e:
131
  return {"execution_error": str(e)}
132
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
133
 
134
  API_KEY = os.getenv("GEMINI_API_KEY")
135
  HF_SPACE_TOKEN = os.getenv("HF_SPACE_TOKEN")
@@ -139,8 +168,10 @@ GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
139
  tools = [
140
  multiply, add, subtract, divide, modulus,
141
  wiki_search, web_search, arvix_search,
142
- read_file_content, # Added new tool
143
- python_interpreter, # Added new tool
 
 
144
  ]
145
 
146
 
@@ -153,7 +184,7 @@ def build_graph(provider: str = "gemini"):
153
  """Build the LangGraph agent with chosen LLM (default: Gemini)."""
154
  if provider == "gemini":
155
  llm = ChatGoogleGenerativeAI(
156
- model= "gemini-2.5-flash-preview-05-20",
157
  temperature=1.0,
158
  max_retries=2,
159
  api_key=GEMINI_API_KEY,
@@ -168,7 +199,7 @@ def build_graph(provider: str = "gemini"):
168
  temperature=0,
169
  )
170
  else:
171
- raise ValueError("Invalid provider. Choose 'openai' or 'huggingface'.")
172
 
173
  llm_with_tools = llm.bind_tools(tools)
174
 
@@ -189,4 +220,3 @@ if __name__ == "__main__":
189
  # This block is intentionally left empty as per user request to remove examples.
190
  # Your agent will interact with the graph by invoking it with messages.
191
  pass
192
-
 
1
  import os
2
  import io
3
  import contextlib
4
+ import pandas as pd
5
+ from typing import Dict, List, Union
6
+
7
+ # New imports for image and audio processing
8
+ from PIL import Image as PILImage # Used for type checking/potential future local processing
9
+ from huggingface_hub import InferenceClient
10
 
11
  from langgraph.graph import START, StateGraph, MessagesState
12
  from langgraph.prebuilt import tools_condition, ToolNode
 
49
  def wiki_search(query: str) -> dict:
50
  """Search Wikipedia for a query and return up to 2 documents."""
51
  try:
52
+ docs = WikipediaLoader(query=query, load_max_docs=2, lang="en").load()
53
  if not docs:
54
  return {"wiki_results": f"No documents found on Wikipedia for '{query}'."}
55
  formatted = "\n\n---\n\n".join(
56
+ f'<Document source="{d.metadata.get("source", "N/A")}"/>\n{d.page_content}'
57
  for d in docs
58
  )
59
  return {"wiki_results": formatted}
60
  except Exception as e:
 
61
  print(f"Error in wiki_search tool: {e}")
62
  return {"wiki_results": f"Error occurred while searching Wikipedia for '{query}'. Details: {str(e)}"}
63
 
64
  @tool
65
  def web_search(query: str) -> dict:
66
  """Perform a web search (via Tavily) and return up to 3 results."""
67
+ try:
68
  docs = TavilySearchResults(max_results=3).invoke(query=query)
69
  formatted = "\n\n---\n\n".join(
70
  f'<Document source="{d.metadata["source"]}"/>\n{d.page_content}'
 
85
  )
86
  return {"arvix_results": formatted}
87
 
88
+ # Initialize Hugging Face Inference Client
89
+ HF_API_TOKEN = os.getenv("HF_API_TOKEN")
90
+ HF_INFERENCE_CLIENT = None
91
+ if HF_API_TOKEN:
92
+ HF_INFERENCE_CLIENT = InferenceClient(token=HF_API_TOKEN)
93
+ else:
94
+ print("WARNING: HF_API_TOKEN not set. Image tools will not function.")
95
+
96
  @tool
97
  def read_file_content(file_path: str) -> Dict[str, str]:
98
  """
99
+ Reads the content of a file and returns its primary information.
100
+ For text/code/excel, returns content. For media, returns a prompt to use specific tools.
 
101
  """
102
  try:
103
  _, file_extension = os.path.splitext(file_path)
104
+ file_extension = file_extension.lower()
105
+
106
+ if file_extension in (".txt", ".py"):
107
  with open(file_path, "r", encoding="utf-8") as f:
108
  content = f.read()
109
+ return {"file_type": "text/code", "file_name": file_path, "file_content": content}
110
+ elif file_extension == ".xlsx":
111
  df = pd.read_excel(file_path)
112
+ content = df.to_string()
113
+ return {"file_type": "excel", "file_name": file_path, "file_content": content}
114
+ elif file_extension in (".jpeg", ".jpg", ".png"):
115
+ # Indicate that it's an image and needs to be described by a specific tool
116
+ return {"file_type": "image", "file_name": file_path, "file_content": f"Image file '{file_path}' detected. Use 'describe_image' tool to get a textual description."}
117
+ elif file_extension == ".mp3":
118
+ # Indicate that it's an audio file and the LLM should process it natively
119
+ return {"file_type": "audio", "file_name": file_path, "file_content": f"Audio file '{file_path}' detected. The LLM should process this natively."}
120
  else:
121
+ return {"file_type": "unsupported", "file_name": file_path, "file_content": f"Unsupported file type: {file_extension}. Only .txt, .py, .xlsx, .jpeg, .jpg, .png, .mp3 files are recognized."}
 
122
  except FileNotFoundError:
123
  return {"file_error": f"File not found: {file_path}. Please ensure the file exists in the environment."}
124
  except Exception as e:
 
131
  If there's an error during execution, it returns the error message.
132
  """
133
  old_stdout = io.StringIO()
 
134
  with contextlib.redirect_stdout(old_stdout):
135
  try:
 
136
  exec_globals = {}
137
  exec_locals = {}
138
  exec(code, exec_globals, exec_locals)
 
141
  except Exception as e:
142
  return {"execution_error": str(e)}
143
 
144
+ @tool
145
+ def describe_image(image_path: str) -> Dict[str, str]:
146
+ """
147
+ Generates a textual description for an image file (JPEG, JPG, PNG) using an image-to-text model
148
+ from the Hugging Face Inference API. Requires HF_API_TOKEN environment variable to be set.
149
+ """
150
+ if not HF_INFERENCE_CLIENT:
151
+ return {"error": "Hugging Face API token not configured for image description. Cannot use this tool."}
152
+ try:
153
+ with open(image_path, "rb") as f:
154
+ image_bytes = f.read()
155
+ description = HF_INFERENCE_CLIENT.image_to_text(image_bytes)
156
+ return {"image_description": description, "image_path": image_path}
157
+ except FileNotFoundError:
158
+ return {"error": f"Image file not found: {image_path}. Please ensure the file exists."}
159
+ except Exception as e:
160
+ return {"error": f"Error describing image {image_path}: {str(e)}"}
161
+
162
 
163
  API_KEY = os.getenv("GEMINI_API_KEY")
164
  HF_SPACE_TOKEN = os.getenv("HF_SPACE_TOKEN")
 
168
  tools = [
169
  multiply, add, subtract, divide, modulus,
170
  wiki_search, web_search, arvix_search,
171
+ read_file_content,
172
+ python_interpreter,
173
+ describe_image, # Added new tool
174
+ # transcribe_audio, # Removed as per user request
175
  ]
176
 
177
 
 
184
  """Build the LangGraph agent with chosen LLM (default: Gemini)."""
185
  if provider == "gemini":
186
  llm = ChatGoogleGenerativeAI(
187
+ model= "gemini-1.5-flash-preview-05-20", # This model is capable of native audio processing
188
  temperature=1.0,
189
  max_retries=2,
190
  api_key=GEMINI_API_KEY,
 
199
  temperature=0,
200
  )
201
  else:
202
+ raise ValueError("Invalid provider. Choose 'gemini' or 'huggingface'.")
203
 
204
  llm_with_tools = llm.bind_tools(tools)
205
 
 
220
  # This block is intentionally left empty as per user request to remove examples.
221
  # Your agent will interact with the graph by invoking it with messages.
222
  pass