commit 00bf2413b82ba725b9eb43ea17ed7e06df461ed3 Author: Даниил Викторов Date: Thu Jul 2 04:03:13 2026 +0000 add: main.py — Агент с RAG-памятью diff --git a/main.py b/main.py new file mode 100644 index 0000000..f4180fa --- /dev/null +++ b/main.py @@ -0,0 +1,169 @@ +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.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 + +load_dotenv() + +# ---------- LLM ---------- +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, +) + +# ---------- 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, +) + +# ---------- Text Splitter ---------- +splitter = RecursiveCharacterTextSplitter( + chunk_size=1000, + chunk_overlap=200, + separators=["\n\n", "\n", " "], +) + +# ---------- 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) + if not docs: + return "No relevant documents found." + return "\n---\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. + """ + chunks = splitter.split_text(content) + docs = [ + Document(page_content=chunk, metadata={"title": title, "chunk_index": i}) + for i, chunk in enumerate(chunks) + ] + vector_store.add_documents(docs) + return f"Added {len(docs)} chunks from '{title}' to the knowledge base." + + +# ---------- 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." + ), +) + +# ---------- 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.") + while True: + user_input = input(">>> ").strip() + if not user_input: + continue + if user_input.lower() == "/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: + response = await agent.ainvoke( + {"messages": [HumanMessage(content=user_input)]}, + {"configurable": {"thread_id": thread_id}}, + ) + answer = response["messages"][-1].content + print(answer) + except Exception as e: + print(f"Agent error: {e}") + + +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