126 lines
4.4 KiB
Python
126 lines
4.4 KiB
Python
from pathlib import Path
|
||
|
||
# ---------- LLM & Embeddings ----------
|
||
from langchain_ollama import Ollama, OllamaEmbeddings
|
||
|
||
llm = Ollama(
|
||
model="openai/gpt-oss-20b", # e.g. "llama3"
|
||
base_url="http://localhost:11434",
|
||
temperature=0.7,
|
||
)
|
||
|
||
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
||
|
||
# ---------- Qdrant Vector Store ----------
|
||
from qdrant_client import QdrantClient
|
||
from qdrant_client.http.models import Distance, VectorParams
|
||
from langchain_qdrant import QdrantVectorStore
|
||
|
||
client = QdrantClient(":memory:") # in‑memory for demo; replace with path or URL as needed
|
||
# Determine vector size from the embedding model
|
||
sample_vector = embeddings.embed_query("test")[0]
|
||
vector_size = len(sample_vector)
|
||
|
||
client.create_collection(
|
||
collection_name="knowledge_base",
|
||
vectors_config=VectorParams(size=vector_size, distance=Distance.COSINE),
|
||
)
|
||
|
||
vector_store = QdrantVectorStore(
|
||
client=client,
|
||
collection_name="knowledge_base",
|
||
embedding=embeddings,
|
||
)
|
||
|
||
# ---------- Text Splitter ----------
|
||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||
|
||
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
|
||
|
||
# ---------- Tools ----------
|
||
from langchain.tools import tool
|
||
from langchain_core.documents import Document
|
||
|
||
@tool("search_knowledge_base")
|
||
def search_knowledge_base(query: str, max_results: int = 5) -> str:
|
||
"""Search the knowledge base for relevant documents."""
|
||
docs = vector_store.similarity_search_with_score(query, k=max_results)
|
||
if not docs:
|
||
return "No results found."
|
||
response_lines = []
|
||
for i, (doc, score) in enumerate(docs, start=1):
|
||
title = doc.metadata.get("title", "Untitled")
|
||
snippet = doc.page_content[:200] + ("..." if len(doc.page_content) > 200 else "")
|
||
response_lines.append(f"{i}. [{score:.2f}] {title}\n{snippet}")
|
||
return "\n\n".join(response_lines)
|
||
|
||
@tool("add_to_knowledge_base")
|
||
def add_to_knowledge_base(content: str, title: str = "Untitled") -> str:
|
||
"""Add a new document to the knowledge base."""
|
||
chunks = splitter.split_text(content)
|
||
docs = [Document(page_content=c, metadata={"title": title}) for c in chunks]
|
||
vector_store.add_documents(docs)
|
||
return f"Added {len(chunks)} chunks under title '{title}'."
|
||
|
||
# ---------- Agent ----------
|
||
from langchain.agents import create_agent
|
||
from langchain_core.messages import HumanMessage
|
||
|
||
agent = create_agent(
|
||
model=llm,
|
||
tools=[search_knowledge_base, add_to_knowledge_base],
|
||
system_message="You are a helpful assistant that can search and update the knowledge base.",
|
||
)
|
||
|
||
# ---------- Document Loader ----------
|
||
def load_documents_from_dir(directory: str):
|
||
"""Load all .txt files from directory into vector store."""
|
||
for file_path in Path(directory).glob("*.txt"):
|
||
text = file_path.read_text(encoding="utf-8")
|
||
add_to_knowledge_base(text, title=file_path.stem)
|
||
|
||
# ---------- CLI ----------
|
||
def main():
|
||
print(
|
||
"Welcome to the RAG Agent. Commands:\n"
|
||
"/add <title> <file>\n"
|
||
"/search <query>\n"
|
||
"/quit\n"
|
||
)
|
||
while True:
|
||
user_input = input("> ").strip()
|
||
if not user_input:
|
||
continue
|
||
if user_input.lower() in ("/quit", "exit"):
|
||
break
|
||
|
||
if user_input.startswith("/add"):
|
||
try:
|
||
_, title, file_path = user_input.split(maxsplit=2)
|
||
content = Path(file_path).read_text(encoding="utf-8")
|
||
print(add_to_knowledge_base(content, title))
|
||
except Exception as e:
|
||
print(f"Error adding document: {e}")
|
||
|
||
elif user_input.startswith("/search"):
|
||
query = user_input[len("/search") :].strip()
|
||
if not query:
|
||
print("Please provide a search query.")
|
||
continue
|
||
result = agent.invoke({"messages": [HumanMessage(content=query)]})
|
||
for msg in result["messages"]:
|
||
if hasattr(msg, "content"):
|
||
print(msg.content)
|
||
|
||
else:
|
||
# Treat as normal chat message
|
||
result = agent.invoke({"messages": [HumanMessage(content=user_input)]})
|
||
for msg in result["messages"]:
|
||
if hasattr(msg, "content"):
|
||
print(msg.content)
|
||
|
||
if __name__ == "__main__":
|
||
# Optional: load initial docs from a folder named 'docs'
|
||
if Path("docs").exists():
|
||
load_documents_from_dir("docs")
|
||
main() |