diff --git a/main.py b/main.py index e3ffeac..afa59fa 100644 --- a/main.py +++ b/main.py @@ -1,119 +1,77 @@ -import os import asyncio -from langchain_ollama import Ollama, OllamaEmbeddings +import os +from pathlib import Path + +from langchain_core.messages import HumanMessage 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_text_splitters import RecursiveCharacterTextSplitter + +from utils import llm, vector_store, splitter from langchain_core.documents import Document -from langchain_core.messages import HumanMessage -# Инициализация эмбеддинговой модели Ollama -embeddings = OllamaEmbeddings(model="nomic-embed-text") - -# Инициализация векторного хранилища Qdrant -vector_store = QdrantVectorStore( - url="http://localhost:6333", - collection_name="knowledge", - embedding_function=embeddings -) - -# Инструмент: поиск в базе знаний @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: {title} ({len(chunks)} chunks)" + return f"Added {len(docs)} chunks for {title}." -# Инициализация 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=system_prompt, + system_prompt="You are a helpful assistant with access to a knowledge base. Use the provided tools to search and add information.", ) -# Загрузка документов из директории в векторное хранилище -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 , /search , /quit") +async def main(): + print("RAG Agent CLI. Commands: /add, /search, /quit") while True: - user_input = input(">> ") - if not user_input.strip(): + user_input = input(">> ").strip() + if not user_input: continue - if user_input.startswith("/add"): - parts = user_input.split(maxsplit=1) - if len(parts) < 2: - print("Usage: /add ") - continue - 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] + if user_input.lower() == "/quit": + print("Goodbye.") + break + if user_input.lower().startswith("/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) - elif user_input.startswith("/search"): - parts = user_input.split(maxsplit=1) - if len(parts) < 2: - print("Usage: /search ") - continue - query = parts[1] - 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: - print("Unknown command. Use /add, /search, /quit") - -async def main(): - # При необходимости загрузить начальные документы - # load_documents_from_dir("./docs") - await interactive_cli() + continue + if user_input.lower().startswith("/search"): + query = input("Query: ").strip() + max_results_str = input("Max results (default 3): ").strip() + max_results = int(max_results_str) if max_results_str.isdigit() else 3 + result = search_knowledge_base(query, max_results) + print("Search results:") + print(result) + continue + # Regular message to agent + response = await agent.ainvoke( + {"messages": [HumanMessage(content=user_input)]}, + {"configurable": {"thread_id": "session-1"}}, + ) + print(response["messages"][-1].content) if __name__ == "__main__": asyncio.run(main()) \ No newline at end of file