ready: update src/agent.py
This commit is contained in:
+3
-4
@@ -1,7 +1,7 @@
|
||||
import os
|
||||
from typing import Dict, List
|
||||
|
||||
from langchain.agents import AgentExecutor, create_agent
|
||||
from langchain.agents import create_agent
|
||||
from langchain_ollama import Ollama
|
||||
from langchain_core.prompts import ChatPromptTemplate
|
||||
from langchain_core.messages import HumanMessage
|
||||
@@ -49,7 +49,6 @@ PROMPT = ChatPromptTemplate.from_messages([
|
||||
# Create agent
|
||||
llm = Ollama(model="llama2")
|
||||
agent = create_agent(llm, TOOL_DEFINITIONS, PROMPT)
|
||||
executor = AgentExecutor(agent=agent, tools=TOOL_DEFINITIONS, verbose=True)
|
||||
|
||||
# CLI entry point
|
||||
if __name__ == "__main__":
|
||||
@@ -58,14 +57,14 @@ if __name__ == "__main__":
|
||||
parser.add_argument("--question", type=str, help="Question to ask the bot")
|
||||
args = parser.parse_args()
|
||||
if args.question:
|
||||
result = executor.invoke({"input": args.question})
|
||||
result = agent.invoke({"input": args.question})
|
||||
print(result["output"])
|
||||
else:
|
||||
print("Enter questions (Ctrl-D to exit):")
|
||||
try:
|
||||
while True:
|
||||
q = input("Q: ")
|
||||
res = executor.invoke({"input": q})
|
||||
res = agent.invoke({"input": q})
|
||||
print("A:", res["output"])
|
||||
except EOFError:
|
||||
pass
|
||||
|
||||
Reference in New Issue
Block a user