diff --git a/main.py b/main.py index 13d7632..fc0eafe 100644 --- a/main.py +++ b/main.py @@ -1,36 +1,16 @@ -""" -# main.py – RAG‑agent with Qdrant, OpenRouter, and deepagents -# ---------------------------------------------------------- -# 1. Imports and configuration -# 2. Qdrant vector store wrapper (embedding, add, search) -# 3. Text splitter (RecursiveCharacterTextSplitter) -# 4. LangChain tools: search_knowledge_base, add_to_knowledge_base -# 5. DeepAgent creation (create_deep_agent) -# 6. CLI client for /add, /search, /quit -# ---------------------------------------------------------- -""" import os import asyncio -import json from pathlib import Path -from typing import List - -from langchain_openai import ChatOpenAI, OpenAIEmbeddings +from langchain_openai import ChatOpenAI +from langchain_ollama import OllamaEmbeddings +from langchain_qdrant import QdrantVectorStore from langchain_core.documents import Document -from langchain_text_splitters import RecursiveCharacterTextSplitter from langchain.tools import tool from deepagents import create_deep_agent from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend -from langchain_qdrant import QdrantVectorStore +from langchain_core.messages import HumanMessage -# ------------------------------------------------------------------ -# 1. Configuration -# ------------------------------------------------------------------ -# Load environment variables (e.g. OPENAI_API_KEY) -from dotenv import load_dotenv -load_dotenv() - -# LLM – OpenRouter (free tier) +# ---------- LLM ---------- llm = ChatOpenAI( model="openai/gpt-oss-20b:free", base_url="https://openrouter.ai/api/v1", @@ -38,110 +18,91 @@ llm = ChatOpenAI( temperature=0.0, ) -# Embeddings – OpenAI via OpenRouter -embeddings = OpenAIEmbeddings( - model="text-embedding-3-small", - base_url="https://openrouter.ai/api/v1", - api_key=os.getenv("OPENAI_API_KEY"), -) +# ---------- Embeddings ---------- +# Using Ollama embeddings as per assignment correction +embeddings = OllamaEmbeddings(model="nomic-embed-text") -# Qdrant client – assumes Qdrant is running locally on default port -qdrant_url = os.getenv("QDRANT_URL", "http://localhost:6333") -collection_name = "knowledge_base" +# ---------- Vector Store (Qdrant) ---------- +# Ensure Qdrant is running locally (default port 6333) vector_store = QdrantVectorStore( - url=qdrant_url, - collection_name=collection_name, - embeddings=embeddings, + url="http://localhost:6333", + collection_name="knowledge", + embedding_function=embeddings, ) -# Text splitter – 1000 chars max, 200 overlap -text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) - -# ------------------------------------------------------------------ -# 2. Tools -# ------------------------------------------------------------------ +# ---------- Tools ---------- @tool def search_knowledge_base(query: str, max_results: int = 3) -> str: - """Semantic search in the Qdrant knowledge base.""" - docs: List[Document] = vector_store.similarity_search(query, k=max_results) + """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---\n\n".join([f"{doc.metadata.get('title', 'Untitled')}:\n{doc.page_content}" for doc in docs]) + return "No results 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 = "Untitled") -> str: - """Add a new document to the knowledge base. - The content is split into chunks before being stored. - """ - chunks = text_splitter.split_text(content) - docs = [Document(page_content=chunk, metadata={"title": title}) for chunk in chunks] - vector_store.add_documents(docs) - return f"Added {len(docs)} chunks for document '{title}'." +def add_to_knowledge_base(content: str, title: str = "untitled") -> str: + """Add a document to the knowledge base.""" + doc = Document(page_content=content, metadata={"title": title}) + vector_store.add_documents([doc]) + return f"Document '{title}' added to the knowledge base." -# ------------------------------------------------------------------ -# 3. DeepAgent setup -# ------------------------------------------------------------------ +# ---------- Backend ---------- backend = CompositeBackend([ LocalShellBackend(workspace_dir="./workspace"), FilesystemBackend(), ]) +# ---------- Agent ---------- agent = create_deep_agent( model=llm, tools=[search_knowledge_base, add_to_knowledge_base], backend=backend, - system_prompt="You are a helpful assistant with access to a knowledge base. Use the provided tools to search and add information.", + system_prompt="You are an assistant with access to a knowledge base. Use the provided tools to search and add information." ) -# ------------------------------------------------------------------ -# 4. CLI client -# ------------------------------------------------------------------ -async def run_cli(): - print("Welcome to the RAG Agent CLI. Commands: /add <file>, /search <query>, /quit") - thread_id = "cli-session" +# ---------- Document Loader ---------- +async def load_documents_from_dir(directory: str): + """Load all text files from a directory into the vector store.""" + for file_path in Path(directory).rglob("*.txt"): + text = file_path.read_text(encoding="utf-8") + title = file_path.stem + await agent.ainvoke( + {"messages": [HumanMessage(content=f"/add {title}")], "content": text}, + {"configurable": {"thread_id": "init"}}, + ) + +# ---------- Interactive CLI ---------- +async def interactive_loop(): + print("Welcome to the RAG Agent. Commands: /add <title>, /search <query>, /quit") while True: - try: - user_input = input("> ") - except EOFError: - break - if not user_input: - continue - if user_input.startswith("/quit"): + user_input = input("> ") + if user_input.strip() == "/quit": print("Goodbye!") break - if user_input.startswith("/add"): - parts = user_input.split(maxsplit=2) - if len(parts) < 3: - print("Usage: /add <title> <file_path>") - continue - title, file_path = parts[1], parts[2] - try: - content = Path(file_path).read_text(encoding="utf-8") - except Exception as e: - print(f"Error reading file: {e}") - continue - # Invoke tool directly - result = add_to_knowledge_base(content, title) - print(result) - continue - if user_input.startswith("/search"): - query = user_input[len("/search"):].strip() - if not query: - print("Usage: /search <query>") - continue - # Use agent to perform search via tool + if user_input.startswith("/add "): + parts = user_input.split(" ", 1) + title = parts[1] if len(parts) > 1 else "untitled" + content = input("Enter content: ") response = await agent.ainvoke( - {"messages": [{"role": "user", "content": f"search {query}"}]}, - {"configurable": {"thread_id": thread_id}}, + {"messages": [HumanMessage(content=f"/add {title}")], "content": content}, + {"configurable": {"thread_id": "session"}}, ) print(response["messages"][-1].content) - continue - # Default: treat as normal user message - response = await agent.ainvoke( - {"messages": [{"role": "user", "content": user_input}]}, - {"configurable": {"thread_id": thread_id}}, - ) - print(response["messages"][-1].content) + elif user_input.startswith("/search "): + query = user_input.split(" ", 1)[1] + response = await agent.ainvoke( + {"messages": [HumanMessage(content=f"/search {query}")]}, + {"configurable": {"thread_id": "session"}}, + ) + print(response["messages"][-1].content) + else: + print("Unknown command. Use /add, /search, or /quit.") + +# ---------- Main ---------- +async def main(): + # Optional: load initial documents + # await load_documents_from_dir("./data") + await interactive_loop() if __name__ == "__main__": - asyncio.run(run_cli()) + asyncio.run(main())