diff --git a/main.py b/main.py index 289b4d7..087f7ef 100644 --- a/main.py +++ b/main.py @@ -1,135 +1,41 @@ -import os import asyncio -from pathlib import Path -from langchain_openai import ChatOpenAI, OpenAIEmbeddings -from langchain_core.documents import Document -from langchain_text_splitters import RecursiveCharacterTextSplitter -from langchain_qdrant import QdrantVectorStore -from langchain.tools import tool -from deepagents import create_deep_agent -from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend from langchain_core.messages import HumanMessage +from rag_agent import agent, add_to_knowledge_base, search_knowledge_base -# Configuration -OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") -if not OPENAI_API_KEY: - raise RuntimeError("OPENAI_API_KEY not set in environment") - -# LLM via OpenRouter -llm = ChatOpenAI( - model="openai/gpt-oss-20b:free", - base_url="https://openrouter.ai/api/v1", - api_key=OPENAI_API_KEY, - temperature=0.0, -) - -# Embeddings via OpenRouter -embeddings = OpenAIEmbeddings( - model="text-embedding-3-small", - base_url="https://openrouter.ai/api/v1", - api_key=OPENAI_API_KEY, -) - -# Qdrant client and vector store -from qdrant_client import QdrantClient -qdrant_client = QdrantClient(url="http://localhost:6333") -collection_name = "knowledge_base" -vector_store = QdrantVectorStore( - client=qdrant_client, - collection_name=collection_name, - embedding_function=embeddings, -) - -# Text splitter -splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) - -# Tool: search knowledge base -@tool -def search_knowledge_base(query: str, max_results: int = 3) -> str: - """Search the knowledge base for relevant information.""" - docs = vector_store.similarity_search(query, k=max_results) - if not docs: - return "No results found." - return "\n\n".join(f"Title: {doc.metadata.get('title', 'unknown')}\n{doc.page_content}" for doc in docs) - -# Tool: add to knowledge base -@tool -def add_to_knowledge_base(content: str, title: str = "document") -> str: - """Add content to the knowledge base.""" - chunks = 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 from '{title}'." - -# Backend for deepagents -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 knowledge assistant. Use the provided tools to search and add information.", -) - -# Helper: load documents from a directory -def load_documents_from_dir(directory: str): - dir_path = Path(directory) - if not dir_path.is_dir(): - raise ValueError(f"Directory {directory} does not exist.") - for file_path in dir_path.rglob("*"): - if file_path.is_file() and file_path.suffix.lower() in {".txt", ".md", ".py", ".json"}: - content = file_path.read_text(encoding="utf-8") - title = file_path.stem - add_to_knowledge_base(content, title) - -# Interactive CLI async def interactive_loop(): - print("RAG Agent CLI. Commands: /add , /search , /quit") + print("Welcome to the RAG agent. Type /add to add a document, /search to query, /quit to exit.") + thread_id = "session-1" while True: user_input = input(">> ").strip() - if not user_input: - continue if user_input.lower() == "/quit": - print("Exiting.") + print("Goodbye!") break - if user_input.startswith("/add "): - _, file_path = user_input.split(maxsplit=1) - try: - content = Path(file_path).read_text(encoding="utf-8") - title = Path(file_path).stem - result = add_to_knowledge_base(content, title) - print(result) - except Exception as e: - print(f"Error adding file: {e}") - continue - if user_input.startswith("/search "): - _, query = user_input.split(maxsplit=1) - result = search_knowledge_base(query, max_results=3) + elif user_input.lower() == "/add": + title = input("Title: ").strip() + print("Enter content (end with a single line containing only 'END'):") + lines = [] + while True: + line = input() + if line.strip() == "END": + break + lines.append(line) + content = "\n".join(lines) + result = add_to_knowledge_base(content, title) print(result) - continue - # Treat as normal message to agent - try: + elif user_input.lower() == "/search": + query = input("Query: ").strip() + result = search_knowledge_base(query) + print("Search results:") + print(result) + else: + messages = [HumanMessage(content=user_input)] response = await agent.ainvoke( - {"messages": [HumanMessage(content=user_input)]}, - {"configurable": {"thread_id": "session-1"}}, + {"messages": messages}, + {"configurable": {"thread_id": thread_id}}, ) print(response["messages"][-1].content) - except Exception as e: - print(f"Agent error: {e}") def main(): - # Optional: load initial docs from a folder - init_dir = os.getenv("INIT_DOCS_DIR") - if init_dir: - try: - load_documents_from_dir(init_dir) - print(f"Loaded documents from {init_dir}") - except Exception as e: - print(f"Failed to load initial docs: {e}") asyncio.run(interactive_loop()) if __name__ == "__main__":