diff --git a/main.py b/main.py new file mode 100644 index 0000000..231ff25 --- /dev/null +++ b/main.py @@ -0,0 +1,112 @@ +""" +Agent with RAG memory using Qdrant and Ollama. +""" + +import os +import sys +from pathlib import Path +from typing import List, Dict, Any + +from langchain_ollama import OllamaEmbeddings, Ollama +from langchain_qdrant import QdrantVectorStore +from langchain.text_splitter import RecursiveCharacterTextSplitter +from langchain.tools import tool +from langchain.agents import create_agent, AgentExecutor, AgentType +from langchain.schema import AgentAction, AgentFinish +from langchain.callbacks import get_openai_callback + +# Configuration +QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333") +QDRANT_COLLECTION = os.getenv("QDRANT_COLLECTION", "rag_agent") +Ollama_EMBED_MODEL = os.getenv("Ollama_EMBED_MODEL", "nomic-embed-text") +Ollama_MODEL = os.getenv("Ollama_MODEL", "llama3") + +# Initialize embeddings and vector store +embeddings = OllamaEmbeddings(model=Ollama_EMBED_MODEL) +vector_store = QdrantVectorStore( + url=QDRANT_URL, + collection_name=QDRANT_COLLECTION, + embeddings=embeddings, +) + +# Text splitter +splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50) + +# Tools +@tool +def search_knowledge_base(query: str, max_results: int = 5) -> str: + """Semantic search in the knowledge base.""" + docs = vector_store.similarity_search(query, k=max_results) + if not docs: + return "No relevant documents found." + return "\n\n".join([f"{i+1}. {doc.page_content[:200]}..." for i, doc in enumerate(docs)]) + +@tool +def add_to_knowledge_base(content: str, title: str) -> str: + """Add a document to the knowledge base.""" + # Split content into chunks + chunks = splitter.split_text(content) + # Create documents with metadata + docs = [] + for i, chunk in enumerate(chunks): + docs.append( + { + "page_content": chunk, + "metadata": {"title": title, "chunk_index": i}, + } + ) + # Add to vector store + vector_store.add_documents(docs) + return f"Document '{title}' added with {len(chunks)} chunks." + +# Agent setup +SYSTEM_PROMPT = ( + "You are an AI assistant with access to a local knowledge base. " + "Use the provided tools to search and add information. " + "When answering, rely on the knowledge base and the LLM." +) + +# Create agent with tools +agent = create_agent( + llm=Ollama(model=Ollama_MODEL), + tools=[search_knowledge_base, add_to_knowledge_base], + system_prompt=SYSTEM_PROMPT, + agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION, +) + +executor = AgentExecutor(agent=agent, tools=[search_knowledge_base, add_to_knowledge_base], verbose=True) + +# CLI +def main(): + print("RAG Agent CLI. Commands: /add