diff --git a/main.py b/main.py index 8449355..e3ffeac 100644 --- a/main.py +++ b/main.py @@ -1,111 +1,119 @@ import os import asyncio -from langchain_openai import ChatOpenAI +from langchain_ollama import Ollama, OllamaEmbeddings 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 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. +from langchain_text_splitters import RecursiveCharacterTextSplitter +from langchain_core.documents import Document +from langchain_core.messages import HumanMessage +# Инициализация эмбеддинговой модели Ollama embeddings = OllamaEmbeddings(model="nomic-embed-text") -# 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) +# Инициализация векторного хранилища Qdrant vector_store = QdrantVectorStore( - client=qdrant_client, + url="http://localhost:6333", collection_name="knowledge", embedding_function=embeddings ) -# 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 = 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.""" + splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) 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 for {title}." + return f"Added: {title} ({len(chunks)} chunks)" -# 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, -) +# Инициализация LLM Ollama +llm = Ollama(model="llama3") +# Backend для deepagents backend = CompositeBackend([ LocalShellBackend(workspace_dir="./workspace"), FilesystemBackend(), ]) +# Системный промпт агента +system_prompt = ( + "You are a helpful agent with access to a knowledge base. " + "Use the provided tools to search and add information. " + "When answering user queries, first search the knowledge base and then provide a concise response." +) + +# Создание агента agent = create_deep_agent( model=llm, tools=[search_knowledge_base, add_to_knowledge_base], backend=backend, - system_prompt="You are a helpful agent with a knowledge base. Use the tools search_knowledge_base and add_to_knowledge_base as needed.", + system_prompt=system_prompt, ) -async def main(): - print("RAG Agent CLI. Commands: /add <content>, /search <query>, /quit") +# Загрузка документов из директории в векторное хранилище +def load_documents_from_dir(dir_path: str): + for root, dirs, files in os.walk(dir_path): + for file in files: + if file.lower().endswith(".txt"): + path = os.path.join(root, file) + with open(path, "r", encoding="utf-8") as f: + content = f.read() + title = os.path.splitext(file)[0] + result = add_to_knowledge_base(content, title) + print(result) + +# Интерактивный клиент +async def interactive_cli(): + print("Welcome to the RAG agent. Commands: /add <file_path>, /search <query>, /quit") while True: - user_input = input(">> ").strip() - if not user_input: + user_input = input(">> ") + if not user_input.strip(): continue - if user_input.lower() == "/quit": - print("Goodbye!") - break - if user_input.lower().startswith("/add"): - parts = user_input.split(maxsplit=2) - if len(parts) < 3: - print("Usage: /add <title> <content>") + if user_input.startswith("/add"): + parts = user_input.split(maxsplit=1) + if len(parts) < 2: + print("Usage: /add <file_path>") continue - title, content = parts[1], parts[2] - message = f"Add document titled {title} with content: {content}" - elif user_input.lower().startswith("/search"): + file_path = parts[1] + if not os.path.isfile(file_path): + print(f"File not found: {file_path}") + continue + with open(file_path, "r", encoding="utf-8") as f: + content = f.read() + title = os.path.splitext(os.path.basename(file_path))[0] + result = add_to_knowledge_base(content, title) + print(result) + elif user_input.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}" + response = await agent.ainvoke( + {"messages": [HumanMessage(content=query)]}, + {"configurable": {"thread_id": "session-1"}}, + ) + print(response["messages"][-1].content) + elif user_input.startswith("/quit"): + print("Goodbye!") + break else: - message = user_input + print("Unknown command. Use /add, /search, /quit") - result = await agent.ainvoke( - {"messages": [HumanMessage(content=message)]}, - {"configurable": {"thread_id": "session-1"}}, - ) - print(result["messages"][-1].content) +async def main(): + # При необходимости загрузить начальные документы + # load_documents_from_dir("./docs") + await interactive_cli() if __name__ == "__main__": asyncio.run(main()) \ No newline at end of file