agent_setup.py

from langchain.agents import initialize_agent, AgentType, Tool
from langchain_google_genai import ChatGoogleGenerativeAI
from tools import get_current_weather, get_current_time, make_work_order

def setup_agent(vectorstore):
    def vector_search(query: str) -> str:
        retriever = vectorstore.as_retriever(
            search_type="similarity", 
            search_kwargs={"k": 2}
        )
        docs = retriever.get_relevant_documents(query)
        if not docs:
            return "No relevant documents found."
        return "\n\n".join([doc.page_content for doc in docs])

    weather_tool = Tool(
        name="GetWeather",
        func=get_current_weather,
        description="Returns current weather for a location."
    )
    time_tool = Tool(
        name="GetTime",
        func=get_current_time,
        description="Returns the current time."
    )
    work_order_tool = Tool(
        name="MakeWorkOrder",
        func=make_work_order,
        description="Creates a work order from JSON data."
    )
    vector_tool = Tool(
        name="VectorSearch",
        func=vector_search,
        description="Searches the content of the uploaded document."
    )

    tools = [weather_tool, time_tool, work_order_tool, vector_tool]
    llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash-latest")
    agent = initialize_agent(
        tools=tools,
        llm=llm,
        agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
        verbose=True,
        handle_parsing_errors=True
    )
    return agent