diff --git a/rag_agent.py b/rag_agent.py deleted file mode 100644 index 89e6b2a..0000000 --- a/rag_agent.py +++ /dev/null @@ -1,99 +0,0 @@ -""" -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"]