141 lines
4.5 KiB
Python
141 lines
4.5 KiB
Python
"""
|
||
Agent and vector store setup for the RAG task.
|
||
|
||
This module defines:
|
||
* `QdrantVectorStore` wrapper that uses Ollama embeddings.
|
||
* Two tools – ``search_knowledge_base`` and ``add_to_knowledge_base``.
|
||
* A helper to create the agent via :func:`langchain.agents.create_agent`.
|
||
"""
|
||
|
||
import os
|
||
from pathlib import Path
|
||
from typing import List, Dict
|
||
|
||
from langchain_ollama import ChatOllama, OllamaEmbeddings
|
||
from langchain_qdrant import QdrantVectorStore
|
||
from langchain.tools import tool
|
||
from langchain.agents import create_agent
|
||
from langchain_core.messages import HumanMessage
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Vector store configuration
|
||
# ---------------------------------------------------------------------------
|
||
|
||
QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost")
|
||
QDRANT_PORT = int(os.getenv("QDRANT_PORT", "6333"))
|
||
COLLECTION_NAME = "knowledge"
|
||
|
||
# Initialize embeddings and vector store. The client is created lazily on first use.
|
||
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
||
vector_store: QdrantVectorStore | None = None
|
||
|
||
|
||
def get_vector_store() -> QdrantVectorStore:
|
||
"""Return a singleton Qdrant vector store instance.
|
||
|
||
The collection is created automatically if it does not exist.
|
||
"""
|
||
global vector_store
|
||
if vector_store is None:
|
||
from qdrant_client import QdrantClient
|
||
client = QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT)
|
||
vector_store = QdrantVectorStore(
|
||
client=client,
|
||
collection_name=COLLECTION_NAME,
|
||
embedding=embeddings,
|
||
)
|
||
return vector_store
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Tools
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@tool("search_knowledge_base")
|
||
def search_knowledge_base(query: str, max_results: int = 5) -> str:
|
||
"""Semantic search in the knowledge base.
|
||
|
||
Parameters
|
||
----------
|
||
query: str
|
||
Search query.
|
||
max_results: int, optional
|
||
Number of top results to return. Defaults to 5.
|
||
|
||
Returns
|
||
-------
|
||
str
|
||
A numbered list of passages or a message if nothing was found.
|
||
"""
|
||
store = get_vector_store()
|
||
docs = store.similarity_search(query, k=max_results)
|
||
if not docs:
|
||
return "No relevant documents found."
|
||
return "\n\n".join(f"{i+1}. {doc.page_content}" for i, doc in enumerate(docs))
|
||
|
||
@tool("add_to_knowledge_base")
|
||
def add_to_knowledge_base(content: str, title: str = "document") -> str:
|
||
"""Add a document to the knowledge base.
|
||
|
||
Parameters
|
||
----------
|
||
content: str
|
||
Text of the document.
|
||
title: str, optional
|
||
Title used as metadata. Defaults to ``"document"``.
|
||
|
||
Returns
|
||
-------
|
||
str
|
||
Confirmation message.
|
||
"""
|
||
store = get_vector_store()
|
||
from langchain_core.documents import Document
|
||
doc = Document(page_content=content, metadata={"title": title})
|
||
store.add_documents([doc])
|
||
return f"Added '{title}' to knowledge base."
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Agent creation helper
|
||
# ---------------------------------------------------------------------------
|
||
|
||
llm = ChatOllama(model="llama3", temperature=0.0)
|
||
|
||
SYSTEM_PROMPT = (
|
||
"You are an assistant that can search and add information to a local knowledge base.\n"
|
||
"Use the tools `search_knowledge_base` and `add_to_knowledge_base`.\n"
|
||
"When searching, return the most relevant passages. When adding, confirm success."
|
||
)
|
||
|
||
|
||
def create_rag_agent():
|
||
"""Return a LangChain agent configured with the RAG tools."""
|
||
agent = create_agent(
|
||
llm=llm,
|
||
tools=[search_knowledge_base, add_to_knowledge_base],
|
||
system_prompt=SYSTEM_PROMPT,
|
||
)
|
||
return agent
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Example usage (can be imported by main.py)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
if __name__ == "__main__":
|
||
# Simple demo: add a short doc and search it.
|
||
agent = create_rag_agent()
|
||
print("Adding sample document...")
|
||
res = agent.invoke(
|
||
{"messages": [HumanMessage(content="Add to knowledge base: content='Python is great' title='Python intro'")]},
|
||
{"configurable": {"thread_id": "demo-1"}},
|
||
)
|
||
print(res["messages"][-1].content)
|
||
|
||
print("Searching for Python...")
|
||
res = agent.invoke(
|
||
{"messages": [HumanMessage(content="Search for Python')"],},
|
||
{"configurable": {"thread_id": "demo-2"}},
|
||
)
|
||
print(res["messages"][-1].content)
|
||
|
||
# End of agent.py
|