80 lines
2.7 KiB
Python
80 lines
2.7 KiB
Python
"""
|
|
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"}})
|