51 lines
1.7 KiB
Python
51 lines
1.7 KiB
Python
"""
|
||
Agent creation using LangChain 1.x `create_agent` API.
|
||
|
||
This module exposes a single helper `create_agent_executor` that builds an
|
||
AgentExecutor configured with the custom tools defined in :mod:`tools`.
|
||
"""
|
||
|
||
from langchain_ollama import ChatOllama
|
||
from langchain.agents import create_agent, AgentExecutor
|
||
from langchain.tools import BaseTool
|
||
from typing import List
|
||
|
||
from .tools import search, write_file
|
||
|
||
# Define the list of tools that the agent can use
|
||
TOOLS: List[BaseTool] = [search, write_file]
|
||
|
||
# LLM configuration – Ollama local model
|
||
LLM = ChatOllama(model="llama3.1:latest", temperature=0.0)
|
||
|
||
# Create the agent executor using the new LangChain 1.x API
|
||
# The `create_agent` function returns an AgentExecutor instance
|
||
# that can be called like a normal function.
|
||
|
||
def create_agent_executor() -> AgentExecutor:
|
||
"""Instantiate and return an AgentExecutor.
|
||
|
||
The executor is configured with:
|
||
* The Ollama Chat model.
|
||
* The custom tools defined in :mod:`tools`.
|
||
* The default agent type "openai-retrieval-qa" is not used – we rely on
|
||
the automatically selected agent type by `create_agent`.
|
||
"""
|
||
agent = create_agent(
|
||
llm=LLM,
|
||
tools=TOOLS,
|
||
verbose=True,
|
||
)
|
||
# The returned object is already an AgentExecutor
|
||
return agent
|
||
|
||
# Expose the executor for external use
|
||
agent_executor = create_agent_executor()
|
||
|
||
# For convenience, a small helper that runs a single prompt
|
||
|
||
def run_prompt(prompt: str) -> str:
|
||
"""Run the prompt through the agent and return the final answer."""
|
||
result = agent_executor.invoke({"input": prompt})
|
||
# The result is a dict with keys: "output" and possibly others
|
||
return result.get("output", "") |