add main.py

main.py

@@ -1,19 +1,22 @@
 import os
 import asyncio
 from pathlib import Path
-from typing import List
+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_chroma import Chroma
+from langchain_ollama import OllamaEmbeddings
 from langchain_tavily import TavilySearchResults
 from deepagents import create_deep_agent
 from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
 
-# --------------------------- LLM ---------------------------
+# Load env vars
+load_dotenv()
+
+# ---------- LLM ----------
 llm = ChatOpenAI(
     model="openai/gpt-oss-20b:free",
     base_url="https://openrouter.ai/api/v1",
@@ -21,79 +24,65 @@ llm = ChatOpenAI(
     temperature=0.0,
 )
 
-# --------------------------- Backend ---------------------------
+# ---------- Backend ----------
 backend = CompositeBackend([
     LocalShellBackend(workspace_dir="./workspace"),
     FilesystemBackend(),
 ])
 
-# --------------------------- Vectorstore ---------------------------
+# ---------- Vector Store ----------
 PERSIST_DIR = Path("./chroma_db")
-PERSIST_DIR.mkdir(parents=True, exist_ok=True)
+PERSIST_DIR.mkdir(exist_ok=True)
 
 embeddings = OllamaEmbeddings(model="nomic-embed-text")
 vectorstore = Chroma(persist_directory=str(PERSIST_DIR), embedding_function=embeddings)
 
 # Load documents from ./documents if not already loaded
-DOCS_DIR = Path("./documents")
-if DOCS_DIR.exists():
-    for file in DOCS_DIR.glob("**/*.*"):
-        if file.suffix.lower() in {".txt", ".md"}:
-            text = file.read_text(encoding="utf-8")
-            splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-            docs = splitter.split_text(text)
-            vectorstore.add_texts(docs, metadatas=[{"source": str(file)} for _ in docs])
-    vectorstore.persist()
+if not vectorstore.get_collection().count():
+    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    docs = []
+    for file in Path("./documents").glob("*.txt"):
+        text = file.read_text(encoding="utf-8")
+        docs.extend(splitter.split_text(text))
+    vectorstore.add_texts(docs)
 
-# --------------------------- Tools ---------------------------
+# ---------- Tools ----------
 @tool
 def search_local_kb(query: str, top_k: int = 3) -> str:
-    """Semantic search in the local ChromaDB knowledge base."""
+    """Semantic search in local ChromaDB knowledge base."""
     retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
     docs = retriever.invoke(query)
-    if not docs:
-        return "No local knowledge found."
-    return "\n---\n".join([f"{d.page_content[:500]}..." for d in docs])
+    return "\n".join(doc.page_content for doc in docs) if docs else "No local results found."
 
 @tool
 def web_search(query: str) -> str:
-    """Web search using Tavily."""
+    """Web search via Tavily."""
     tavily = TavilySearchResults(api_key=os.getenv("TAVILY_API_KEY"))
     results = tavily.invoke(query)
-    if not results:
-        return "No web results found."
-    return "\n---\n".join([f"{r['title']}: {r['content'][:500]}..." for r in results])
-
-# --------------------------- Agent ---------------------------
-SYSTEM_PROMPT = (
-    "You are an AI assistant that can answer questions using either a local knowledge base or the web. "
-    "If the answer can be found in the local documents, use the `search_local_kb` tool and prefix the response with `[Local KB]`. "
-    "If the answer requires up‑to‑date information, use the `web_search` tool and prefix the response with `[Web Search]`. "
-    "Always indicate the source in the response."
-)
+    return "\n".join(f"{r['title']}: {r['url']}" for r in results) if results else "No web results found."
 
+# ---------- Agent ----------
 agent = create_deep_agent(
     model=llm,
     tools=[search_local_kb, web_search],
     backend=backend,
-    system_prompt=SYSTEM_PROMPT,
+    system_prompt="You are a helpful RAG agent. For questions about local documents use search_local_kb, for up‑to‑date facts use web_search. Always state the source (chromadb or tavily) in the answer.",
 )
 
-# --------------------------- CLI ---------------------------
+# ---------- CLI ----------
 async def main():
     print("RAG Agent ready. Type 'exit' to quit.")
     while True:
         user_input = input("\nЗапрос: ")
         if user_input.lower() in {"exit", "quit"}:
-            print("Goodbye!")
             break
         result = await agent.ainvoke(
             {"messages": [HumanMessage(content=user_input)]},
             {"configurable": {"thread_id": "session-1"}},
         )
-        # The last message is the assistant's reply
+        # The agent returns a list of messages; last is the assistant reply
         reply = result["messages"][-1].content
-        print(f"\n{reply}")
+        print(f"\nОтвет: {reply}")
 
 if __name__ == "__main__":
     asyncio.run(main())