Final_Assignment_Template / langgraph_agent.py
philincloud's picture
Update langgraph_agent.py
514b444 verified
import os
import io
import contextlib
import pandas as pd
from typing import Dict, List, Union
import re
from PIL import Image as PILImage # Keep PIL for potential future use or if other parts depend on it, but describe_image is removed.
from huggingface_hub import InferenceClient # Keep InferenceClient for other potential HF uses, but describe_image is removed.
from langgraph.graph import START, StateGraph, MessagesState
from langgraph.prebuilt import tools_condition, ToolNode
from langchain_openai import ChatOpenAI
from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
from langchain_community.document_loaders import WikipediaLoader
from langchain_core.messages import SystemMessage, HumanMessage, ToolMessage
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.tools import tool
from langchain_google_community import GoogleSearchAPIWrapper
@tool
def multiply(a: int, b: int) -> int:
"""Multiply two integers."""
return a * b
@tool
def add(a: int, b: int) -> int:
"""Add two integers."""
return a + b
@tool
def subtract(a: int, b: int) -> int:
"""Subtract the second integer from the first."""
return a - b
@tool
def divide(a: int, b: int) -> float:
"""Divide first integer by second; error if divisor is zero."""
if b == 0:
raise ValueError("Cannot divide by zero.")
return a / b
@tool
def modulus(a: int, b: int) -> int:
"""Return the remainder of dividing first integer by second."""
return a % b
@tool
def wiki_search(query: str) -> dict:
"""Search Wikipedia for a query and return up to 2 documents."""
try:
docs = WikipediaLoader(query=query, load_max_docs=5, lang="en", doc_content_chars_max=7000).load()
if not docs:
return {"wiki_results": f"No documents found on Wikipedia for '{query}'."}
formatted = "\n\n---\n\n".join(
f'<Document source="{d.metadata.get("source", "N/A")}"/>\n{d.page_content}'
for d in docs
)
return {"wiki_results": formatted}
except Exception as e:
print(f"Error in wiki_search tool: {e}")
return {"wiki_results": f"Error occurred while searching Wikipedia for '{query}'. Details: {str(e)}"}
search = GoogleSearchAPIWrapper()
@tool
def google_web_search(query: str) -> str:
"""Perform a web search (via Google Custom Search) and return results."""
try:
return search.run(query)
except Exception as e:
print(f"Error in google_web_search tool: {e}")
return f"Error occurred while searching the web for '{query}'. Details: {str(e)}"
HF_API_TOKEN = os.getenv("HF_API_TOKEN")
MODEL = os.getenv("MODEL")
HF_INFERENCE_CLIENT = None
if HF_API_TOKEN:
HF_INFERENCE_CLIENT = InferenceClient(token=HF_API_TOKEN)
else:
print("WARNING: HF_API_TOKEN not set. If any other HF tools are used, they might not function.")
@tool
def read_file_content(file_path: str) -> Dict[str, str]:
"""Reads the content of a file and returns its primary information. For text/code/excel, returns content. For media, indicates it's a blob for LLM processing."""
try:
_, file_extension = os.path.splitext(file_path)
file_extension = file_extension.lower()
# Prioritize handling of video, audio, and image files for direct LLM processing
if file_extension in (".mp4", ".avi", ".mov", ".mkv", ".webm"):
return {"file_type": "video", "file_name": file_path, "file_content": f"Video file '{file_path}' detected. The LLM (Gemini 2.5 Pro) can process this video content directly as a blob."}
elif file_extension == ".mp3":
return {"file_type": "audio", "file_name": file_path, "file_content": f"Audio file '{file_path}' detected. The LLM (Gemini 2.5 Pro) can process this audio content directly as a blob."}
elif file_extension in (".jpeg", ".jpg", ".png"):
return {"file_type": "image", "file_name": file_path, "file_content": f"Image file '{file_path}' detected. The LLM (Gemini 2.5 Pro) can process this image content directly as a blob."}
# Handle text and code files
elif file_extension in (".txt", ".py"):
with open(file_path, "r", encoding="utf-8") as f:
content = f.read()
return {"file_type": "text/code", "file_name": file_path, "file_content": content}
# Handle Excel files
elif file_extension == ".xlsx":
df = pd.read_excel(file_path)
content = df.to_string()
return {"file_type": "excel", "file_name": file_path, "file_content": content}
else:
return {"file_type": "unsupported", "file_name": file_path, "file_content": f"Unsupported file type: {file_extension}. Only .txt, .py, .xlsx, .jpeg, .jpg, .png, .mp3, .mp4, .avi, .mov, .mkv, .webm files are recognized."}
except FileNotFoundError:
return {"file_error": f"File not found: {file_path}. Please ensure the file exists in the environment."}
except Exception as e:
return {"file_error": f"Error reading file {file_path}: {e}"}
@tool
def python_interpreter(code: str) -> Dict[str, str]:
"""Executes Python code and returns its standard output. If there's an error during execution, it returns the error message."""
old_stdout = io.StringIO()
with contextlib.redirect_stdout(old_stdout):
try:
exec_globals = {}
exec_locals = {}
exec(code, exec_globals, exec_locals)
output = old_stdout.getvalue()
return {"execution_result": output.strip()}
except Exception as e:
return {"execution_error": str(e)}
API_KEY = os.getenv("GEMINI_API_KEY")
HF_API_TOKEN = os.getenv("HF_SPACE_TOKEN") # Kept for potential future HF uses, but not for describe_image
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
# Updated tools list without Youtube
tools = [
multiply, add, subtract, divide, modulus,
wiki_search,
google_web_search,
read_file_content,
python_interpreter,
]
with open("prompt.txt", "r", encoding="utf-8") as f:
system_prompt = f.read()
sys_msg = SystemMessage(content=system_prompt)
def build_graph(provider: str = "gemini"):
if provider == "gemini":
llm = ChatGoogleGenerativeAI(
model=MODEL,
temperature=1.0,
max_retries=2,
api_key=GEMINI_API_KEY,
max_tokens=5000
)
elif provider == "huggingface":
llm = ChatHuggingFace(
llm=HuggingFaceEndpoint(
url="https://api-inference.huggingface.co/models/Meta-DeepLearning/llama-2-7b-chat-hf",
),
temperature=0,
)
else:
raise ValueError("Invalid provider. Choose 'gemini' or 'huggingface'.")
llm_with_tools = llm.bind_tools(tools)
def assistant(state: MessagesState):
messages_to_send = [sys_msg] + state["messages"]
llm_response = llm_with_tools.invoke(messages_to_send, {"recursion_limit": 25})
print(f"LLM Raw Response: {llm_response}")
return {"messages": [llm_response]}
builder = StateGraph(MessagesState)
builder.add_node("assistant", assistant)
builder.add_node("tools", ToolNode(tools))
builder.add_edge(START, "assistant")
builder.add_conditional_edges("assistant", tools_condition)
builder.add_edge("tools", "assistant")
return builder.compile()
if __name__ == "__main__":
pass