diff --git a/agent.py b/agent.py index d51ed54..dd244a5 100644 --- a/agent.py +++ b/agent.py @@ -1,24 +1,56 @@ -from langchain import ChatPromptTemplate +from langchain_ollama import ChatOllama +from langchain.prompts import ChatPromptTemplate from langchain.agents import create_agent from langchain.tools import tool -from .tools import add_content, search_content import os -# LLM configuration – use Ollama via langchain-ollama -MODEL = os.getenv("OLLAMA_MODEL", "llama3.1") +MODEL = os.getenv("OLLAMA_MODEL", "llama3") BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434/v1") -from langchain_ollama import ChatOllama +# Initialize LLM with local Ollama endpoint. llm = ChatOllama(model=MODEL, base_url=BASE_URL) -# System prompt for the agent SYSTEM_PROMPT = """You are an assistant that answers user queries using a knowledge base. Use the provided tools to search and add content.""" prompt = ChatPromptTemplate.from_messages([ ("system", SYSTEM_PROMPT), ]) -tools = [add_content, search_content] +from .rag_tools import add_to_knowledge_base as rag_add, search_knowledge_base as rag_search +from .init_loader import load_documents + +# Load existing documents from ./data directory at startup +load_documents("./data") + +@tool +def add_to_knowledge_base(content: str, title: str = "Document") -> str: + """Add content to the knowledge base.""" + return rag_add(content=content, title=title) + +@tool +def search_knowledge_base(query: str, max_results: int = 5) -> str: + """Search the knowledge base for relevant chunks.""" + results = rag_search(query=query, max_results=max_results) + formatted = "\n".join([f"{score:.4f}: {text[:200]}..." for text, score in results]) + return formatted if formatted else "No results found." + +tools = [add_to_knowledge_base, search_knowledge_base] agent = create_agent(llm=llm, prompt=prompt, tools=tools) +executor = agent + if __name__ == "__main__": - print("Agent initialized.") \ No newline at end of file + print("RAG Agent Interactive Mode. Type /quit to exit.") + while True: + try: + user_input = input("User: ") + except (EOFError, KeyboardInterrupt): + print("\nGoodbye!") + break + if user_input.strip().lower() in {"/quit", "quit"}: + print("Goodbye!") + break + try: + response = executor.invoke({"input": user_input}) + print("Assistant:", response) + except Exception as e: + print("Error:", str(e))