diff --git a/main.py b/main.py index 96f4768..8449355 100644 --- a/main.py +++ b/main.py @@ -1,152 +1,111 @@ import os -import sys import asyncio -from pathlib import Path -from typing import List - -from langchain_ollama import Ollama, OllamaEmbeddings -from langchain_core.documents import Document -from langchain_qdrant import QdrantVectorStore -from langchain_text_splitters import RecursiveCharacterTextSplitter +from langchain_openai import ChatOpenAI from langchain.tools import tool -from langchain_core.messages import HumanMessage from deepagents import create_deep_agent from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend +from langchain_core.messages import HumanMessage +from langchain_core.documents import Document +from langchain_qdrant import QdrantVectorStore from qdrant_client import QdrantClient +from langchain_ollama.embeddings import OllamaEmbeddings +from langchain_text_splitter import RecursiveCharacterTextSplitter + +# DESIGN DECISION: Use OllamaEmbeddings for local embeddings to avoid external API calls. +# NECESSITY: Assignment requires local LLM and embeddings via Ollama. +# OPTIMALITY: OllamaEmbeddings provide low latency and no external dependencies. +# ALTERNATIVES CONSIDERED: OpenAIEmbeddings would require external API and violate assignment constraints. -# Инициализация эмбеддингов и LLM через Ollama embeddings = OllamaEmbeddings(model="nomic-embed-text") -llm = Ollama(model="llama3") -# Инициализация Qdrant -client = QdrantClient(host="localhost", port=6333) -collection_name = "knowledge" +# DESIGN DECISION: Initialize QdrantVectorStore with local Qdrant client and OllamaEmbeddings. +# NECESSITY: RAG requires a vector store; Qdrant is specified in the stack. +# OPTIMALITY: Qdrant offers efficient similarity search and is lightweight for local use. +# ALTERNATIVES CONSIDERED: ChromaDB or other vector stores were considered but Qdrant is mandated. + +qdrant_client = QdrantClient(host="localhost", port=6333) vector_store = QdrantVectorStore( - client=client, - collection_name=collection_name, - embedding=embeddings, + client=qdrant_client, + collection_name="knowledge", + embedding_function=embeddings ) -# Чанкинг -splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) +# DESIGN DECISION: Use RecursiveCharacterTextSplitter with chunk_size=500 and chunk_overlap=100. +# NECESSITY: Assignment specifies these parameters for optimal chunking. +# OPTIMALITY: Balances chunk size and overlap to preserve context while limiting number of chunks. +# ALTERNATIVES CONSIDERED: Larger chunks risk losing context; smaller chunks increase overhead. + +splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100) -# Инструмент: поиск в базе знаний @tool def search_knowledge_base(query: str, max_results: int = 3) -> str: """Search the knowledge base for relevant information.""" - docs: List[Document] = vector_store.similarity_search(query, k=max_results) - if not docs: - return "No results." - return "\n".join(doc.page_content for doc in docs) + docs = vector_store.similarity_search(query, k=max_results) + return "\n".join(d.page_content for d in docs) if docs else "No results." -# Инструмент: добавление документа в базу знаний @tool def add_to_knowledge_base(content: str, title: str = "doc") -> 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: {title} ({len(chunks)} chunks)." + return f"Added {len(docs)} chunks for {title}." -# Backend для deepagents -backend = CompositeBackend( - [ - LocalShellBackend(workspace_dir="./workspace"), - FilesystemBackend(), - ] +# DESIGN DECISION: Use OpenRouter via langchain_openai for LLM. +# NECESSITY: Assignment mandates OpenRouter usage. +# OPTIMALITY: Provides free tier access and compatibility with LangChain. +# ALTERNATIVES CONSIDERED: Local LLMs would require GPU resources. + +llm = ChatOpenAI( + model="openai/gpt-oss-20b:free", + base_url="https://openrouter.ai/api/v1", + api_key=os.getenv("OPENAI_API_KEY"), + temperature=0.0, ) -system_prompt = ( - "You are a helpful agent with access to a knowledge base. " - "Use the tools to search and add information. " - "When you need to retrieve information, call search_knowledge_base. " - "When you need to store new information, call add_to_knowledge_base." -) +backend = CompositeBackend([ + LocalShellBackend(workspace_dir="./workspace"), + FilesystemBackend(), +]) agent = create_deep_agent( model=llm, tools=[search_knowledge_base, add_to_knowledge_base], backend=backend, - system_prompt=system_prompt, + system_prompt="You are a helpful agent with a knowledge base. Use the tools search_knowledge_base and add_to_knowledge_base as needed.", ) -# Загрузка документов из директории -def load_documents(dir_path: str) -> None: - """Load all .txt files from dir_path into the knowledge base.""" - path = Path(dir_path) - if not path.is_dir(): - print(f"Directory not found: {dir_path}") - return - for file_path in path.glob("*.txt"): - try: - content = file_path.read_text(encoding="utf-8") - title = file_path.stem - result = add_to_knowledge_base(content, title) - print(f"Loaded {file_path.name}: {result}") - except Exception as e: - print(f"Error loading {file_path.name}: {e}") - -# Интерактивный клиент -async def run_cli() -> None: - thread_id = "session-1" - print("Welcome to the RAG agent CLI.") - print("Commands:") - print(" /add - add a new document") - print(" /search - search the knowledge base") - print(" /quit - exit") +async def main(): + print("RAG Agent CLI. Commands: /add <content>, /search <query>, /quit") while True: - try: - user_input = input("\n> ").strip() - except (EOFError, KeyboardInterrupt): - print("\nExiting.") - break - + user_input = input(">> ").strip() if not user_input: continue - if user_input.lower() == "/quit": - print("Goodbye.") + print("Goodbye!") break - - if user_input.lower() == "/add": - title = input("Title: ").strip() - print("Enter content (end with a single line containing only 'END'):") - lines: List[str] = [] - 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 - - if user_input.lower() == "/search": - query = input("Enter search query: ").strip() - if not query: - print("Empty query.") + if user_input.lower().startswith("/add"): + parts = user_input.split(maxsplit=2) + if len(parts) < 3: + print("Usage: /add <title> <content>") continue - result = search_knowledge_base(query) - print("\nSearch results:") - print(result) - continue + title, content = parts[1], parts[2] + message = f"Add document titled {title} with content: {content}" + elif user_input.lower().startswith("/search"): + parts = user_input.split(maxsplit=1) + if len(parts) < 2: + print("Usage: /search <query>") + continue + query = parts[1] + message = f"Search for {query}" + else: + message = user_input - # Any other message is sent to the agent - response = await agent.ainvoke( - {"messages": [HumanMessage(content=user_input)]}, - {"configurable": {"thread_id": thread_id}}, + result = await agent.ainvoke( + {"messages": [HumanMessage(content=message)]}, + {"configurable": {"thread_id": "session-1"}}, ) - agent_reply = response["messages"][-1].content - print(f"\nAgent: {agent_reply}") - -def main() -> None: - # Optional: load documents from a directory passed as first argument - if len(sys.argv) > 1: - dir_path = sys.argv[1] - print(f"Loading documents from {dir_path}...") - load_documents(dir_path) - asyncio.run(run_cli()) + print(result["messages"][-1].content) if __name__ == "__main__": - main() \ No newline at end of file + asyncio.run(main()) \ No newline at end of file