32 lines
904 B
Python
32 lines
904 B
Python
"""
|
|
Agent creation with RAG integration.
|
|
"""
|
|
import os
|
|
from typing import List
|
|
|
|
from langchain_openai import ChatOpenAI
|
|
from langchain.agents import create_agent
|
|
from langchain_core.messages import HumanMessage
|
|
from langgraph.checkpoint.memory import MemorySaver
|
|
from tools import search_knowledge_base, add_to_knowledge_base
|
|
|
|
# LLM via Ollama (llama3)
|
|
llm = ChatOpenAI(
|
|
model="ollama/llama3",
|
|
base_url="http://localhost:11434/v1",
|
|
api_key=None,
|
|
temperature=0.2,
|
|
)
|
|
|
|
agent = create_agent(
|
|
llm=llm,
|
|
tools=[search_knowledge_base, add_to_knowledge_base],
|
|
system_prompt="You are a helpful assistant that can search and add documents to the knowledge base.",
|
|
)
|
|
|
|
memory = MemorySaver()
|
|
|
|
async def run_agent(messages: List[HumanMessage]):
|
|
result = await agent.ainvoke({"messages": messages}, {"configurable": {"thread_id": "session-1"}})
|
|
return result["messages"][-1].content
|