diff --git a/agent.py b/agent.py deleted file mode 100644 index 758022a..0000000 --- a/agent.py +++ /dev/null @@ -1,79 +0,0 @@ -""" -Agent definition for the RAG system. - -Provides two tools: -* search_knowledge_base(query, max_results) -* add_to_knowledge_base(content, title) - -The agent is created with create_agent from langchain.agents. -""" - -import os -from typing import List -from langchain_ollama import OllamaEmbeddings -from chromadb import PersistentClient -from chromadb.utils import embedding_functions as ef -from langchain_text_splitters import RecursiveCharacterTextSplitter -from langchain.tools import tool -from langchain.agents import create_agent -from langchain_core.messages import HumanMessage - -# Initialize embeddings and Chroma client -EMBEDDINGS = OllamaEmbeddings(model="nomic-embed-text") -CHROMA_PATH = os.path.join(os.getcwd(), "chromadb_store") -CLIENT = PersistentClient(path=CHROMA_PATH) -COLLECTION_NAME = "knowledge" -if COLLECTION_NAME not in CLIENT.list_collections(): - CLIENT.create_collection(name=COLLECTION_NAME, embedding_function=EMBEDDINGS) -COLL = CLIENT.get_or_create_collection(name=COLLECTION_NAME, embedding_function=EMBEDDINGS) - -# Chunker from chunker.py -from chunker import CHUNKER - -@tool -def add_to_knowledge_base(content: str, title: str) -> str: - """ - Add a document to the knowledge base. - The content is split into chunks and stored with metadata. - Returns confirmation message. - """ - # Split content - chunks = CHUNKER.split_text(content) - ids = [f"{title}_{i}" for i in range(len(chunks))] - metadatas = [{"title": title} for _ in chunks] - COLL.add(ids=ids, documents=chunks, metadatas=metadatas) - return f"Added {len(chunks)} chunks from '{title}'." - -@tool -def search_knowledge_base(query: str, max_results: int = 5) -> str: - """ - Search the knowledge base for relevant documents. - Returns a formatted string of results. - """ - results = COLL.query( - query_texts=[query], - n_results=max_results, - include=['documents', 'distances'], - ) - docs = results.get("documents", [])[0] - dists = results.get("distances", [])[0] - if not docs: - return "No relevant documents found." - output_lines = [] - for i, (doc, dist) in enumerate(zip(docs, dists), 1): - output_lines.append(f"{i}. (score: {dist:.4f})\n{doc[:200]}...") - return "\n\n".join(output_lines) - -# Create agent -SYSTEM_PROMPT = ( - "You are an assistant that can search and add to a knowledge base. Use the provided tools." -) -AGENT = create_agent( - llm=None, # No LLM needed for tool calls; agent will use system prompt only - tools=[add_to_knowledge_base, search_knowledge_base], - system_prompt=SYSTEM_PROMPT, -) - -# Expose a simple invoke function -async def run_agent(messages: List[HumanMessage]): - return await AGENT.ainvoke({"messages": messages}, {"configurable": {"thread_id": "rag-agent"}})