add rag_agent.py

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"]