agent.py

from typing import List, Optional, Tuple
from langchain.agents import AgentExecutor, AgentType, ConversationalAgent
from langchain.agents import initialize_agent  # type: ignore
from langchain.agents.agent_types import AgentType
from langchain.tools import BaseTool
from langchain.memory import ConversationBufferMemory

# from langchain_openai import ChatOpenAI  # type: ignore

from tools.fallback import FallbackTool
from tools.python_REPL import PythonREPLTool

""" Tool Imports"""
from tools.guarded_retriever import GuardedRetrieverTool
from tools.reasoning import ReasoningTool
from tools.retriever import RetrieverTool
from tools.websearch import WebSearchTool
from tools.calculator import CalculatorTool
from tools.longterm_memory import LongTermMemoryTool
from longterm_memory import LongTermMemory

# from transformers import pipeline
from langchain_ollama import OllamaLLM
from langchain_core.language_models import BaseLanguageModel  # type: ignore

import warnings
from langchain_core._api.deprecation import LangChainDeprecationWarning

warnings.filterwarnings("ignore", category=LangChainDeprecationWarning)

from dotenv import load_dotenv

load_dotenv()


def load_llm():
    # return OllamaLLM(model="gemma3", temperature=0.3)
    return OllamaLLM(model="llama3.1", temperature=0.3)
    # return ChatOpenAI(model="gpt-3.5-turbo", temperature=0.3)


def get_tools(memory: ConversationBufferMemory) -> List[BaseTool]:
    retriever = RetrieverTool()
    guarded_retriever_tool = GuardedRetrieverTool(retriever=retriever)
    websearch_tool = WebSearchTool()
    calculator_tool = CalculatorTool()
    fallback_tool = FallbackTool()
    longterm_store = LongTermMemory()
    longterm_tool = LongTermMemoryTool(memory_store=longterm_store)
    reasoning_tool = ReasoningTool(name="Reasoning", memory=memory)
    python_repl_tool = PythonREPLTool()

    return [
        websearch_tool,
        python_repl_tool,
        calculator_tool,
        guarded_retriever_tool,
        longterm_tool,
        reasoning_tool,
        fallback_tool,
    ]


# Initialize the LLM and tools
def get_agent(
    llm: Optional[BaseLanguageModel] = None,
    memory: Optional[ConversationBufferMemory] = None,
) -> Tuple[AgentExecutor, ConversationBufferMemory]:
    buffer = memory or ConversationBufferMemory(
        memory_key="chat_history",
        return_messages=False,
        output_key="output",
    )
    tools = get_tools(buffer)
    llm_to_use = llm or load_llm()

    format_instructions = (
        "You MUST follow this format:\n"
        "Thought: brief reflection on whether you need a tool.\n"
        "Action: tool name from [{tool_names}] (omit if not using a tool).\n"
        "Action Input: exact JSON or text arguments for the tool.\n"
        "Observation: (filled automatically after tool execution).\n"
        "Repeat Thought/Action/Observation until you can answer directly.\n"
        "When you are ready to reply to the user, finish with:\n"
        "Thought: Do I need to use a tool? No\n"
        "Final Answer: <concise response to the user>"
    )

    guidance = (
        "You are a resilient assistant named Cortex that can consult tools when needed. "
        "Adhere to these guardrails:\n"
        "- Prefer to answer simple conversational prompts directly without tools.\n"
        "- Only use python_repl when the user explicitly asks for code execution or complex calculation.\n"
        "- If a web search is rate limited or fails twice, stop retrying and explain the limitation.\n"
        "- Use calculator for straightforward math instead of python_repl.\n"
        "- When tools fail or data is unavailable, give your best effort answer with the information you have.\n"
        "- Never leave the user without a Final Answer."
    )

    agent_prompt = ConversationalAgent.create_prompt(
        tools=tools,
        prefix=(
            f"{guidance}\n\n"
            "Conversation so far:\n"
            "{chat_history}\n\n"
            "Current user request:\n"
        ),
        format_instructions=format_instructions,
        suffix="{input}",
        input_variables=["input", "chat_history", "agent_scratchpad"],
    )

    agent = initialize_agent(
        tools=tools,
        llm=llm_to_use,
        agent_type=AgentType.CONVERSATIONAL_REACT_DESCRIPTION,
        memory=buffer,
        agent_kwargs={"prompt": agent_prompt},
        verbose=True,
        handle_parsing_errors=True,
        max_iterations=4,
        early_stopping_method="generate",
    )

    return agent, buffer