From bf14e4ed93b590303814682df1e3b996f6dea4a4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9A=D0=B8=D1=80=D0=B8=D0=BB=D0=BB=20=D0=9A=D1=83=D1=82?= =?UTF-8?q?=D0=BB=D0=B0=D1=85=D0=BC=D0=B5=D1=82=D0=BE=D0=B2?= Date: Thu, 28 May 2026 12:38:43 +0000 Subject: [PATCH] add rag_agent.py --- rag_agent.py | 99 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 99 insertions(+) create mode 100644 rag_agent.py diff --git a/rag_agent.py b/rag_agent.py new file mode 100644 index 0000000..89e6b2a --- /dev/null +++ b/rag_agent.py @@ -0,0 +1,99 @@ +""" +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"]