79 lines
2.5 KiB
Python
79 lines
2.5 KiB
Python
"""RAG agent: Qdrant vector store + Ollama LLM + LangGraph."""
|
|
import os
|
|
from langchain_ollama import ChatOllama, OllamaEmbeddings
|
|
from langchain_qdrant import QdrantVectorStore
|
|
from langchain_core.documents import Document
|
|
from langchain.tools import tool
|
|
from langgraph.prebuilt import create_react_agent
|
|
|
|
QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost")
|
|
QDRANT_PORT = int(os.getenv("QDRANT_PORT", "6333"))
|
|
COLLECTION_NAME = "knowledge"
|
|
|
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
|
_vs = None
|
|
|
|
|
|
def get_vector_store() -> QdrantVectorStore:
|
|
"""Return singleton Qdrant vector store, creating collection if needed."""
|
|
global _vs
|
|
if _vs is None:
|
|
from qdrant_client import QdrantClient
|
|
from qdrant_client.models import Distance, VectorParams
|
|
client = QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT)
|
|
existing = [c.name for c in client.get_collections().collections]
|
|
if COLLECTION_NAME not in existing:
|
|
client.create_collection(
|
|
collection_name=COLLECTION_NAME,
|
|
vectors_config=VectorParams(size=768, distance=Distance.COSINE),
|
|
)
|
|
_vs = QdrantVectorStore(
|
|
client=client,
|
|
collection_name=COLLECTION_NAME,
|
|
embedding=embeddings,
|
|
)
|
|
return _vs
|
|
|
|
|
|
@tool
|
|
def search_knowledge_base(query: str) -> str:
|
|
"""Search the knowledge base for relevant passages.
|
|
|
|
Args:
|
|
query: search query string
|
|
"""
|
|
docs = get_vector_store().similarity_search(query, k=5)
|
|
if not docs:
|
|
return "No results found."
|
|
return "\n\n".join(f"{i+1}. {d.page_content}" for i, d in enumerate(docs))
|
|
|
|
|
|
@tool
|
|
def add_to_knowledge_base(text: str) -> str:
|
|
"""Add a text passage to the knowledge base.
|
|
|
|
Args:
|
|
text: text to store
|
|
"""
|
|
get_vector_store().add_documents([Document(page_content=text)])
|
|
return "Added to knowledge base."
|
|
|
|
llm = ChatOllama(model="llama3", temperature=0.0)
|
|
|
|
agent = create_react_agent(
|
|
model=llm,
|
|
tools=[search_knowledge_base, add_to_knowledge_base],
|
|
state_modifier=(
|
|
"You are a helpful assistant with access to a local knowledge base. "
|
|
"Use search_knowledge_base to find information. "
|
|
"Use add_to_knowledge_base to store new information."
|
|
),
|
|
)
|
|
|
|
|
|
def run_agent(user_input: str) -> str:
|
|
"""Run the agent and return the last message content."""
|
|
from langchain_core.messages import HumanMessage
|
|
result = agent.invoke({"messages": [HumanMessage(content=user_input)]})
|
|
return result["messages"][-1].content
|