add: main.py — Агент с RAG-памятью

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 <path>   - add a text file or all txt files in a directory
+      /search <q>   - 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 <path>")
+                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 <query>")
+                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())