fix: main.py

main.py

@@ -2,20 +2,18 @@ import os
 import asyncio
 from pathlib import Path
 from dotenv import load_dotenv
-
 from langchain_openai import ChatOpenAI, OpenAIEmbeddings
 from langchain_chroma import Chroma
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_core.documents import Document
-from langchain_tavily import TavilySearchResults
 from langchain.tools import tool
 from deepagents import create_deep_agent
 from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
 
-# --------------------- 1. Загрузка переменных окружения ---------------------
+# Load environment variables
 load_dotenv()
 
-# --------------------- 2. LLM и Embeddings ---------------------
+# ---------- LLM and Embeddings ----------
 llm = ChatOpenAI(
     model="openai/gpt-oss-20b:free",
     base_url="https://openrouter.ai/api/v1",
@@ -29,78 +27,79 @@ embeddings = OpenAIEmbeddings(
     api_key=os.getenv("OPENAI_API_KEY"),
 )
 
-# --------------------- 3. Векторное хранилище Chroma ---------------------
+# ---------- Vector Store ----------
 CHROMA_DIR = Path("./chroma_db")
-CHROMA_DIR.mkdir(parents=True, exist_ok=True)
 vector_store = Chroma(
     collection_name="knowledge",
     embedding_function=embeddings,
     persist_directory=str(CHROMA_DIR),
 )
 
-# --------------------- 4. Загрузка документов ---------------------
+# ---------- Document Loader ----------
-DOCS_DIR = Path("./documents")
-if DOCS_DIR.exists():
-    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-    for file_path in DOCS_DIR.rglob("*.txt"):
-        text = file_path.read_text(encoding="utf-8")
-        docs = splitter.split_text(text)
-        documents = [Document(page_content=chunk, metadata={"source": str(file_path)}) for chunk in docs]
-        vector_store.add_documents(documents)
-    vector_store.persist()
 
-# --------------------- 5. Инструменты ---------------------
+def load_documents(directory: str, vectorstore: Chroma):
+    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    docs = []
+    for file_path in Path(directory).glob("**/*.*"):
+        if file_path.suffix.lower() not in {".txt", ".md"}:
+            continue
+        text = file_path.read_text(encoding="utf-8")
+        chunks = splitter.split_text(text)
+        docs.extend([Document(page_content=c, metadata={"source": str(file_path)}) for c in chunks])
+    if docs:
+        vectorstore.add_documents(docs)
+        vectorstore.persist()
+
+# Load initial documents if collection is empty
+if not CHROMA_DIR.exists() or not list(CHROMA_DIR.iterdir()):
+    load_documents("documents", vector_store)
+
+# ---------- Tools ----------
 @tool
 def search_local_kb(query: str, top_k: int = 3) -> str:
     """Semantic search in the local knowledge base."""
     docs = vector_store.similarity_search(query, k=top_k)
     if not docs:
-        return "No relevant local knowledge found."
+        return "No relevant information found in local knowledge base."
-    return "\n---\n".join(f"{i+1}. {doc.page_content[:200]}..." for i, doc in enumerate(docs))
+    return "\n---\n".join([f"{d.metadata.get('source', 'unknown')}\n{d.page_content}" for d in docs])
 
 @tool
 def web_search(query: str) -> str:
     """Web search using Tavily."""
-    tavily = TavilySearchResults(api_key=os.getenv("TAVILY_API_KEY"))
+    from langchain_tavily import TavilySearchResults
+    tavily = TavilySearchResults()
     results = tavily.run(query)
-    if not results:
+    return "\n---\n".join([f"{r['title']}\n{r['content']}" for r in results])
-        return "No web results found."
-    return "\n---\n".join(f"{i+1}. {res['title']}\n{res['url']}\n{res['content'][:200]}..." for i, res in enumerate(results))
 
-# --------------------- 6. Backend ---------------------
+# ---------- Backend ----------
 backend = CompositeBackend([
     LocalShellBackend(workspace_dir="./workspace"),
     FilesystemBackend(),
 ])
 
-# --------------------- 7. Создание агента ---------------------
+# ---------- Agent ----------
 agent = create_deep_agent(
     model=llm,
     tools=[search_local_kb, web_search],
     backend=backend,
-    system_prompt=(
+    system_prompt="You are a helpful assistant that can search both a local knowledge base and the web.\nWhen answering, always indicate the source: either 'chromadb' or 'tavily'.\nUse the appropriate tool based on the query context.",
-        "You are an assistant that answers user questions.\n"
-        "If the question is about information that should be in the local knowledge base, use the tool `search_local_kb`.\n"
-        "If the question requires up‑to‑date information from the web, use the tool `web_search`.\n"
-        "Always indicate the source of the answer in the format: `Источник: chromadb` or `Источник: tavily`."
-    ),
 )
 
-# --------------------- 8. CLI ---------------------
+# ---------- CLI ----------
-async def chat_loop():
+async def main():
-    print("Добро пожаловать в RAG‑агент. Введите 'exit' для выхода.")
+    print("RAG Agent with ChromaDB and Tavily. Type 'exit' to quit.")
     while True:
         user_input = input("\nЗапрос: ")
-        if user_input.lower() in {"exit", "quit", "q"}:
+        if user_input.lower() in {"exit", "quit"}:
-            print("До свидания!")
+            print("Goodbye!")
             break
         result = await agent.ainvoke(
             {"messages": [{"role": "user", "content": user_input}]},
             {"configurable": {"thread_id": "session-1"}},
         )
-        # Последнее сообщение агента
+        # Extract last message content
         content = result["messages"][-1]["content"]
         print(content)
 
 if __name__ == "__main__":
-    asyncio.run(chat_loop())
+    asyncio.run(main())