41 lines
1.2 KiB
Python
41 lines
1.2 KiB
Python
"""
|
|
Agent creation using LangChain create_agent.
|
|
"""
|
|
import os
|
|
from langchain_ollama import OllamaLLM
|
|
from langchain_core.messages import HumanMessage
|
|
from langchain.agents import create_agent
|
|
from langchain.tools import tool
|
|
from tools import search_knowledge_base, add_to_knowledge_base
|
|
|
|
# LLM via Ollama
|
|
llm = OllamaLLM(model="llama3")
|
|
|
|
# System prompt instructing to use knowledge base tools
|
|
SYSTEM_PROMPT = """
|
|
You are a helpful assistant that can store and retrieve information.
|
|
Use the provided tools search_knowledge_base and add_to_knowledge_base.
|
|
When answering, prefer to call the tools if needed.
|
|
"""
|
|
|
|
def create_rag_agent():
|
|
agent = create_agent(
|
|
llm=llm,
|
|
tools=[search_knowledge_base, add_to_knowledge_base],
|
|
system_prompt=SYSTEM_PROMPT,
|
|
)
|
|
return agent
|
|
|
|
if __name__ == "__main__":
|
|
ag = create_rag_agent()
|
|
# Simple demo loop
|
|
while True:
|
|
user_input = input("User: ")
|
|
if user_input.lower() in ("quit", "exit"):
|
|
break
|
|
result = ag.ainvoke(
|
|
{"messages": [HumanMessage(content=user_input)]},
|
|
{"configurable": {"thread_id": "demo"}},
|
|
)
|
|
print(result["messages"][-1].content)
|