diff --git a/agent.py b/agent.py index 9b2f790..7900644 100644 --- a/agent.py +++ b/agent.py @@ -1,32 +1,28 @@ -""" -Create a LangChain agent that chooses between local KB and web search. -""" - -from typing import List, Dict - from langchain.agents import create_agent -from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder -from langchain.schema.document import Document +from langchain_ollama import ChatOllama -# Import tools from tools import search_local_kb, web_search +OLLAMA_BASE_URL = "http://127.0.0.1:11434" +LLM_MODEL = "llama3" + SYSTEM_PROMPT = """ -You are an assistant that can answer questions using either a local knowledge base or the web. Use the tool `search_local_kb` when the question is about content already in your documents. Use `web_search` for up‑to‑date facts. -When you provide an answer, include the source: either "chromadb" or "tavily". +You are a RAG assistant that must choose the best source before answering. + +Rules: +- Use search_local_kb for questions about local notes, internal documents, course materials, or any topic that may already exist in the local knowledge base. +- Use web_search for current events, recent news, live facts, or questions that clearly require the internet. +- Do not call both tools unless the first one clearly failed to provide enough information. +- Every final answer must end with a separate line in the exact format: "Источник: chromadb" or "Источник: tavily". +- Answer in Russian. """ -# Prompt template with tool messages -prompt = ChatPromptTemplate.from_messages([ - ("system", SYSTEM_PROMPT), - MessagesPlaceholder("history"), - ("human", "{input}"), -]) - -# Create agent harness agent = create_agent( - model="ollama:llama3", + model=ChatOllama( + model=LLM_MODEL, + base_url=OLLAMA_BASE_URL, + temperature=0, + ), tools=[search_local_kb, web_search], system_prompt=SYSTEM_PROMPT, - prompt_template=prompt, )