add main.py

main.py

@@ -1,93 +1,109 @@
 import os
-import sys
+import asyncio
 from pathlib import Path
-from dotenv import load_dotenv
-
 from langchain_openai import ChatOpenAI
+from langchain_core.messages import HumanMessage
+from langchain.tools import tool
 from langchain_ollama import OllamaEmbeddings
 from langchain_chroma import Chroma
 from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain.tools import tool
-from langchain.agents import AgentExecutor, create_openai_tools_agent
-from langchain_community.tools.tavily import TavilySearchResults
+from langchain_tavily import TavilySearchResults
+from deepagents import create_deep_agent
+from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
 
-# Load env variables
-load_dotenv()
-
-# ---------- LLM and embeddings ----------
+# ---------- 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,
 )
-embeddings = OllamaEmbeddings(model="nomic-embed-text")
 
-# ---------- Vectorstore ----------
+# ---------- Backend ----------
+backend = CompositeBackend([
+    LocalShellBackend(workspace_dir="./workspace"),
+    FilesystemBackend(),
+])
+
+# ---------- Vectorstore utilities ----------
 PERSIST_DIR = Path("./chroma_db")
 PERSIST_DIR.mkdir(parents=True, exist_ok=True)
-vectorstore = Chroma(persist_directory=str(PERSIST_DIR), embedding_function=embeddings)
 
-# Load documents if collection empty
-if not vectorstore.get_collection().count():
-    docs_dir = Path("./documents")
-    if docs_dir.exists():
+# Create or load Chroma vectorstore
+vectorstore = Chroma(
+    persist_directory=str(PERSIST_DIR),
+    embedding_function=OllamaEmbeddings(model="nomic-embed-text"),
+)
+
+# Load documents from a directory into the vectorstore
+
+def load_documents(directory: str, vectorstore):
+    docs = []
     splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-        for file in docs_dir.glob("**/*.*"):
-            if file.suffix.lower() in {".txt", ".md"}:
-                text = file.read_text(encoding="utf-8")
-                chunks = splitter.split_text(text)
-                vectorstore.add_texts(chunks, ids=[f"{file.name}_{i}" for i in range(len(chunks))])
+    for file_path in Path(directory).glob("**/*"):
+        if file_path.suffix.lower() in {".txt", ".md"}:
+            text = file_path.read_text(encoding="utf-8")
+            docs.extend(splitter.split_text(text))
+    # Convert to LangChain Documents
+    from langchain_core.documents import Document
+    documents = [Document(page_content=chunk) for chunk in docs]
+    vectorstore.add_documents(documents)
     vectorstore.persist()
 
+# Load documents once at startup (if not already loaded)
+if not any(PERSIST_DIR.iterdir()):
+    load_documents("./documents", vectorstore)
+
 # ---------- Tools ----------
-@tool("search_local_kb")
+@tool
 def search_local_kb(query: str, top_k: int = 3) -> str:
-    """Semantic search in local ChromaDB knowledge base."""
+    """Semantic search in the local ChromaDB knowledge base."""
     retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
     docs = retriever.get_relevant_documents(query)
     if not docs:
         return "No relevant documents found in local KB."
     return "\n---\n".join(doc.page_content for doc in docs)
 
-@tool("web_search")
+@tool
 def web_search(query: str) -> str:
     """Web search using Tavily."""
-    tavily = TavilySearchResults(max_results=3)
+    tavily = TavilySearchResults(api_key=os.getenv("TAVILY_API_KEY"))
     results = tavily.run(query)
     if not results:
         return "No web results found."
     return "\n---\n".join(f"{r['title']}\n{r['url']}\n{r.get('content', '')}" for r in results)
 
 # ---------- Agent ----------
-SYSTEM_PROMPT = (
-    "You are an AI assistant. For questions about local documents, use the tool 'search_local_kb'. "
-    "For up-to-date information or news, use 'web_search'. Always indicate the source in your answer."
+agent = create_deep_agent(
+    model=llm,
+    tools=[search_local_kb, web_search],
+    backend=backend,
+    system_prompt=(
+        "You are an AI assistant that answers user questions.\n"
+        "If the question is about information that should be in the local knowledge base,\n"
+        "use the search_local_kb tool.\n"
+        "If the question requires up‑to‑date information from the web,\n"
+        "use the web_search tool.\n"
+        "Always indicate the source of the answer in the format:\n"
+        "[Source: chromadb] or [Source: tavily] before the answer."
+    ),
 )
 
-tools = [search_local_kb, web_search]
-agent = create_openai_tools_agent(llm, tools, system_message=SYSTEM_PROMPT)
-agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
-
 # ---------- CLI ----------
-
-def main():
+async def main():
     print("RAG Agent with ChromaDB and Tavily. Type 'exit' to quit.")
     while True:
-        try:
-            query = input("\nЗапрос: ")
-        except (EOFError, KeyboardInterrupt):
-            print("\nBye!")
+        user_input = input("\nЗапрос: ")
+        if user_input.lower() in {"exit", "quit"}:
+            print("Goodbye!")
             break
-        if query.strip().lower() in {"exit", "quit"}:
-            print("Bye!")
-            break
-        result = agent_executor.invoke({"input": query})
-        # The tool name is stored in the tool_calls field of the result
-        tool_name = result.get("tool_calls", [{}])[0].get("name", "unknown")
-        source = "tavily" if tool_name == "web_search" else "chromadb"
-        print(f"\n[{'Web Search' if source=='tavily' else 'Local KB'}] {result.get('output', '')}\n")
-        print(f"Источник: {source}\n")
+        result = await agent.ainvoke(
+            {"messages": [HumanMessage(content=user_input)]},
+            {"configurable": {"thread_id": "session-1"}},
+        )
+        # The agent returns a list of messages; the last is the assistant reply
+        reply = result["messages"][-1].content
+        print(f"\n{reply}")
 
 if __name__ == "__main__":
-    main()
+    asyncio.run(main())