100 lines
3.0 KiB
Python
100 lines
3.0 KiB
Python
"""
|
||
RAG agent implementation with Qdrant and Ollama.
|
||
"""
|
||
|
||
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_core.messages import HumanMessage
|
||
from langchain.agents import create_agent
|
||
|
||
# Initialize LLM and embeddings using Ollama
|
||
LLM_MODEL = "llama3"
|
||
EMBEDDING_MODEL = "nomic-embed-text"
|
||
|
||
llm = ChatOllama(model=LLM_MODEL, temperature=0.0)
|
||
embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
|
||
|
||
# Qdrant client (in‑memory for simplicity)
|
||
from qdrant_client import QdrantClient
|
||
from qdrant_client.models import Distance, VectorParams
|
||
|
||
client = QdrantClient(":memory:")
|
||
COLLECTION_NAME = "knowledge"
|
||
if not client.collection_exists(COLLECTION_NAME):
|
||
client.create_collection(
|
||
COLLECTION_NAME,
|
||
vectors_config=VectorParams(size=embeddings.embedding_size, distance=Distance.COSINE),
|
||
)
|
||
vector_store = QdrantVectorStore(client=client, collection_name=COLLECTION_NAME, embedding=embeddings)
|
||
|
||
# Tool: add to knowledge base
|
||
@tool
|
||
def add_to_knowledge_base(content: str, title: str = "document") -> str:
|
||
"""Add a document to the vector store.
|
||
|
||
Parameters
|
||
----------
|
||
content: str
|
||
Raw text of the document.
|
||
title: str, optional
|
||
Title or identifier for the document.
|
||
Returns
|
||
-------
|
||
str
|
||
Confirmation message.
|
||
"""
|
||
# Split into chunks using chunker module
|
||
from chunker import split_text
|
||
chunks = split_text(content)
|
||
docs = []
|
||
for i, chunk in enumerate(chunks):
|
||
meta = {"title": title, "chunk_index": str(i)}
|
||
docs.append({"page_content": chunk, "metadata": meta})
|
||
vector_store.add_documents(docs)
|
||
return f"Added {len(chunks)} chunks from '{title}'."
|
||
|
||
# Tool: search knowledge base
|
||
@tool
|
||
def search_knowledge_base(query: str, max_results: int = 5) -> str:
|
||
"""Semantic search in the vector store.
|
||
|
||
Parameters
|
||
----------
|
||
query: str
|
||
Search query.
|
||
max_results: int, optional
|
||
Number of top results to return.
|
||
Returns
|
||
-------
|
||
str
|
||
Formatted search results.
|
||
"""
|
||
docs = vector_store.similarity_search(query, k=max_results)
|
||
if not docs:
|
||
return "No relevant documents found."
|
||
lines = []
|
||
for i, doc in enumerate(docs, 1):
|
||
title = doc.metadata.get("title", "unknown")
|
||
chunk_idx = doc.metadata.get("chunk_index", "0")
|
||
lines.append(f"{i}. [{title} - chunk {chunk_idx}]\n{doc.page_content[:200]}...")
|
||
return "\n\n".join(lines)
|
||
|
||
# Create agent with tools
|
||
SYSTEM_PROMPT = (
|
||
"You are an assistant that can search and add documents to a knowledge base."
|
||
" Use the provided tools to manage the knowledge base."
|
||
)
|
||
agent = create_agent(
|
||
llm=llm,
|
||
tools=[add_to_knowledge_base, search_knowledge_base],
|
||
system_prompt=SYSTEM_PROMPT,
|
||
)
|
||
|
||
# Expose agent for external use
|
||
__all__ = ["agent", "add_to_knowledge_base", "search_knowledge_base"]
|