116 lines
4.2 KiB
Python
116 lines
4.2 KiB
Python
import os
|
|
from functools import lru_cache
|
|
from typing import Any
|
|
from uuid import uuid4
|
|
|
|
from langchain.agents import create_agent
|
|
from langchain_core.documents import Document
|
|
from langchain_core.tools import tool
|
|
from langchain_ollama import ChatOllama, OllamaEmbeddings
|
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
from chromadb import Client
|
|
|
|
OLLAMA_CHAT_MODEL = os.getenv("OLLAMA_CHAT_MODEL", "llama3")
|
|
OLLAMA_EMBED_MODEL = os.getenv("OLLAMA_EMBED_MODEL", "nomic-embed-text")
|
|
EMBEDDING_SIZE = int(os.getenv("OLLAMA_EMBEDDING_SIZE", "768"))
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def get_embeddings() -> OllamaEmbeddings:
|
|
return OllamaEmbeddings(model=OLLAMA_EMBED_MODEL)
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def get_vector_store() -> Client:
|
|
client = Client()
|
|
# Ensure collection exists
|
|
client.get_or_create_collection(name="rag_memory")
|
|
return client
|
|
|
|
|
|
def chunk_document(content: str, title: str) -> list[Document]:
|
|
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)
|
|
return splitter.create_documents([content], metadatas=[{"title": title}])
|
|
|
|
|
|
@tool
|
|
def search_knowledge_base(query: str, max_results: int = 3) -> str:
|
|
"""Search relevant chunks in the local ChromaDB knowledge base."""
|
|
client = get_vector_store()
|
|
collection = client.get_collection(name="rag_memory")
|
|
results = collection.query(query_texts=[query], n_results=max_results)
|
|
if not results.get("documents"):
|
|
return "No relevant documents found."
|
|
|
|
lines: list[str] = []
|
|
for idx, (doc, dist) in enumerate(zip(results["documents"], results["distances"]), start=1):
|
|
ids = results["ids"]
|
|
metadata = collection.get(ids=[ids[idx-1]])["metadatas"][0]
|
|
title = metadata.get("title", "untitled") if metadata else "untitled"
|
|
snippet = doc.replace("\n", " ")[:300]
|
|
lines.append(f"{idx}. {title} (score={dist:.4f}): {snippet}")
|
|
return "\n".join(lines)
|
|
|
|
|
|
@tool
|
|
def add_to_knowledge_base(content: str, title: str) -> str:
|
|
"""Split content into chunks and store it in the local ChromaDB knowledge base."""
|
|
documents = chunk_document(content, title)
|
|
ids = [str(uuid4()) for _ in documents]
|
|
embeddings = get_embeddings().embed_documents([doc.page_content for doc in documents])
|
|
client = get_vector_store()
|
|
collection = client.get_collection(name="rag_memory")
|
|
collection.add(ids=ids, documents=[doc.page_content for doc in documents], embeddings=embeddings, metadatas=[doc.metadata for doc in documents])
|
|
return f"Added {len(documents)} chunk(s) from '{title}' to the knowledge base."
|
|
|
|
|
|
def build_agent() -> Any:
|
|
llm = ChatOllama(model=OLLAMA_CHAT_MODEL)
|
|
return create_agent(
|
|
model=llm,
|
|
tools=[search_knowledge_base, add_to_knowledge_base],
|
|
system_prompt=(
|
|
"You are a RAG assistant. Use search_knowledge_base before answering "
|
|
"questions that may depend on stored knowledge, and use "
|
|
"add_to_knowledge_base when the user asks to remember information."
|
|
),
|
|
)
|
|
|
|
|
|
def run_cli() -> None:
|
|
print("RAG agent. Commands: /add <title> | <content>, /search <query>, /quit")
|
|
agent = None
|
|
|
|
while True:
|
|
try:
|
|
user_input = input(">>> ").strip()
|
|
except EOFError:
|
|
break
|
|
|
|
if not user_input:
|
|
continue
|
|
if user_input == "/quit":
|
|
break
|
|
if user_input.startswith("/add "):
|
|
payload = user_input[5:]
|
|
if "|" in payload:
|
|
title, content = [part.strip() for part in payload.split("|", 1)]
|
|
else:
|
|
title, content = "note", payload.strip()
|
|
print(add_to_knowledge_base.invoke({"content": content, "title": title}))
|
|
continue
|
|
if user_input.startswith("/search "):
|
|
query = user_input[8:].strip()
|
|
print(search_knowledge_base.invoke({"query": query, "max_results": 3}))
|
|
continue
|
|
|
|
if agent is None:
|
|
agent = build_agent()
|
|
response = agent.invoke({"messages": [{"role": "user", "content": user_input}]})
|
|
messages = response.get("messages", [])
|
|
print(messages[-1].content if messages else response)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
run_cli()
|