From fcafdafdc2cdac25ee1bfbe5db6ecd31fcc25122 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=94=D0=B0=D0=BD=D0=B8=D0=B8=D0=BB=20=D0=92=D0=B8=D0=BA?= =?UTF-8?q?=D1=82=D0=BE=D1=80=D0=BE=D0=B2?= Date: Thu, 2 Jul 2026 08:17:26 +0000 Subject: [PATCH] =?UTF-8?q?fix:=20main.py=20=E2=80=94=20=D0=90=D0=B3=D0=B5?= =?UTF-8?q?=D0=BD=D1=82=20=D1=81=20RAG-=D0=BF=D0=B0=D0=BC=D1=8F=D1=82?= =?UTF-8?q?=D1=8C=D1=8E?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- main.py | 178 +++++++++++++++++++------------------------------------- 1 file changed, 60 insertions(+), 118 deletions(-) diff --git a/main.py b/main.py index f4180fa..f67d0c0 100644 --- a/main.py +++ b/main.py @@ -1,21 +1,20 @@ import os import asyncio -from pathlib import Path -from typing import List - from dotenv import load_dotenv from langchain_openai import ChatOpenAI, OpenAIEmbeddings -from langchain_chroma import Chroma +from langchain_core.messages import HumanMessage from langchain_core.documents import Document -from langchain.text_splitter import RecursiveCharacterTextSplitter 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_qdrant import QdrantVectorStore +from langchain_text_splitters import RecursiveCharacterTextSplitter +from qdrant_client import QdrantClient +# Загрузка переменных окружения load_dotenv() -# ---------- LLM ---------- +# Инициализация LLM llm = ChatOpenAI( model="openai/gpt-oss-20b:free", base_url="https://openrouter.ai/api/v1", @@ -23,147 +22,90 @@ llm = ChatOpenAI( temperature=0.0, ) -# ---------- Vector Store ---------- +# Инициализация эмбеддингов embeddings = OpenAIEmbeddings( model="text-embedding-3-small", base_url="https://openrouter.ai/api/v1", api_key=os.getenv("OPENAI_API_KEY"), ) -vector_store = Chroma( - collection_name="knowledge", - embedding_function=embeddings, +# Инициализация Qdrant +client = QdrantClient(url="http://localhost:6333") +collection_name = "knowledge_base" +vector_store = QdrantVectorStore( + client=client, + collection_name=collection_name, + embeddings=embeddings, ) -# ---------- Text Splitter ---------- -splitter = RecursiveCharacterTextSplitter( - chunk_size=1000, - chunk_overlap=200, - separators=["\n\n", "\n", " "], -) +# Чанкинг +splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) -# ---------- RAG Tools ---------- +# Инструмент поиска @tool -def search_knowledge_base(query: str, max_results: int = 3) -> str: - """ - Perform a semantic search in the knowledge base. - Returns the concatenated contents of the most relevant documents. - """ - docs: List[Document] = vector_store.similarity_search(query, k=max_results) +def search_knowledge_base(query: str, max_results: int) -> str: + """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".join(doc.page_content for doc in docs) - + return "No results found." + return "\n".join(doc.page_content for doc in docs) +# Инструмент добавления @tool -def add_to_knowledge_base(content: str, title: str = "document") -> str: - """ - Add a new document to the knowledge base. - The content will be split into chunks before indexing. - """ +def add_to_knowledge_base(content: str, title: str) -> str: + """Add content to the knowledge base.""" chunks = splitter.split_text(content) - docs = [ - Document(page_content=chunk, metadata={"title": title, "chunk_index": i}) - for i, chunk in enumerate(chunks) - ] + 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}' to the knowledge base." + return f"Added {len(docs)} chunks for {title}." +# Backend для deepagents +backend = CompositeBackend([ + LocalShellBackend(workspace_dir="./workspace"), + FilesystemBackend(), +]) -# ---------- 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 an AI assistant with access to a local knowledge base. " - "When you need factual information, use the provided tools: " - "`search_knowledge_base` to retrieve data and `add_to_knowledge_base` to store new documents. " - "Always cite sources from the knowledge base in your answers." - ), + system_prompt="You are a helpful agent with access to a knowledge base. Use the tools to search and add information.", ) -# ---------- Helper Functions ---------- -def load_documents_from_directory(directory: Path) -> None: - """ - Recursively read .txt files from the given directory and add them to the knowledge base. - """ - for file_path in directory.rglob("*.txt"): - try: - content = file_path.read_text(encoding="utf-8") - title = file_path.stem - add_to_knowledge_base(content, title) - print(f"Loaded {file_path}") - except Exception as e: - print(f"Failed to load {file_path}: {e}") - - -async def chat_loop() -> None: - """ - Simple CLI loop. - Commands: - /add - add a text file or all txt files in a directory - /search - search the knowledge base - /quit - exit - Anything else is sent to the agent as a user message. - """ - thread_id = "cli-session" - print("AI assistant ready. Type /quit to exit.") +async def main(): + print("Interactive agent. Commands: /add, /search, /quit") while True: - user_input = input(">>> ").strip() + try: + user_input = input("> ").strip() + except EOFError: + break if not user_input: continue - if user_input.lower() == "/quit": + if user_input.startswith("/add"): + title = input("Title: ").strip() + content = input("Content: ").strip() + result = add_to_knowledge_base(content, title) + print(result) + elif user_input.startswith("/search"): + query = input("Query: ").strip() + max_str = input("Max results (int): ").strip() + try: + max_results = int(max_str) + except ValueError: + max_results = 3 + result = search_knowledge_base(query, max_results) + print(result) + elif user_input.startswith("/quit"): print("Goodbye!") break - if user_input.startswith("/add"): - parts = user_input.split(maxsplit=1) - if len(parts) != 2: - print("Usage: /add ") - continue - path = Path(parts[1]).expanduser().resolve() - if path.is_dir(): - load_documents_from_directory(path) - elif path.is_file() and path.suffix.lower() == ".txt": - content = path.read_text(encoding="utf-8") - add_to_knowledge_base(content, path.stem) - print(f"Added file {path}") - else: - print("Provide a .txt file or a directory containing .txt files.") - continue - if user_input.startswith("/search"): - parts = user_input.split(maxsplit=1) - if len(parts) != 2: - print("Usage: /search ") - continue - query = parts[1] - result = search_knowledge_base(query) - print(f"Search results:\n{result}") - continue - - # Normal conversation with the agent - try: + else: + # обычный диалог с агентом response = await agent.ainvoke( {"messages": [HumanMessage(content=user_input)]}, - {"configurable": {"thread_id": thread_id}}, + {"configurable": {"thread_id": "session-1"}}, ) - answer = response["messages"][-1].content - print(answer) - except Exception as e: - print(f"Agent error: {e}") - + print(response["messages"][-1].content) if __name__ == "__main__": - # Optional: preload a default docs folder - default_dir = Path("./docs") - if default_dir.is_dir(): - load_documents_from_directory(default_dir) - asyncio.run(chat_loop()) \ No newline at end of file + asyncio.run(main()) \ No newline at end of file